[4.x] Schema queries with configurable server-side timeouts

core/src/main/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/CassandraSchemaQueries.java

@@ -15,6 +15,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+/*
+ * Copyright (C) 2024 ScyllaDB
+ *
+ * Modified by ScyllaDB
+ */
 package com.datastax.oss.driver.internal.core.metadata.schema.queries;
 
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
@@ -54,6 +59,7 @@ public abstract class CassandraSchemaQueries implements SchemaQueries {
   // The future we return from execute, completes when all the queries are done.
   private final CompletableFuture<SchemaRows> schemaRowsFuture = new CompletableFuture<>();
   private final long startTimeNs = System.nanoTime();
+  private final String usingTimeoutClause;
 
   // All non-final fields are accessed exclusively on adminExecutor
   private CassandraSchemaRows.Builder schemaRowsBuilder;
@@ -73,6 +79,10 @@ protected CassandraSchemaQueries(
             DefaultDriverOption.METADATA_SCHEMA_REFRESHED_KEYSPACES, Collections.emptyList());
     assert refreshedKeyspaces != null; // per the default value
     this.keyspaceFilter = KeyspaceFilter.newInstance(logPrefix, refreshedKeyspaces);
+    this.usingTimeoutClause =
+        " USING TIMEOUT "
+            + config.getDuration(DefaultDriverOption.METADATA_SCHEMA_REQUEST_TIMEOUT).toMillis()
+            + "ms";
   }
 
   protected abstract String selectKeyspacesQuery();
@@ -112,29 +122,47 @@ private void executeOnAdminExecutor() {
 
     schemaRowsBuilder = new CassandraSchemaRows.Builder(node, keyspaceFilter, logPrefix);
     String whereClause = keyspaceFilter.getWhereClause();
+    String usingClause = shouldApplyUsingTimeout() ? usingTimeoutClause : "";
 
-    query(selectKeyspacesQuery() + whereClause, schemaRowsBuilder::withKeyspaces);
-    query(selectTypesQuery() + whereClause, schemaRowsBuilder::withTypes);
-    query(selectTablesQuery() + whereClause, schemaRowsBuilder::withTables);
-    query(selectColumnsQuery() + whereClause, schemaRowsBuilder::withColumns);
+    query(selectKeyspacesQuery() + whereClause + usingClause, schemaRowsBuilder::withKeyspaces);
+    query(selectTypesQuery() + whereClause + usingClause, schemaRowsBuilder::withTypes);
+    query(selectTablesQuery() + whereClause + usingClause, schemaRowsBuilder::withTables);
+    query(selectColumnsQuery() + whereClause + usingClause, schemaRowsBuilder::withColumns);
     selectIndexesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withIndexes));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withIndexes));
     selectViewsQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withViews));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withViews));
     selectFunctionsQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withFunctions));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withFunctions));
     selectAggregatesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withAggregates));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withAggregates));
     selectVirtualKeyspacesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualKeyspaces));
+        .ifPresent(
+            select ->
+                query(select + whereClause + usingClause, schemaRowsBuilder::withVirtualKeyspaces));
     selectVirtualTablesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualTables));
+        .ifPresent(
+            select ->
+                query(select + whereClause + usingClause, schemaRowsBuilder::withVirtualTables));
     selectVirtualColumnsQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualColumns));
+        .ifPresent(
+            select ->
+                query(select + whereClause + usingClause, schemaRowsBuilder::withVirtualColumns));
     selectEdgesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withEdges));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withEdges));
     selectVerticiesQuery()
-        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVertices));
+        .ifPresent(
+            select -> query(select + whereClause + usingClause, schemaRowsBuilder::withVertices));
+  }
+
+  protected boolean shouldApplyUsingTimeout() {
+    // We use non-null sharding info as a proxy check for cluster being a ScyllaDB cluster
+    return (channel.getShardingInfo() != null);
   }
 
   private void query(

core/src/test/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/Cassandra3SchemaQueriesTest.java

@@ -15,6 +15,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+/*
+ * Copyright (C) 2024 ScyllaDB
+ *
+ * Modified by ScyllaDB
+ */
 package com.datastax.oss.driver.internal.core.metadata.schema.queries;
 
 import static com.datastax.oss.driver.Assertions.assertThat;
@@ -28,6 +33,7 @@
 import com.datastax.oss.driver.internal.core.adminrequest.AdminResult;
 import com.datastax.oss.driver.internal.core.channel.DriverChannel;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Queue;
 import java.util.concurrent.CompletionStage;
@@ -63,52 +69,72 @@ public void should_query_with_keyspace_filter() {
     should_query_with_where_clause(" WHERE keyspace_name IN ('ks1','ks2')");
   }
 
+  @Test
+  public void should_query_with_using_clause() {
+    when(config.getDuration(DefaultDriverOption.METADATA_SCHEMA_REQUEST_TIMEOUT))
+        .thenReturn(Duration.ofMillis(100));
+    should_query_with_clauses("", " USING TIMEOUT 100ms");
+  }
+
   private void should_query_with_where_clause(String whereClause) {
+    should_query_with_clauses(whereClause, "");
+  }
+
+  private void should_query_with_clauses(String whereClause, String usingClause) {
     SchemaQueriesWithMockedChannel queries =
-        new SchemaQueriesWithMockedChannel(driverChannel, node, config, "test");
+        new SchemaQueriesWithMockedChannel(
+            driverChannel, node, config, "test", !usingClause.equals(""));
     CompletionStage<SchemaRows> result = queries.execute();
 
     // Keyspace
     Call call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.keyspaces" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.keyspaces" + whereClause + usingClause);
     call.result.complete(
         mockResult(mockRow("keyspace_name", "ks1"), mockRow("keyspace_name", "ks2")));
 
     // Types
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.types" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.types" + whereClause + usingClause);
     call.result.complete(mockResult(mockRow("keyspace_name", "ks1", "type_name", "type")));
 
     // Tables
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.tables" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.tables" + whereClause + usingClause);
     call.result.complete(mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo")));
 
     // Columns
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.columns" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.columns" + whereClause + usingClause);
     call.result.complete(
         mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "k")));
 
     // Indexes
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.indexes" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.indexes" + whereClause + usingClause);
     call.result.complete(
         mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "index_name", "index")));
 
     // Views
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.views" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.views" + whereClause + usingClause);
     call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "view_name", "foo")));
 
     // Functions
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.functions" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.functions" + whereClause + usingClause);
     call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "function_name", "add")));
 
     // Aggregates
     call = queries.calls.poll();
