From 79e48c6b45a1b80012673ea9d8828fc718ae9890 Mon Sep 17 00:00:00 2001
From: Akshat Jain <akjn11@gmail.com>
Date: Thu, 18 Apr 2024 21:52:02 +0530
Subject: [PATCH] Fix NPE while loading lookups from empty JDBC source (#16307)

---
 .../lookup/namespace/JdbcCacheGenerator.java  | 24 ++++++++-------
 .../NamespacedExtractorModuleTest.java        |  2 +-
 .../cache/JdbcExtractionNamespaceTest.java    | 30 +++++++++++++++++++
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/extensions-core/lookups-cached-global/src/main/java/org/apache/druid/server/lookup/namespace/JdbcCacheGenerator.java b/extensions-core/lookups-cached-global/src/main/java/org/apache/druid/server/lookup/namespace/JdbcCacheGenerator.java
index f05ff702504..ccf8504875d 100644
--- a/extensions-core/lookups-cached-global/src/main/java/org/apache/druid/server/lookup/namespace/JdbcCacheGenerator.java
+++ b/extensions-core/lookups-cached-global/src/main/java/org/apache/druid/server/lookup/namespace/JdbcCacheGenerator.java
@@ -204,18 +204,20 @@ public final class JdbcCacheGenerator implements CacheGenerator<JdbcExtractionNa
     if (tsColumn == null) {
       return null;
     }
-    final Timestamp update = dbi.withHandle(
-        handle -> {
-          final String query = StringUtils.format(
-              "SELECT MAX(%s) FROM %s",
-              tsColumn, table
-          );
-          return handle
-              .createQuery(query)
-              .map(TimestampMapper.FIRST)
-              .first();
-        }
+    final String query = StringUtils.format(
+        "SELECT MAX(%s) FROM %s",
+        tsColumn, table
     );
+    final Timestamp update = dbi.withHandle(
+        handle -> handle
+            .createQuery(query)
+            .map(TimestampMapper.FIRST)
+            .first()
+    );
+    if (update == null) {
+      LOG.info("Lookup table[%s] is empty. No rows returned for the query[%s].", table, query);
+      return null;
+    }
     return update.getTime();
   }
 }
diff --git a/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/NamespacedExtractorModuleTest.java b/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/NamespacedExtractorModuleTest.java
index 715da7bed92..4e04f034cab 100644
--- a/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/NamespacedExtractorModuleTest.java
+++ b/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/NamespacedExtractorModuleTest.java
@@ -121,7 +121,7 @@ public class NamespacedExtractorModuleTest
     Assert.assertNotNull(version);
     Map<String, String> map = cache.getCache();
     Assert.assertEquals("bar", map.get("foo"));
-    Assert.assertEquals(null, map.get("baz"));
+    Assert.assertNull(map.get("baz"));
   }
 
   @Test
diff --git a/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/cache/JdbcExtractionNamespaceTest.java b/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/cache/JdbcExtractionNamespaceTest.java
index cf71a5da490..ada8e2ee926 100644
--- a/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/cache/JdbcExtractionNamespaceTest.java
+++ b/extensions-core/lookups-cached-global/src/test/java/org/apache/druid/server/lookup/namespace/cache/JdbcExtractionNamespaceTest.java
@@ -390,6 +390,36 @@ public class JdbcExtractionNamespaceTest
     }
   }
 
+  @Test(timeout = 60_000L)
+  public void testEmptyTable()
+      throws InterruptedException
+  {
+    // Delete existing rows from table.
+    final Handle handle = derbyConnectorRule.getConnector().getDBI().open();
+    handle.createStatement(
+        StringUtils.format("DELETE FROM %s", TABLE_NAME)
+    ).setQueryTimeout(1).execute();
+
+    final JdbcExtractionNamespace extractionNamespace = new JdbcExtractionNamespace(
+        derbyConnectorRule.getMetadataConnectorConfig(),
+        TABLE_NAME,
+        KEY_NAME,
+        VAL_NAME,
+        tsColumn,
+        null,
+        new Period(0),
+        null,
+        0,
+        null,
+        new JdbcAccessSecurityConfig()
+    );
+    try (CacheScheduler.Entry entry = scheduler.schedule(extractionNamespace)) {
+      CacheSchedulerTest.waitFor(entry);
+      final Map<String, String> map = entry.getCache();
+      Assert.assertTrue(map.isEmpty());
+    }
+  }
+
   @Test(timeout = 60_000L)
   public void testSkipOld()
       throws InterruptedException