Wire secondary-schema stores into the execution context at connection time

TransactionBoundDatabase now populates storesBySchemaName from the
additional stores map on RecordStoreAndRecordContextTransaction, so that
loadRecordStore(schemaId) can return the secondary FDBRecordStoreBase
that RecordQueryStoreBindingPlan looks up at execution time.

Author: arnaud-lacurie

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/RecordStoreAndRecordContextTransaction.java

@@ -33,6 +33,8 @@
 import com.google.protobuf.Message;
 
 import javax.annotation.Nonnull;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -45,6 +47,7 @@
 public class RecordStoreAndRecordContextTransaction implements Transaction {
     FDBRecordStoreBase<Message> store;
     RecordContextTransaction transaction;
+    Map<String, FDBRecordStoreBase<Message>> additionalStores = new HashMap<>();
 
     /**
      * the schema template this transaction is bound to. This is mainly needed when accessing the plan cache
@@ -105,6 +108,18 @@ public FDBRecordStoreBase<Message> getRecordStore() {
         return store;
     }
 
+    @Nonnull
+    public Map<String, FDBRecordStoreBase<Message>> getAdditionalStores() {
+        return additionalStores;
+    }
+
+    @Nonnull
+    public RecordStoreAndRecordContextTransaction withAdditionalStore(@Nonnull String schemaName,
+                                                                      @Nonnull FDBRecordStoreBase<Message> additionalStore) {
+        this.additionalStores.put(schemaName, additionalStore);
+        return this;
+    }
+
     @Nonnull
     public SchemaTemplate getBoundSchemaTemplate() {
         return boundSchemaTemplate;

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/catalog/TransactionBoundDatabase.java

@@ -49,6 +49,8 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * There can only be 1 Database object per Connection instance, and its lifecycle is managed by the connection
@@ -66,6 +68,7 @@ public class TransactionBoundDatabase extends AbstractDatabase {
     @Nullable
     private final KeySpace keySpace;
     BackingStore store;
+    Map<String, BackingStore> storesBySchemaName = new HashMap<>();
     URI uri;
 
     private static final MetadataOperationsFactory onlyTemporaryFunctionOperationsFactory = new AbstractMetadataOperationsFactory() {
@@ -97,6 +100,9 @@ public RelationalConnection connect(@Nullable Transaction transaction) throws Re
         }
         final var recordStoreAndRecordContextTx = transaction.unwrap(RecordStoreAndRecordContextTransaction.class);
         store = BackingRecordStore.fromTransactionWithStore(recordStoreAndRecordContextTx);
+        storesBySchemaName.clear();
+        recordStoreAndRecordContextTx.getAdditionalStores().forEach((schemaName, additionalStore) ->
+                storesBySchemaName.put(schemaName, BackingRecordStore.fromTransactionAndStore(recordStoreAndRecordContextTx, additionalStore)));
         final var boundSchemaTemplate = recordStoreAndRecordContextTx.getBoundSchemaTemplate();
         EmbeddedRelationalConnection connection = new EmbeddedRelationalConnection(this, new HollowStoreCatalog(boundSchemaTemplate, keySpace),
                 ((RecordStoreAndRecordContextTransaction) transaction).getRecordContextTransaction(), options);
@@ -106,7 +112,7 @@ public RelationalConnection connect(@Nullable Transaction transaction) throws Re
 
     @Override
     public BackingStore loadRecordStore(@Nonnull String schemaId, @Nonnull FDBRecordStoreBase.StoreExistenceCheck existenceCheck) {
-        return store;
+        return storesBySchemaName.getOrDefault(schemaId, store);
     }
 
     @Override

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/storage/BackingRecordStore.java

@@ -221,6 +221,10 @@ public static BackingRecordStore fromTransactionWithStore(@Nonnull RecordStoreAn
         return new BackingRecordStore(txn, txn.getRecordStore());
     }
 
+    public static BackingRecordStore fromTransactionAndStore(@Nonnull Transaction txn, @Nonnull FDBRecordStoreBase<Message> store) {
+        return new BackingRecordStore(txn, store);
+    }
+
     @SuppressWarnings("PMD.PreserveStackTrace")
     public static BackingRecordStore load(@Nonnull Transaction txn, @Nonnull StoreConfig config, @Nonnull FDBRecordStoreBase.StoreExistenceCheck existenceCheck) throws RelationalException {
         //TODO(bfines) error handling if this store doesn't exist