-    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.aggregates" + whereClause);
+    assertThat(call.query)
+        .isEqualTo("SELECT * FROM system_schema.aggregates" + whereClause + usingClause);
     call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "aggregate_name", "add")));
 
     channel.runPendingTasks();
@@ -349,10 +375,26 @@ public void should_abort_if_query_fails() {
   static class SchemaQueriesWithMockedChannel extends Cassandra3SchemaQueries {
 
     final Queue<Call> calls = new LinkedBlockingDeque<>();
+    final boolean shouldApplyUsingTimeout;
 
     SchemaQueriesWithMockedChannel(
         DriverChannel channel, Node node, DriverExecutionProfile config, String logPrefix) {
+      this(channel, node, config, logPrefix, false);
+    }
+
+    SchemaQueriesWithMockedChannel(
+        DriverChannel channel,
+        Node node,
+        DriverExecutionProfile config,
+        String logPrefix,
+        boolean shouldApplyUsingTimeout) {
       super(channel, node, config, logPrefix);
+      this.shouldApplyUsingTimeout = shouldApplyUsingTimeout;
+    }
+
+    @Override
+    protected boolean shouldApplyUsingTimeout() {
+      return shouldApplyUsingTimeout;
     }
 
     @Override

core/src/test/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/DefaultSchemaQueriesFactoryTest.java

@@ -15,6 +15,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+/*
+ * Copyright (C) 2024 ScyllaDB
+ *
+ * Modified by ScyllaDB
+ */
 package com.datastax.oss.driver.internal.core.metadata.schema.queries;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -23,6 +28,7 @@
 
 import com.datastax.dse.driver.api.core.metadata.DseNodeProperties;
 import com.datastax.oss.driver.api.core.Version;
+import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverConfig;
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.metadata.Node;
@@ -33,6 +39,8 @@
 import com.tngtech.java.junit.dataprovider.DataProvider;
 import com.tngtech.java.junit.dataprovider.DataProviderRunner;
 import com.tngtech.java.junit.dataprovider.UseDataProvider;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.Optional;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -130,6 +138,8 @@ private DefaultSchemaQueriesFactory buildFactory() {
     when(mockConfig.getDefaultProfile()).thenReturn(mockProfile);
     final InternalDriverContext mockInternalCtx = mock(InternalDriverContext.class);
     when(mockInternalCtx.getConfig()).thenReturn(mockConfig);
+    when(mockProfile.getDuration(DefaultDriverOption.METADATA_SCHEMA_REQUEST_TIMEOUT))
+        .thenReturn(Duration.of(5, ChronoUnit.SECONDS));
 
     return new DefaultSchemaQueriesFactory(mockInternalCtx);
   }

core/src/test/java/com/datastax/oss/driver/internal/core/metadata/schema/queries/SchemaQueriesTest.java

@@ -15,6 +15,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+/*
+ * Copyright (C) 2024 ScyllaDB
+ *
+ * Modified by ScyllaDB
+ */
 package com.datastax.oss.driver.internal.core.metadata.schema.queries;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -54,7 +59,7 @@ public abstract class SchemaQueriesTest {
   public void setup() {
     // Whatever, not actually used because the requests are mocked
     when(config.getDuration(DefaultDriverOption.METADATA_SCHEMA_REQUEST_TIMEOUT))
-        .thenReturn(Duration.ZERO);
+        .thenReturn(Duration.ofSeconds(2));
     when(config.getInt(DefaultDriverOption.METADATA_SCHEMA_REQUEST_PAGE_SIZE)).thenReturn(5000);
 
     channel = new EmbeddedChannel();