Multi-Cluster support for yamsql & support not-registering the driver with FRL

fdb-relational-server/src/main/java/com/apple/foundationdb/relational/server/FRL.java

@@ -79,7 +79,7 @@
 @API(API.Status.EXPERIMENTAL)
 public class FRL implements AutoCloseable {
     private final FdbConnection fdbDatabase;
-    private final RelationalDriver registeredDriver;
+    private final RelationalDriver driver;
     private boolean registeredJDBCEmbedDriver;
 
     public FRL() throws RelationalException {
@@ -91,6 +91,10 @@ public FRL(@Nonnull Options options) throws RelationalException {
     }
 
     public FRL(@Nonnull Options options, @Nullable String clusterFile) throws RelationalException {
+        this(options, clusterFile, true);
+    }
+
+    public FRL(@Nonnull Options options, @Nullable String clusterFile, boolean registerDriver) throws RelationalException {
         final FDBDatabase fdbDb = FDBDatabaseFactory.instance().getDatabase(clusterFile);
         final Long asyncToSyncTimeout = options.getOption(Options.Name.ASYNC_OPERATIONS_TIMEOUT_MILLIS);
         if (asyncToSyncTimeout > 0) {
@@ -114,7 +118,7 @@ public FRL(@Nonnull Options options, @Nullable String clusterFile) throws Relati
                 .setStoreCatalog(storeCatalog).build();
 
         try {
-            this.registeredDriver = new EmbeddedRelationalDriver(RecordLayerEngine.makeEngine(
+            this.driver = new EmbeddedRelationalDriver(RecordLayerEngine.makeEngine(
                     rlConfig,
                     Collections.singletonList(fdbDb),
                     keySpace,
@@ -130,13 +134,19 @@ public FRL(@Nonnull Options options, @Nullable String clusterFile) throws Relati
                             .setTertiarySize(options.getOption(Options.Name.PLAN_CACHE_TERTIARY_MAX_ENTRIES))
                             .build()));
 
-            DriverManager.registerDriver(this.registeredDriver);
-            this.registeredJDBCEmbedDriver = true;
+            if (registerDriver) {
+                DriverManager.registerDriver(this.driver);
+                this.registeredJDBCEmbedDriver = true;
+            }
         } catch (SQLException ve) {
             throw new RelationalException(ve);
         }
     }
 
+    public RelationalDriver getDriver() {
+        return driver;
+    }
+
     @SuppressWarnings("AbbreviationAsWordInName") // allow JDBCURI, though perhaps we should update this to make it clearer
     private static String createEmbeddedJDBCURI(String database, String schema)  {
         return EmbeddedRelationalDriver.JDBC_URL_PREFIX + database + (schema != null ? "?schema=" + schema : "");
@@ -204,7 +214,6 @@ public Response execute(String database, String schema, String sql, List<Paramet
     }
 
     private RelationalConnection connect(String database, String schema, Options options) throws SQLException {
-        final var driver = (RelationalDriver) DriverManager.getDriver(createEmbeddedJDBCURI(database, schema));
         return driver.connect(URI.create(createEmbeddedJDBCURI(database, schema)), options);
     }
 
@@ -331,7 +340,6 @@ public RelationalResultSet scan(String database, String schema, String tableName
     }
 
     public TransactionalToken createTransactionalToken(String database, String schema, Options options) throws SQLException {
-        final var driver = (RelationalDriver) DriverManager.getDriver(createEmbeddedJDBCURI(database, schema));
         RelationalConnection transactionalConnection = driver.connect(URI.create(createEmbeddedJDBCURI(database, schema)), options);
         transactionalConnection.setAutoCommit(false);
         return new TransactionalToken(transactionalConnection);
@@ -395,7 +403,7 @@ public void close() throws SQLException, RelationalException {
         }
         // We registered the Relational embed driver... cleanup.
         if (this.registeredJDBCEmbedDriver) {
-            DriverManager.deregisterDriver(registeredDriver);
+            DriverManager.deregisterDriver(driver);
         }
     }
 }

fdb-test-utils/src/main/java/com/apple/foundationdb/test/FDBTestEnvironment.java

@@ -26,6 +26,7 @@
 import javax.annotation.Nullable;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -73,6 +74,12 @@ public static List<String> allClusterFiles() {
         return clusterFiles;
     }
 
+    public static List<String> allClusterFilesInRandomOrder() {
+        final List<String> randomized = new ArrayList<>(clusterFiles);
+        Collections.shuffle(randomized);
+        return randomized;
+    }
+
     public static String randomClusterFile() {
         return clusterFiles.get(ThreadLocalRandom.current().nextInt(clusterFiles.size()));
     }

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/ConnectionTarget.java

@@ -0,0 +1,70 @@
+/*
+ * ConnectionTarget.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2026 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb.relational.yamltests;
+
+import javax.annotation.Nonnull;
+import java.net.URI;
+
+/**
+ * A resolved connection target consisting of a URI and a cluster index.
+ *
+ * <p>The cluster index identifies which FDB cluster to connect to. Index 0 is the default.
+ * Additional clusters can be accessed in YAMSQL files using the map form
+ * of the {@code connect} directive:
+ * <pre>{@code
+ * connect: { cluster: 1, uri: 0 }
+ * }</pre>
+ * The uri component can be a few different things:
+ * <ul>
+ *     <li>An index, in which case {@code 0} is the catalog, and other positive numbers refer to the schemas created
+ *     automatically with the {@code schema_template:} block</li>
+ *     <li>A fully qualified uri such as {@code "jdbc:embed:/FRL/MCI_DB?schema=S1"} or
+ *     {@code "jdbc:embed:/__SYS?schema=CATALOG"}. The scheme should always be {@code jdbc:embed:}, and the framework
+ *     will update it to control whether it goes to the embedded connection, or one of the servers.</li>
+ * </ul>
+ */
+public class ConnectionTarget {
+    @Nonnull
+    private final URI uri;
+    private final int clusterIndex;
+
+    public ConnectionTarget(@Nonnull URI uri, int clusterIndex) {
+        this.uri = uri;
+        this.clusterIndex = clusterIndex;
+    }
+
+    @Nonnull
+    public URI getUri() {
+        return uri;
+    }
+
+    public int getClusterIndex() {
+        return clusterIndex;
+    }
+
+    @Override
+    public String toString() {
+        if (clusterIndex == 0) {
+            return uri.toString();
+        }
+        return uri + " [cluster=" + clusterIndex + "]";
+    }
+}

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/YamlConnectionFactory.java

@@ -33,15 +33,16 @@
  */
 public interface YamlConnectionFactory {
     /**
-     * Convert a connection uri into an actual connection.
+     * Convert a connection uri into an actual connection on a specific cluster.
      *
      * @param connectPath the path to connect to
+     * @param clusterIndex the cluster to connect to (0 is the default cluster)
      *
-     * @return A new {@link RelationalConnection} for the given path appropriate for this test class
+     * @return A new {@link RelationalConnection} for the given path on the specified cluster
      *
-     * @throws SQLException if we cannot connect
+     * @throws SQLException if we cannot connect or the cluster index is not supported
      */
-    YamlConnection getNewConnection(@Nonnull URI connectPath) throws SQLException;
+    YamlConnection getNewConnection(@Nonnull URI connectPath, int clusterIndex) throws SQLException;
 
     /**
      * The versions that the connection has, other than the current code.
@@ -66,4 +67,13 @@ public interface YamlConnectionFactory {
     default boolean isMultiServer() {
         return false;
     }
+
+    /**
+     * Returns the number of clusters available for testing.
+     *
+     * @return the number of available clusters (1 means only the default cluster)
+     */
+    default int getAvailableClusterCount() {
+        return 1;
+    }
 }

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/YamlConnectionFactoryWithOptions.java

@@ -42,8 +42,8 @@ public YamlConnectionFactoryWithOptions(@Nonnull final YamlConnectionFactory und
     }
 
     @Override
-    public YamlConnection getNewConnection(@Nonnull final URI connectPath) throws SQLException {
-        final var connection = underlying.getNewConnection(connectPath);
+    public YamlConnection getNewConnection(@Nonnull final URI connectPath, int clusterIndex) throws SQLException {
+        final var connection = underlying.getNewConnection(connectPath, clusterIndex);
         connection.setConnectionOptions(options);
         return connection;
     }
@@ -53,6 +53,11 @@ public Set<SemanticVersion> getVersionsUnderTest() {
         return underlying.getVersionsUnderTest();
     }
 
+    @Override
+    public int getAvailableClusterCount() {
+        return underlying.getAvailableClusterCount();
+    }
+
     @Override
     public boolean isMultiServer() {
         return underlying.isMultiServer();

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/YamlExecutionContext.java

@@ -286,29 +286,50 @@ public void registerConnectionURI(@Nonnull YamlReference.YamlResource resource,
     }
 
     /**
-     * Infers the URI of the database to which a block should connect to.
+     * Infers the connection target (URI and cluster index) for a block.
      * <br>
-     * A block can declare a connection in multiple ways:
-     * 1. no explicit declaration: Try to connect to the only registered connection URI in the local {@link YamlReference.YamlResource}.
+     * <ul>
+     *   <li>no explicit declaration: Try to connect to the only registered connection URI in the local {@link YamlReference.YamlResource}.
      *    If not, try to connect to the only connection across all parent resources.
      *    A URI can be registered by defining a "schema_template" block before that, which sets up the database and schema for a provided schema template.
-     * 2. Parameter 0: connects to the system tables (catalog).
-     * 3. Parameter One-based Number: connects to the registered connection URI, number denotes the sequence of definitions in the local YamlResource.
+     *    </li>
+     *    <li>Parameter 0: connects to the system tables (catalog). </li>
+     *    <li>Parameter One-based Number: connects to the registered connection URI, number denotes the sequence of definitions in the local YamlResource.
      *    To access parent connection URIs, this number should be prepended by `(global)` tag.
-     * 4. Parameter String: connects to the defined String
+     *    </li>
+     *    <li>Parameter String: connects to the defined String</li>
+     *    <li>A map form for specifying the cluster:
+     * <pre>{@code
+     * connect: { cluster: 1, uri: 0 }
+     * connect: { cluster: 1 }
+     * }</pre>
+     * </li>
+     * </ul>
      *
-     * @param connectObject can be {@code null}, an {@link Integer} value or a {@link String}.
+     * @param connectObject can be {@code null}, an {@link Integer}, a {@link String}, or a {@link Map} with
+     *                      optional {@code cluster} and {@code uri} keys.
      *
-     * @return a valid connection URI
+     * @return a valid connection target
      */
-    public URI inferConnectionURI(@Nonnull final YamlReference.YamlResource resource, @Nullable Object connectObject) {
+    public ConnectionTarget inferConnectionTarget(@Nonnull final YamlReference.YamlResource resource, @Nullable Object connectObject) {
         Assert.thatUnchecked(registeredResources.contains(resource), "A YamlResource should be registered before registering available connection URIs");
+        if (connectObject instanceof Map) {
+            final Map<?, ?> connectMap = CustomYamlConstructor.LinedObject.unlineKeys(Matchers.map(connectObject, "connect"));
+            final int clusterIndex = connectMap.containsKey("cluster")
+                    ? ((Number) connectMap.get("cluster")).intValue() : 0;
+            final Object uriSpec = connectMap.getOrDefault("uri", null);
+            return new ConnectionTarget(resolveConnectionURI(resource, uriSpec), clusterIndex);
+        }
+        return new ConnectionTarget(resolveConnectionURI(resource, connectObject), 0);
+    }
+
+    private URI resolveConnectionURI(@Nonnull final YamlReference.YamlResource resource, @Nullable Object connectObject) {
         if (connectObject == null) {
             return getConnectionFromConnectionURIList(resource, true, -1, true);
         } else if (connectObject instanceof Integer) {
             return getConnectionFromConnectionURIList(resource, false, (Integer) connectObject, false);
         } else {
-            final var stringURI = Matchers.string(connectObject);
+            final var stringURI = Matchers.string(connectObject, "connection object");
             if (stringURI.startsWith("(global)")) {
                 return getConnectionFromConnectionURIList(resource, false, Integer.parseInt(stringURI.substring(8).trim()), true);
             }