kfaraz commented on code in PR #18805:
URL: https://github.com/apache/druid/pull/18805#discussion_r2585838884


##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/JdbcQueryTest.java:
##########
@@ -47,69 +40,55 @@
 import java.util.Properties;
 import java.util.Set;
 
-@Test(groups = {TestNGGroup.QUERY, TestNGGroup.CENTRALIZED_DATASOURCE_SCHEMA})
-@Guice(moduleFactory = DruidTestModuleFactory.class)
-public class ITJdbcQueryTest
+public class JdbcQueryTest extends QueryTestBase
 {
-  private static final Logger LOG = new Logger(ITJdbcQueryTest.class);
-  private static final String WIKIPEDIA_DATA_SOURCE = "wikipedia_editstream";
+  private static final Logger LOG = new Logger(JdbcQueryTest.class);
   private static final String CONNECTION_TEMPLATE = 
"jdbc:avatica:remote:url=%s/druid/v2/sql/avatica/";
   private static final String TLS_CONNECTION_TEMPLATE =
       
"jdbc:avatica:remote:url=%s/druid/v2/sql/avatica/;truststore=%s;truststore_password=%s;keystore=%s;keystore_password=%s;key_password=%s";
 
   private static final String QUERY_TEMPLATE =
-      "SELECT \"user\", SUM(\"added\"), COUNT(*)" +
-      "FROM \"wikipedia\" " +
-      "WHERE \"__time\" >= CURRENT_TIMESTAMP - INTERVAL '99' YEAR AND 
\"language\" = %s" +
-      "GROUP BY 1 ORDER BY 3 DESC LIMIT 10";
-  private static final String QUERY = StringUtils.format(QUERY_TEMPLATE, 
"'en'");
-
-  private static final String QUERY_PARAMETERIZED = 
StringUtils.format(QUERY_TEMPLATE, "?");
+      "SELECT \"item\", SUM(\"value\"), COUNT(*) "
+      + "FROM \"%s\" "
+      + "WHERE \"__time\" >= CURRENT_TIMESTAMP - INTERVAL '99' YEAR AND 
\"value\" < %s \n"
+      + "GROUP BY 1 ORDER BY 3 DESC LIMIT 10";
 
   private String[] connections;
   private Properties connectionProperties;
 
-  @Inject
-  private IntegrationTestingConfig config;
-
-  @Inject
-  SSLClientConfig sslConfig;
-
-  @Inject
-  private DataLoaderHelper dataLoaderHelper;
+  private String tableName;
 
-  @BeforeMethod
-  public void before()
+  @Override
+  protected void beforeAll()
   {
     connectionProperties = new Properties();
     connectionProperties.setProperty("user", "admin");
     connectionProperties.setProperty("password", "priest");
     connections = new String[]{
-        StringUtils.format(CONNECTION_TEMPLATE, config.getRouterUrl()),
-        StringUtils.format(CONNECTION_TEMPLATE, config.getBrokerUrl()),
-        StringUtils.format(
-            TLS_CONNECTION_TEMPLATE,
-            config.getRouterTLSUrl(),
-            sslConfig.getTrustStorePath(),
-            sslConfig.getTrustStorePasswordProvider().getPassword(),
-            sslConfig.getKeyStorePath(),
-            sslConfig.getKeyStorePasswordProvider().getPassword(),
-            sslConfig.getKeyManagerPasswordProvider().getPassword()
-        ),
-        StringUtils.format(
-            TLS_CONNECTION_TEMPLATE,
-            config.getBrokerTLSUrl(),
-            sslConfig.getTrustStorePath(),
-            sslConfig.getTrustStorePasswordProvider().getPassword(),
-            sslConfig.getKeyStorePath(),
-            sslConfig.getKeyStorePasswordProvider().getPassword(),
-            sslConfig.getKeyManagerPasswordProvider().getPassword()
-        )
+        StringUtils.format(CONNECTION_TEMPLATE, getServerUrl(router)),
+        StringUtils.format(CONNECTION_TEMPLATE, getServerUrl(broker)),
+        // Add in the consecutive patch
+        // StringUtils.format(
+        //    TLS_CONNECTION_TEMPLATE,
+        //    config.getRouterTLSUrl(),
+        //    sslConfig.getTrustStorePath(),
+        //    sslConfig.getTrustStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyStorePath(),
+        //    sslConfig.getKeyStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyManagerPasswordProvider().getPassword()
+        //),
+        // StringUtils.format(
+        //    TLS_CONNECTION_TEMPLATE,
+        //    config.getBrokerTLSUrl(),
+        //    sslConfig.getTrustStorePath(),
+        //    sslConfig.getTrustStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyStorePath(),
+        //    sslConfig.getKeyStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyManagerPasswordProvider().getPassword()
+        // )
     };
-    // ensure that wikipedia segments are loaded completely
-    dataLoaderHelper.waitUntilDatasourceIsReady(WIKIPEDIA_DATA_SOURCE);
-    dataLoaderHelper.waitUntilDatasourceIsReady("wikipedia");
-    dataLoaderHelper.waitUntilDatasourceIsReady("twitterstream");
+
+    tableName = ingestBasicData();

Review Comment:
   Nit: In Druid, we generally use the term datasource. Maybe rename this field 
to `testDataSource` or similar?



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/JdbcQueryTest.java:
##########
@@ -47,69 +40,55 @@
 import java.util.Properties;
 import java.util.Set;
 
-@Test(groups = {TestNGGroup.QUERY, TestNGGroup.CENTRALIZED_DATASOURCE_SCHEMA})
-@Guice(moduleFactory = DruidTestModuleFactory.class)
-public class ITJdbcQueryTest
+public class JdbcQueryTest extends QueryTestBase
 {
-  private static final Logger LOG = new Logger(ITJdbcQueryTest.class);
-  private static final String WIKIPEDIA_DATA_SOURCE = "wikipedia_editstream";
+  private static final Logger LOG = new Logger(JdbcQueryTest.class);
   private static final String CONNECTION_TEMPLATE = 
"jdbc:avatica:remote:url=%s/druid/v2/sql/avatica/";
   private static final String TLS_CONNECTION_TEMPLATE =
       
"jdbc:avatica:remote:url=%s/druid/v2/sql/avatica/;truststore=%s;truststore_password=%s;keystore=%s;keystore_password=%s;key_password=%s";
 
   private static final String QUERY_TEMPLATE =
-      "SELECT \"user\", SUM(\"added\"), COUNT(*)" +
-      "FROM \"wikipedia\" " +
-      "WHERE \"__time\" >= CURRENT_TIMESTAMP - INTERVAL '99' YEAR AND 
\"language\" = %s" +
-      "GROUP BY 1 ORDER BY 3 DESC LIMIT 10";
-  private static final String QUERY = StringUtils.format(QUERY_TEMPLATE, 
"'en'");
-
-  private static final String QUERY_PARAMETERIZED = 
StringUtils.format(QUERY_TEMPLATE, "?");
+      "SELECT \"item\", SUM(\"value\"), COUNT(*) "
+      + "FROM \"%s\" "
+      + "WHERE \"__time\" >= CURRENT_TIMESTAMP - INTERVAL '99' YEAR AND 
\"value\" < %s \n"
+      + "GROUP BY 1 ORDER BY 3 DESC LIMIT 10";
 
   private String[] connections;
   private Properties connectionProperties;
 
-  @Inject
-  private IntegrationTestingConfig config;
-
-  @Inject
-  SSLClientConfig sslConfig;
-
-  @Inject
-  private DataLoaderHelper dataLoaderHelper;
+  private String tableName;
 
-  @BeforeMethod
-  public void before()
+  @Override
+  protected void beforeAll()
   {
     connectionProperties = new Properties();
     connectionProperties.setProperty("user", "admin");
     connectionProperties.setProperty("password", "priest");
     connections = new String[]{
-        StringUtils.format(CONNECTION_TEMPLATE, config.getRouterUrl()),
-        StringUtils.format(CONNECTION_TEMPLATE, config.getBrokerUrl()),
-        StringUtils.format(
-            TLS_CONNECTION_TEMPLATE,
-            config.getRouterTLSUrl(),
-            sslConfig.getTrustStorePath(),
-            sslConfig.getTrustStorePasswordProvider().getPassword(),
-            sslConfig.getKeyStorePath(),
-            sslConfig.getKeyStorePasswordProvider().getPassword(),
-            sslConfig.getKeyManagerPasswordProvider().getPassword()
-        ),
-        StringUtils.format(
-            TLS_CONNECTION_TEMPLATE,
-            config.getBrokerTLSUrl(),
-            sslConfig.getTrustStorePath(),
-            sslConfig.getTrustStorePasswordProvider().getPassword(),
-            sslConfig.getKeyStorePath(),
-            sslConfig.getKeyStorePasswordProvider().getPassword(),
-            sslConfig.getKeyManagerPasswordProvider().getPassword()
-        )
+        StringUtils.format(CONNECTION_TEMPLATE, getServerUrl(router)),
+        StringUtils.format(CONNECTION_TEMPLATE, getServerUrl(broker)),
+        // Add in the consecutive patch
+        // StringUtils.format(
+        //    TLS_CONNECTION_TEMPLATE,
+        //    config.getRouterTLSUrl(),
+        //    sslConfig.getTrustStorePath(),
+        //    sslConfig.getTrustStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyStorePath(),
+        //    sslConfig.getKeyStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyManagerPasswordProvider().getPassword()
+        //),
+        // StringUtils.format(
+        //    TLS_CONNECTION_TEMPLATE,
+        //    config.getBrokerTLSUrl(),
+        //    sslConfig.getTrustStorePath(),
+        //    sslConfig.getTrustStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyStorePath(),
+        //    sslConfig.getKeyStorePasswordProvider().getPassword(),
+        //    sslConfig.getKeyManagerPasswordProvider().getPassword()
+        // )

Review Comment:
   Please remove the commented out code. We can just add one line to the 
javadoc of this class that mentions that the TLS flavor for this test will be 
added later.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryTestBase.java:
##########
@@ -102,10 +111,46 @@ void setUp()
     }
   }
 
+  /**
+   * Ingests Druid with some the data from {@link 
MoreResources.Task#BASIC_INDEX} in a synchronous manner.

Review Comment:
   ```suggestion
      * Ingests test data using the task template {@link 
MoreResources.Task#BASIC_INDEX} in a synchronous manner.
   ```



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryTestBase.java:
##########
@@ -102,10 +111,46 @@ void setUp()
     }
   }
 
+  /**
+   * Ingests Druid with some the data from {@link 
MoreResources.Task#BASIC_INDEX} in a synchronous manner.
+   *
+   * @return ingested datasource name
+   */
+  protected String ingestBasicData()
+  {
+    String datasourceName = EmbeddedClusterApis.createTestDatasourceName();
+
+    final String taskId = IdUtils.getRandomId();
+    final IndexTask task = 
MoreResources.Task.BASIC_INDEX.get().dataSource(datasourceName).withId(taskId);
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
+    cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(datasourceName, 
coordinator, broker);
+    return datasourceName;
+  }
+
+  /**
+   * Execute an async SQL query against the given endpoint via the HTTP client.
+   */
+  protected ListenableFuture<StatusResponseHolder> executeQueryAsync(String 
endpoint, String query)
+  {
+    URL url;
+    try {
+      url = new URL(endpoint);
+    }
+    catch (MalformedURLException e) {
+      throw new AssertionError("Malformed URL");
+    }
+
+    Request request = new Request(HttpMethod.POST, url);
+    request.addHeader("Content-Type", MediaType.APPLICATION_JSON);
+    request.setContent(query.getBytes(StandardCharsets.UTF_8));
+    return httpClientRef.go(request, StatusResponseHandler.getInstance());

Review Comment:
   I am wondering if we couldn't just do the following instead:
   
   ```
   return overlord.bindings().anyBroker().submitSqlQuery(clientSqlQuery);
   ```
   
   this also returns a `ListenableFuture<String>`.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/JdbcQueryTest.java:
##########
@@ -145,84 +124,91 @@ public void testJdbcMetadata()
           druidTables.add(table);
         }
         LOG.info("'druid' schema tables %s", druidTables);
-        // maybe more tables than this, but at least should have these
-        Assert.assertTrue(
-            druidTables.containsAll(ImmutableList.of("twitterstream", 
"wikipedia", WIKIPEDIA_DATA_SOURCE))
+        // maybe more tables than this, but at least should have @tableName

Review Comment:
   ```suggestion
           // There maybe more tables than this, but at least should have 
@tableName
   ```



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryTestBase.java:
##########
@@ -150,4 +195,31 @@ protected void executeQuery(
         response.getContent().trim()
     );
   }
+
+  /**
+   * Execute a SQL query against the given endpoint via the HTTP client.
+   */
+  protected void cancelQuery(String endpoint, String queryId, 
Consumer<StatusResponseHolder> onResponse)

Review Comment:
   I think this adheres to the style used by the integration tests.
   We can update this to just return the `StatusResponseHolder` rather than 
accepting a consumer `onResponse`.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SqlQueryCancelTest.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
+import org.apache.druid.query.BaseQuery;
+import org.apache.druid.query.http.ClientSqlQuery;
+import org.jboss.netty.handler.codec.http.HttpResponseStatus;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class SqlQueryCancelTest extends QueryTestBase
+{
+  private static final String QUERY = " SELECT sleep(4) FROM %s LIMIT 4";
+
+  private ObjectMapper jsonMapper;
+  private String tableName;
+
+  @Override
+  public void beforeAll()
+  {
+    jsonMapper = overlord.bindings().jsonMapper();
+    tableName = ingestBasicData();
+  }
+
+  @Test
+  public void testCancelValidQuery() throws Exception
+  {
+    final String sqlQuery = StringUtils.format(QUERY, tableName);
+    final String queryId = "sql-cancel-test";
+    final ClientSqlQuery query = new ClientSqlQuery(
+        sqlQuery,
+        null,
+        false,
+        false,
+        false,
+        ImmutableMap.of(BaseQuery.SQL_QUERY_ID, queryId),
+        List.of()
+    );
+
+    ListenableFuture<StatusResponseHolder> f = 
executeQueryAsync(routerEndpoint, jsonMapper.writeValueAsString(query));
+
+    // Wait until the sqlLifecycle is authorized and registered
+    Thread.sleep(500L);
+    cancelQuery(
+        routerEndpoint,
+        queryId,
+        (r) -> Assertions.assertEquals(HttpResponseStatus.ACCEPTED, 
r.getStatus())
+    );
+
+    StatusResponseHolder srh = f.get();

Review Comment:
   Nit: Please use more descriptive var names, e.g. `responseHolder`.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SqlQueryCancelTest.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
+import org.apache.druid.query.BaseQuery;
+import org.apache.druid.query.http.ClientSqlQuery;
+import org.jboss.netty.handler.codec.http.HttpResponseStatus;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class SqlQueryCancelTest extends QueryTestBase
+{
+  private static final String QUERY = " SELECT sleep(4) FROM %s LIMIT 4";
+
+  private ObjectMapper jsonMapper;
+  private String tableName;
+
+  @Override
+  public void beforeAll()
+  {
+    jsonMapper = overlord.bindings().jsonMapper();
+    tableName = ingestBasicData();
+  }
+
+  @Test
+  public void testCancelValidQuery() throws Exception
+  {
+    final String sqlQuery = StringUtils.format(QUERY, tableName);
+    final String queryId = "sql-cancel-test";
+    final ClientSqlQuery query = new ClientSqlQuery(
+        sqlQuery,
+        null,
+        false,
+        false,
+        false,
+        ImmutableMap.of(BaseQuery.SQL_QUERY_ID, queryId),
+        List.of()
+    );
+
+    ListenableFuture<StatusResponseHolder> f = 
executeQueryAsync(routerEndpoint, jsonMapper.writeValueAsString(query));
+
+    // Wait until the sqlLifecycle is authorized and registered
+    Thread.sleep(500L);
+    cancelQuery(
+        routerEndpoint,
+        queryId,
+        (r) -> Assertions.assertEquals(HttpResponseStatus.ACCEPTED, 
r.getStatus())
+    );
+
+    StatusResponseHolder srh = f.get();
+    
Assertions.assertEquals(HttpResponseStatus.INTERNAL_SERVER_ERROR.getCode(), 
srh.getStatus().getCode());
+  }
+
+  @Test
+  public void testCancelInvalidQuery() throws Exception

Review Comment:
   ```suggestion
     public void test_cancelInvalidQuery_returnsNotFound() throws Exception
   ```



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryTestBase.java:
##########
@@ -102,10 +111,46 @@ void setUp()
     }
   }
 
+  /**
+   * Ingests Druid with some the data from {@link 
MoreResources.Task#BASIC_INDEX} in a synchronous manner.
+   *
+   * @return ingested datasource name
+   */
+  protected String ingestBasicData()
+  {
+    String datasourceName = EmbeddedClusterApis.createTestDatasourceName();
+
+    final String taskId = IdUtils.getRandomId();
+    final IndexTask task = 
MoreResources.Task.BASIC_INDEX.get().dataSource(datasourceName).withId(taskId);
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
+    cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(datasourceName, 
coordinator, broker);
+    return datasourceName;
+  }
+
+  /**
+   * Execute an async SQL query against the given endpoint via the HTTP client.
+   */
+  protected ListenableFuture<StatusResponseHolder> executeQueryAsync(String 
endpoint, String query)

Review Comment:
   It would probably be cleaner to pass the `ClientSqlQuery` instead of the 
serialized `String query` here. We could then perform the JSON serialization 
directly to bytes without going through the intermediate string.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SqlQueryCancelTest.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
+import org.apache.druid.query.BaseQuery;
+import org.apache.druid.query.http.ClientSqlQuery;
+import org.jboss.netty.handler.codec.http.HttpResponseStatus;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class SqlQueryCancelTest extends QueryTestBase

Review Comment:
   Thanks! The tests are much easier to follow now than the original 
`ITSqlQueryCancelTest`.



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SqlQueryCancelTest.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
+import org.apache.druid.query.BaseQuery;
+import org.apache.druid.query.http.ClientSqlQuery;
+import org.jboss.netty.handler.codec.http.HttpResponseStatus;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class SqlQueryCancelTest extends QueryTestBase
+{
+  private static final String QUERY = " SELECT sleep(4) FROM %s LIMIT 4";
+
+  private ObjectMapper jsonMapper;
+  private String tableName;
+
+  @Override
+  public void beforeAll()
+  {
+    jsonMapper = overlord.bindings().jsonMapper();
+    tableName = ingestBasicData();
+  }
+
+  @Test
+  public void testCancelValidQuery() throws Exception
+  {
+    final String sqlQuery = StringUtils.format(QUERY, tableName);
+    final String queryId = "sql-cancel-test";
+    final ClientSqlQuery query = new ClientSqlQuery(
+        sqlQuery,
+        null,
+        false,
+        false,
+        false,
+        ImmutableMap.of(BaseQuery.SQL_QUERY_ID, queryId),
+        List.of()
+    );
+
+    ListenableFuture<StatusResponseHolder> f = 
executeQueryAsync(routerEndpoint, jsonMapper.writeValueAsString(query));
+
+    // Wait until the sqlLifecycle is authorized and registered
+    Thread.sleep(500L);

Review Comment:
   Is this needed in the embedded test framework too?



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/JdbcQueryTest.java:
##########
@@ -145,84 +124,91 @@ public void testJdbcMetadata()
           druidTables.add(table);
         }
         LOG.info("'druid' schema tables %s", druidTables);
-        // maybe more tables than this, but at least should have these
-        Assert.assertTrue(
-            druidTables.containsAll(ImmutableList.of("twitterstream", 
"wikipedia", WIKIPEDIA_DATA_SOURCE))
+        // maybe more tables than this, but at least should have @tableName
+        Assertions.assertTrue(
+            druidTables.containsAll(ImmutableList.of(tableName))
         );
 
         Set<String> wikiColumns = new HashSet<>();
-        ResultSet columnsMetadata = metadata.getColumns("druid", "druid", 
WIKIPEDIA_DATA_SOURCE, null);
+        ResultSet columnsMetadata = metadata.getColumns("druid", "druid", 
tableName, null);
         while (columnsMetadata.next()) {
           final String column = columnsMetadata.getString(4);
           wikiColumns.add(column);
         }
-        LOG.info("'%s' columns %s", WIKIPEDIA_DATA_SOURCE, wikiColumns);
+        LOG.info("'%s' columns %s", tableName, wikiColumns);
         // a lot more columns than this, but at least should have these
-        Assert.assertTrue(
-            wikiColumns.containsAll(ImmutableList.of("added", "city", "delta", 
"language"))
+        Assertions.assertTrue(
+            wikiColumns.containsAll(ImmutableList.of("__time", "item", 
"value"))
         );
       }
       catch (SQLException throwables) {
-        Assert.fail(throwables.getMessage());
+        Assertions.fail(throwables.getMessage());
       }
     }
   }
 
   @Test
   public void testJdbcStatementQuery()
   {
+    String query = StringUtils.format(QUERY_TEMPLATE, tableName, "1000");
     for (String url : connections) {
       try (Connection connection = DriverManager.getConnection(url, 
connectionProperties)) {
         try (Statement statement = connection.createStatement()) {
-          final ResultSet resultSet = statement.executeQuery(QUERY);
+          final ResultSet resultSet = statement.executeQuery(query);
           int resultRowCount = 0;
           while (resultSet.next()) {
             resultRowCount++;
             LOG.info("%s,%s,%s", resultSet.getString(1), resultSet.getLong(2), 
resultSet.getLong(3));
           }
-          Assert.assertEquals(resultRowCount, 10);
+          Assertions.assertEquals(7, resultRowCount);
           resultSet.close();
         }
       }
       catch (SQLException throwables) {
-        Assert.fail(throwables.getMessage());
+        Assertions.fail(throwables.getMessage());
       }
     }
   }
 
   @Test
   public void testJdbcPrepareStatementQuery()
   {
+    String query = StringUtils.format(QUERY_TEMPLATE, tableName, "?");
     for (String url : connections) {
       try (Connection connection = DriverManager.getConnection(url, 
connectionProperties)) {
-        try (PreparedStatement statement = 
connection.prepareStatement(QUERY_PARAMETERIZED)) {
-          statement.setString(1, "en");
+        try (PreparedStatement statement = connection.prepareStatement(query)) 
{
+          statement.setLong(1, 1000);
           final ResultSet resultSet = statement.executeQuery();
           int resultRowCount = 0;
           while (resultSet.next()) {
             resultRowCount++;
             LOG.info("%s,%s,%s", resultSet.getString(1), resultSet.getLong(2), 
resultSet.getLong(3));
           }
-          Assert.assertEquals(resultRowCount, 10);
+          Assertions.assertEquals(7, resultRowCount);
           resultSet.close();
         }
       }
       catch (SQLException throwables) {
-        Assert.fail(throwables.getMessage());
+        Assertions.fail(throwables.getMessage());
       }
     }
   }
 
-  @Test(expectedExceptions = AvaticaSqlException.class, 
expectedExceptionsMessageRegExp = ".* No value bound for parameter \\(position 
\\[1]\\)")
-  public void testJdbcPrepareStatementQueryMissingParameters() throws 
SQLException
+  @Test
+  public void testJdbcPrepareStatementQueryMissingParameters()
   {
+    String query = StringUtils.format(QUERY_TEMPLATE, tableName, "?");
     for (String url : connections) {
       try (Connection connection = DriverManager.getConnection(url, 
connectionProperties);
-           PreparedStatement statement = 
connection.prepareStatement(QUERY_PARAMETERIZED);
+           PreparedStatement statement = connection.prepareStatement(query);
            ResultSet resultSet = statement.executeQuery()) {
         // This won't actually run as we expect the exception to be thrown 
before it gets here
         throw new IllegalStateException(resultSet.toString());

Review Comment:
   Should we use something like `Assertions.fail()` instead?



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SystemTableQueryTest.java:
##########
@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.indexing.common.task.IndexTask;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.testing.embedded.EmbeddedClusterApis;
+import org.apache.druid.testing.embedded.indexing.MoreResources;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SystemTableQueryTest extends QueryTestBase
+{
+  private final String tableName1 = 
EmbeddedClusterApis.createTestDatasourceName();
+  private final String tableName2 = 
EmbeddedClusterApis.createTestDatasourceName();
+
+  @Override
+  public void beforeAll()
+  {
+    final String taskId1 = IdUtils.getRandomId();
+    final String taskId2 = IdUtils.getRandomId();
+    final IndexTask task1 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName1).withId(taskId1);
+    final IndexTask task2 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName2).withId(taskId2);
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId1, task1));
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId2, task2));
+
+    cluster.callApi().waitForTaskToSucceed(taskId1, overlord);
+    cluster.callApi().waitForTaskToSucceed(taskId2, overlord);
+
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName1, coordinator, 
broker);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName2, coordinator, 
broker);

Review Comment:
   Seems like this can be simplified to:
   
   ```
   tableName1 = ingestBasicData();
   tableName2 = ingestBasicData();
   ```



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SystemTableQueryTest.java:
##########
@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.indexing.common.task.IndexTask;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.testing.embedded.EmbeddedClusterApis;
+import org.apache.druid.testing.embedded.indexing.MoreResources;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SystemTableQueryTest extends QueryTestBase
+{
+  private final String tableName1 = 
EmbeddedClusterApis.createTestDatasourceName();
+  private final String tableName2 = 
EmbeddedClusterApis.createTestDatasourceName();
+
+  @Override
+  public void beforeAll()
+  {
+    final String taskId1 = IdUtils.getRandomId();
+    final String taskId2 = IdUtils.getRandomId();
+    final IndexTask task1 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName1).withId(taskId1);
+    final IndexTask task2 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName2).withId(taskId2);
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId1, task1));
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId2, task2));
+
+    cluster.callApi().waitForTaskToSucceed(taskId1, overlord);
+    cluster.callApi().waitForTaskToSucceed(taskId2, overlord);
+
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName1, coordinator, 
broker);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName2, coordinator, 
broker);
+  }
+
+  @Test
+  public void testSystemTableQueries_segmentsCount()
+  {
+    String query = StringUtils.format(
+        "SELECT datasource, count(*) \n"
+        + "FROM sys.segments \n"
+        + "WHERE datasource='%s' \n"
+        + "OR    datasource='%s' \n"
+        + "GROUP BY 1", tableName1, tableName2

Review Comment:
   formatting:
   ```suggestion
           + "GROUP BY 1",
           tableName1, tableName2
   ```



##########
embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/SystemTableQueryTest.java:
##########
@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.query;
+
+
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.indexing.common.task.IndexTask;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.testing.embedded.EmbeddedClusterApis;
+import org.apache.druid.testing.embedded.indexing.MoreResources;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SystemTableQueryTest extends QueryTestBase
+{
+  private final String tableName1 = 
EmbeddedClusterApis.createTestDatasourceName();
+  private final String tableName2 = 
EmbeddedClusterApis.createTestDatasourceName();
+
+  @Override
+  public void beforeAll()
+  {
+    final String taskId1 = IdUtils.getRandomId();
+    final String taskId2 = IdUtils.getRandomId();
+    final IndexTask task1 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName1).withId(taskId1);
+    final IndexTask task2 = 
MoreResources.Task.BASIC_INDEX.get().dataSource(tableName2).withId(taskId2);
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId1, task1));
+    cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId2, task2));
+
+    cluster.callApi().waitForTaskToSucceed(taskId1, overlord);
+    cluster.callApi().waitForTaskToSucceed(taskId2, overlord);
+
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName1, coordinator, 
broker);
+    cluster.callApi().waitForAllSegmentsToBeAvailable(tableName2, coordinator, 
broker);
+  }
+
+  @Test
+  public void testSystemTableQueries_segmentsCount()
+  {
+    String query = StringUtils.format(
+        "SELECT datasource, count(*) \n"

Review Comment:
   Nit: Are the `\n` required?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]


Reply via email to