feat: Support for using multiple databases in datastore

src/index.ts

@@ -66,6 +66,7 @@ import {Transaction} from './transaction';
 import {promisifyAll} from '@google-cloud/promisify';
 import {google} from '../protos/protos';
 import {AggregateQuery} from './aggregate';
+import {addDatabaseIdToRequest} from './util';
 
 const {grpc} = new GrpcClient();
 
@@ -699,6 +700,16 @@ class Datastore extends DatastoreRequest {
     );
   }
 
+  /**
+   * Gets the database id that all requests will be run against.
+   *
+   * @returns {string} The database id that the current client is set to that
+   *    requests will run against.
+   */
+  getDatabaseId(): string | undefined {
+    return this.options.databaseId;
+  }
+
   getProjectId(): Promise<string> {
     return this.auth.getProjectId();
   }
@@ -1817,7 +1828,9 @@ promisifyAll(Datastore, {
     'double',
     'isDouble',
     'geoPoint',
+    'getDatabaseId',
     'getProjectId',
+    'getRequestWithDatabaseId',
     'getSharedQueryOptions',
     'isGeoPoint',
     'index',
@@ -1898,6 +1911,7 @@ export interface DatastoreOptions extends GoogleAuthOptions {
   namespace?: string;
   apiEndpoint?: string;
   sslCreds?: ChannelCredentials;
+  databaseId?: string;
 }
 
 export interface KeyToLegacyUrlSafeCallback {

src/request.ts

@@ -57,6 +57,7 @@ import {
 import {Datastore} from '.';
 import ITimestamp = google.protobuf.ITimestamp;
 import {AggregateQuery} from './aggregate';
+import {addDatabaseIdToRequest} from './util';
 
 /**
  * A map of read consistency values to proto codes.
@@ -995,6 +996,8 @@ class DatastoreRequest {
       }
     }
 
+    addDatabaseIdToRequest(datastore, reqOpts);
+
     if (method === 'rollback') {
       reqOpts.transaction = this.id;
     }
@@ -1155,6 +1158,7 @@ export interface SharedQueryOptions {
   };
 }
 export interface RequestOptions extends SharedQueryOptions {
+  databaseId?: string;
   mutations?: google.datastore.v1.IMutation[];
   keys?: Entity;
   transactionOptions?: {

src/util.ts

@@ -0,0 +1,19 @@
+// Copyright 2023 Google LLC
+//
+// 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
+//
+//     https://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.
+
+export function addDatabaseIdToRequest(datastore: any, reqOpts: any) {
+  if (datastore.options && datastore.options.databaseId) {
+    reqOpts.databaseId = datastore.options.databaseId;
+  }
+}

system-test/datastore.ts

@@ -21,9 +21,11 @@ import {Datastore, Index} from '../src';
 import {google} from '../protos/protos';
 import {Storage} from '@google-cloud/storage';
 import {AggregateField} from '../src/aggregate';
-import {PropertyFilter, EntityFilter, and, or} from '../src/filter';
 import {entity} from '../src/entity';
 import KEY_SYMBOL = entity.KEY_SYMBOL;
+import {PropertyFilter, and, or} from '../src/filter';
+
+const SECOND_DATABASE_ID = 'multidb-test';
 
 describe('Datastore', () => {
   const testKinds: string[] = [];
@@ -312,6 +314,115 @@ describe('Datastore', () => {
       await datastore.delete(postKey);
     });
 
+    describe('multi-db support for read and write operations', () => {
+      it('should run a query with another database', async () => {
+        // First verify that a record gets written to datastore
+        const postKey = datastore.key(['Post', 'post1']);
+        await datastore.save({key: postKey, data: post});
+        const query = datastore.createQuery('Post').hasAncestor(postKey);
+        const [defaultDatastoreResults] = await datastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 1);
+        const [entity] = await datastore.get(postKey);
+        assert.strictEqual(entity.author, 'Silvano');
+        // With another database, verify that a query returns no results
+        const otherDatastore = new Datastore({
+          namespace: `${Date.now()}`,
+          databaseId: SECOND_DATABASE_ID,
+        });
+        const [secondDatastoreResults] = await otherDatastore.runQuery(query);
+        assert.strictEqual(secondDatastoreResults.length, 0);
+        const [otherEntity] = await otherDatastore.get(postKey);
+        assert(typeof otherEntity === 'undefined');
+        // Cleanup
+        await datastore.delete(postKey);
+      });
+      it('should ensure save works with another database', async () => {
+        // First verify that the default database is empty
+        const postKey = datastore.key(['Post', 'post1']);
+        const query = datastore.createQuery('Post').hasAncestor(postKey);
+        const [defaultDatastoreResults] = await datastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 0);
+        const [originalSecondaryResults] = await datastore.runQuery(query);
+        assert.strictEqual(originalSecondaryResults.length, 0);
+        const [entity] = await datastore.get(postKey);
+        assert(typeof entity === 'undefined');
+        // With another database, verify that saving to the database works
+        const otherDatastore = new Datastore({
+          namespace: `${Date.now()}`,
+          databaseId: SECOND_DATABASE_ID,
+        });
+        await otherDatastore.save({key: postKey, data: post});
+        const [secondDatastoreResults] = await otherDatastore.runQuery(query);
+        assert.strictEqual(secondDatastoreResults.length, 1);
+        const [originalResults] = await datastore.runQuery(query);
+        assert.strictEqual(originalResults.length, 0);
+        const [otherEntity] = await otherDatastore.get(postKey);
+        assert.strictEqual(otherEntity.author, 'Silvano');
+        // Cleanup
+        await otherDatastore.delete(postKey);
+      });
+      it('should ensure save respects the databaseId parameter per key', async () => {
+        // First write entities to the database by specifying the database in the key
+        const otherDatastore = new Datastore({
+          namespace: `${Date.now()}`,
+          databaseId: SECOND_DATABASE_ID,
+        });
+        const dataD1 = Object.assign({}, post);
+        dataD1.author = 'D1';
+        const dataS1 = Object.assign({}, post);
+        dataS1.author = 'S1';
+        const dataS2 = Object.assign({}, post);
+        dataS2.author = 'S2';
+        const dataS3 = Object.assign({}, post);
+        dataS3.author = 'S3';
+        const postKeyDefault1 = datastore.key(['Post', 'postD1']);
+        const postKeySecondary1 = otherDatastore.key(['Post', 'postS1']);
+        const postKeySecondary2 = otherDatastore.key(['Post', 'postS2']);
+        const postKeySecondary3 = otherDatastore.key(['Post', 'postS3']);
+        await datastore.save({key: postKeyDefault1, data: dataD1});
+        await otherDatastore.save({key: postKeySecondary1, data: dataS1});
+        await otherDatastore.save({key: postKeySecondary2, data: dataS2});
+        await otherDatastore.save({key: postKeySecondary3, data: dataS3});
+        // Next, ensure that the default database has the right records
+        const query = datastore
+          .createQuery('Post')
+          .hasAncestor(postKeyDefault1);
+        const [defaultDatastoreResults] = await datastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 1);
+        assert.strictEqual(defaultDatastoreResults[0].author, 'D1');
+        // Next, ensure that the other database has the right records
+        const queryS1 = otherDatastore
+          .createQuery('Post')
+          .hasAncestor(postKeySecondary1);
+        const [secondDatastoreResults1] = await otherDatastore.runQuery(
+          queryS1
+        );
+        assert.strictEqual(secondDatastoreResults1.length, 1);
+        assert.strictEqual(secondDatastoreResults1[0].author, 'S1');
+        const queryS2 = otherDatastore
+          .createQuery('Post')
+          .hasAncestor(postKeySecondary2);
+        const [secondDatastoreResults2] = await otherDatastore.runQuery(
+          queryS2
+        );
+        assert.strictEqual(secondDatastoreResults2.length, 1);
+        assert.strictEqual(secondDatastoreResults2[0].author, 'S2');
+        const queryS3 = otherDatastore
+          .createQuery('Post')
+          .hasAncestor(postKeySecondary3);
+        const [secondDatastoreResults3] = await otherDatastore.runQuery(
+          queryS3
+        );
+        assert.strictEqual(secondDatastoreResults3.length, 1);
+        assert.strictEqual(secondDatastoreResults3[0].author, 'S3');
+        // Cleanup
+        await datastore.delete(postKeyDefault1);
+        await otherDatastore.delete(postKeySecondary1);
+        await otherDatastore.delete(postKeySecondary2);
+        await otherDatastore.delete(postKeySecondary3);
+      });
+    });
+
     it('should save/get/delete from a snapshot', async () => {
       function sleep(ms: number) {
         return new Promise(resolve => setTimeout(resolve, ms));

test/index.ts

@@ -19,7 +19,7 @@ import * as proxyquire from 'proxyquire';
 import {PassThrough, Readable} from 'stream';
 
 import * as ds from '../src';
-import {DatastoreOptions} from '../src';
+import {Datastore, DatastoreOptions} from '../src';
 import {entity, Entity, EntityProto, EntityObject} from '../src/entity';
 import {RequestConfig} from '../src/request';
 import * as is from 'is';
@@ -91,6 +91,10 @@ function fakeGoogleAuth(...args: Array<{}>) {
 
 let createInsecureOverride: Function | null;
 
+const SECOND_DATABASE_ID = 'multidb-test';
+
+export {SECOND_DATABASE_ID};
+
 const fakeGoogleGax = {
   GoogleAuth: fakeGoogleAuth,
   GrpcClient: class extends gax.GrpcClient {
@@ -2150,4 +2154,14 @@ describe('Datastore', () => {
       assert.strictEqual(key.name, 'Test');
     });
   });
+
+  describe('multi-db support', () => {
+    it('should get the database id from the client', async () => {
+      const otherDatastore = new Datastore({
+        namespace: `${Date.now()}`,
+        databaseId: SECOND_DATABASE_ID,
+      });
+      assert.strictEqual(otherDatastore.getDatabaseId(), SECOND_DATABASE_ID);
+    });
+  });
 });

test/query.ts

@@ -20,6 +20,8 @@ const {Query} = require('../src/query');
 import {Datastore} from '../src';
 import {AggregateField, AggregateQuery} from '../src/aggregate';
 import {PropertyFilter, EntityFilter, or} from '../src/filter';
+import {entity} from '../src/entity';
+import {SECOND_DATABASE_ID} from './index';
 
 describe('Query', () => {
   const SCOPE = {} as Datastore;
@@ -432,4 +434,42 @@ describe('Query', () => {
       assert.strictEqual(results, runQueryReturnValue);
     });
   });
+
+  it('should pass the database id to the generated layer', async () => {
+    const options = {
+      namespace: `${Date.now()}`,
+      databaseId: SECOND_DATABASE_ID,
+      projectId: 'test-project-id',
+    };
+    const clientName = 'DatastoreClient';
+    const otherDatastore = new Datastore(options);
+    const postKey = new entity.Key({path: ['Post', 'post1']});
+    // Initialize the generated client so that we can mock it out
+    const gapic = Object.freeze({
+      v1: require('../src/v1'),
+    });
+    otherDatastore.clients_.set(clientName, new gapic.v1[clientName](options));
+    const dataClient = otherDatastore.clients_.get(clientName);
+    const projectId = await otherDatastore.getProjectId();
+    if (dataClient) {
+      dataClient['commit'] = (
+        request: any,
+        options: any,
+        callback: (err?: unknown) => void
+      ) => {
+        try {
+          assert.strictEqual(request.databaseId, SECOND_DATABASE_ID);
+          assert.strictEqual(request.projectId, projectId);
+          assert.strictEqual(
+            options.headers['google-cloud-resource-prefix'],
+            `projects/${projectId}`
+          );
+        } catch (e) {
+          callback(e);
+        }
+        callback();
+      };
+    }
+    await otherDatastore.save({key: postKey, data: {title: 'test'}});
+  });
 });

test/request.ts

@@ -36,6 +36,7 @@ import {
   GetResponse,
   RequestCallback,
 } from '../src/request';
+import {addDatabaseIdToRequest} from '../src/util';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type Any = any;
@@ -87,6 +88,11 @@ describe('Request', () => {
     });
     v1FakeClientOverride = null;
     request = new Request();
+    request.datastore = {
+      addDatabaseIdToRequest(d: any, r: any) {
+        addDatabaseIdToRequest(d, r);
+      },
+    };
   });
 
   afterEach(() => sandbox.restore());
@@ -1625,6 +1631,9 @@ describe('Request', () => {
             callback(null, PROJECT_ID);
           },
         },
+        addDatabaseIdToRequest(d: any, r: any) {
+          addDatabaseIdToRequest(d, r);
+        },
       };
     });
 

test/transaction.ts

@@ -23,6 +23,7 @@ import {Datastore, DatastoreRequest, Query, TransactionOptions} from '../src';
 import {Entity} from '../src/entity';
 import * as tsTypes from '../src/transaction';
 import * as sinon from 'sinon';
+import {addDatabaseIdToRequest} from '../src/util';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type Any = any;
@@ -61,6 +62,9 @@ describe('Transaction', () => {
     request_() {},
     projectId: PROJECT_ID,
     namespace: NAMESPACE,
+    addDatabaseIdToRequest(d: any, r: any) {
+      addDatabaseIdToRequest(d, r);
+    },
   } as {} as Datastore;
 
   function key(path: Path) {