Wait for SQL table visibility in embedded tests. (#19095)

* Wait for SQL table visibility in embedded tests.

Several embedded tests wait for segment/schemaCache/refresh/count as a
proxy for when a datasource is queryable via SQL. However, this metric
is emitted before the SQL table is set up. This creates a window where
SQL queries can fail with "Object 'datasource' not found".

This patch waits for segment/schemaCache/rowSignature/column/count,
which is emitted after a SQL table is set up and visible. This is done
either as a replacement for segment/schemaCache/refresh/count (if we
only care about table visibility) or in addition to it (if we also want
to wait for a specific number of segments to be detected).

* Fix import order.

* Fix import order.

Author: gianm

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/IngestionSmokeTest.java

@@ -40,6 +40,7 @@
 import org.apache.druid.metadata.storage.postgresql.PostgreSQLMetadataStorageModule;
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.http.SqlTaskStatus;
+import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.tasklogs.TaskLogStreamer;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
 import org.apache.druid.testing.embedded.EmbeddedCoordinator;
@@ -402,6 +403,11 @@ private void waitForSegmentsToBeQueryable(int numSegments)
                       .hasDimension(DruidMetrics.DATASOURCE, dataSource),
         agg -> agg.hasSumAtLeast(numSegments)
     );
+    eventCollector.latchableEmitter().waitForEvent(
+        event -> event.hasMetricName(Metric.SCHEMA_ROW_SIGNATURE_COLUMN_COUNT)
+                      .hasService("druid/broker")
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource)
+    );
   }
 
   /**

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/StreamIndexDataFormatsTestBase.java

@@ -45,6 +45,7 @@
 import org.apache.druid.indexing.overlord.supervisor.SupervisorSpec;
 import org.apache.druid.java.util.common.parsers.JSONPathSpec;
 import org.apache.druid.query.DruidMetrics;
+import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
 import org.apache.druid.testing.embedded.EmbeddedCoordinator;
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
@@ -552,7 +553,7 @@ private void waitForDataAndVerifyIngestedEvents(String dataSource, int expectedC
     );
     // Wait for the schema cache to refresh for the datasource under test
     broker.latchableEmitter().waitForEvent(
-        event -> event.hasMetricName("segment/schemaCache/refresh/count")
+        event -> event.hasMetricName(Metric.SCHEMA_ROW_SIGNATURE_COLUMN_COUNT)
                       .hasDimension(DruidMetrics.DATASOURCE, dataSource)
     );
 

embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java

@@ -26,6 +26,7 @@
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.http.SqlTaskStatus;
 import org.apache.druid.segment.TestIndex;
+import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
 import org.apache.druid.testing.embedded.EmbeddedClusterApis;
@@ -149,7 +150,7 @@ protected void setupLookups() throws Exception
         agg -> agg.hasSumAtLeast(totalRows)
     );
     broker.latchableEmitter().waitForEvent(
-        event -> event.hasMetricName("segment/schemaCache/refresh/count")
+        event -> event.hasMetricName(Metric.SCHEMA_ROW_SIGNATURE_COLUMN_COUNT)
                       .hasDimension(DruidMetrics.DATASOURCE, dataSource)
     );
   }

embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/HttpEmitterEventCollectorTest.java

@@ -22,6 +22,7 @@
 import org.apache.druid.common.utils.IdUtils;
 import org.apache.druid.indexing.common.task.IndexTask;
 import org.apache.druid.query.DruidMetrics;
+import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.testing.embedded.EmbeddedBroker;
 import org.apache.druid.testing.embedded.EmbeddedCoordinator;
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
@@ -89,6 +90,11 @@ public void test_runIndexTask_andQueryData()
                       .hasService("druid/broker"),
         agg -> agg.hasSumAtLeast(10)
     );
+    eventCollector.latchableEmitter().waitForEvent(
+        event -> event.hasMetricName(Metric.SCHEMA_ROW_SIGNATURE_COLUMN_COUNT)
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource)
+                      .hasService("druid/broker")
+    );
 
     cluster.callApi().verifySqlQuery("SELECT * FROM %s", dataSource, Resources.InlineData.CSV_10_DAYS);
   }

services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java

@@ -44,6 +44,7 @@
 import org.apache.druid.query.http.ClientSqlQuery;
 import org.apache.druid.rpc.indexing.OverlordClient;
 import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.metadata.Metric;
 import org.apache.druid.server.metrics.LatchableEmitter;
 import org.apache.druid.sql.http.ResultFormat;
 import org.apache.druid.timeline.DataSegment;
@@ -344,6 +345,10 @@ public void waitForAllSegmentsToBeAvailable(String dataSource, EmbeddedCoordinat
           agg -> agg.hasSumAtLeast(numTombstones)
       );
     }
+    broker.latchableEmitter().waitForEvent(
+        event -> event.hasMetricName(Metric.SCHEMA_ROW_SIGNATURE_COLUMN_COUNT)
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource)
+    );
   }
 
   /**