From b9c7664ac34f695e78e9fbbcd54e729df93a6668 Mon Sep 17 00:00:00 2001
From: Rishabh Singh <6513075+findingrish@users.noreply.github.com>
Date: Fri, 28 Jun 2024 11:06:56 +0530
Subject: [PATCH] Fix empty datasource schema on the Broker when metadata query
 is disabled (#16645)

* Fix build

* Fix empty datasource schema on the broker

* review comment

* Remove unused import
---
 ...-compose.centralized-datasource-schema.yml |  2 +-
 .../druid/testing/utils/DataLoaderHelper.java |  4 ++++
 .../testing/utils/SqlTestQueryHelper.java     | 23 ++++++++++++++++++
 .../schema/BrokerSegmentMetadataCache.java    | 10 ++++++++
 .../BrokerSegmentMetadataCacheTest.java       | 24 +++++++++++++++++++
 5 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/integration-tests/docker/docker-compose.centralized-datasource-schema.yml b/integration-tests/docker/docker-compose.centralized-datasource-schema.yml
index 39ce98b1302..e89e49bc132 100644
--- a/integration-tests/docker/docker-compose.centralized-datasource-schema.yml
+++ b/integration-tests/docker/docker-compose.centralized-datasource-schema.yml
@@ -81,7 +81,7 @@ services:
       service: druid-broker
     environment:
       - DRUID_INTEGRATION_TEST_GROUP=${DRUID_INTEGRATION_TEST_GROUP}
-      - druid_sql_planner_metadataRefreshPeriod=PT20S
+      - druid_sql_planner_metadataRefreshPeriod=PT30S
       - druid_sql_planner_disableSegmentMetadataQueries=true
     depends_on:
       - druid-coordinator
diff --git a/integration-tests/src/main/java/org/apache/druid/testing/utils/DataLoaderHelper.java b/integration-tests/src/main/java/org/apache/druid/testing/utils/DataLoaderHelper.java
index 692ab962e62..67317f3e911 100644
--- a/integration-tests/src/main/java/org/apache/druid/testing/utils/DataLoaderHelper.java
+++ b/integration-tests/src/main/java/org/apache/druid/testing/utils/DataLoaderHelper.java
@@ -23,6 +23,7 @@ import com.google.inject.Inject;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.testing.clients.CoordinatorResourceTestClient;
+import org.testng.Assert;
 
 public final class DataLoaderHelper
 {
@@ -50,6 +51,9 @@ public final class DataLoaderHelper
         () -> sqlTestQueryHelper.isDatasourceLoadedInSQL(datasource),
         StringUtils.format("Waiting for [%s] to be ready for SQL queries", datasource)
     );
+
+    Assert.assertTrue(sqlTestQueryHelper.verifyTimeColumnIsPresent(datasource));
+
     LOG.info("Datasource [%s] ready for SQL queries", datasource);
   }
 }
diff --git a/integration-tests/src/main/java/org/apache/druid/testing/utils/SqlTestQueryHelper.java b/integration-tests/src/main/java/org/apache/druid/testing/utils/SqlTestQueryHelper.java
index 962b4a103d0..06a1f680b70 100644
--- a/integration-tests/src/main/java/org/apache/druid/testing/utils/SqlTestQueryHelper.java
+++ b/integration-tests/src/main/java/org/apache/druid/testing/utils/SqlTestQueryHelper.java
@@ -69,4 +69,27 @@ public class SqlTestQueryHelper extends AbstractTestQueryHelper<SqlQueryWithResu
       return false;
     }
   }
+
+  public boolean verifyTimeColumnIsPresent(String datasource)
+  {
+    final SqlQuery query = new SqlQuery(
+        "SELECT __time FROM \"" + datasource + "\" LIMIT 1",
+        null,
+        false,
+        false,
+        false,
+        null,
+        null
+    );
+
+    try {
+      //noinspection unchecked
+      queryClient.query(getQueryURL(broker), query);
+      return true;
+    }
+    catch (Exception e) {
+      LOG.debug(e, "Check query failed");
+      return false;
+    }
+  }
 }
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java
index 0573d8d49ee..7974ed460eb 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java
@@ -246,6 +246,16 @@ public class BrokerSegmentMetadataCache extends AbstractSegmentMetadataCache<Phy
         continue;
       }
 
+      if (rowSignature.getColumnNames().isEmpty()) {
+        // this case could arise when metadata refresh is disabled on broker
+        // and a new datasource is added
+        log.info("datasource [%s] schema has not been initialized yet, "
+                 + "check coordinator logs if this message is persistent.", dataSource);
+        // this is a harmless call
+        tables.remove(dataSource);
+        continue;
+      }
+
       final PhysicalDatasourceMetadata physicalDatasourceMetadata = dataSourceMetadataFactory.build(dataSource, rowSignature);
       updateDSMetadata(dataSource, physicalDatasourceMetadata);
     }
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCacheTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCacheTest.java
index f8660b63494..23b2759286c 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCacheTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCacheTest.java
@@ -1027,4 +1027,28 @@ public class BrokerSegmentMetadataCacheTest extends BrokerSegmentMetadataCacheTe
     buildSchemaMarkAndTableLatch();
     serverView.invokeSegmentSchemasAnnouncedDummy();
   }
+
+  @Test
+  public void testNoDatasourceSchemaWhenNoSegmentMetadata() throws InterruptedException, IOException
+  {
+    BrokerSegmentMetadataCacheConfig config = new BrokerSegmentMetadataCacheConfig();
+    config.setDisableSegmentMetadataQueries(true);
+
+    BrokerSegmentMetadataCache schema = buildSchemaMarkAndTableLatch(
+        config,
+        new NoopCoordinatorClient()
+    );
+
+    schema.start();
+    schema.awaitInitialization();
+
+    List<DataSegment> segments = schema.getSegmentMetadataSnapshot().values()
+                                .stream()
+                                .map(AvailableSegmentMetadata::getSegment)
+                                .collect(Collectors.toList());
+
+    schema.refresh(segments.stream().map(DataSegment::getId).collect(Collectors.toSet()), Collections.singleton("foo"));
+
+    Assert.assertNull(schema.getDatasource("foo"));
+  }
 }