feat: Support for using multiple databases in datastore

package.json

@@ -59,6 +59,7 @@
     "@types/node": "^20.4.9",
     "@types/proxyquire": "^1.3.28",
     "@types/sinon": "^10.0.0",
+    "async": "^3.2.4",
     "c8": "^8.0.1",
     "gapic-tools": "^0.1.8",
     "gts": "^5.0.0",

src/index.ts

@@ -700,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();
   }
@@ -1818,6 +1828,7 @@ promisifyAll(Datastore, {
     'double',
     'isDouble',
     'geoPoint',
+    'getDatabaseId',
     'getProjectId',
     'isGeoPoint',
     'index',
@@ -1898,6 +1909,7 @@ export interface DatastoreOptions extends GoogleAuthOptions {
   namespace?: string;
   apiEndpoint?: string;
   sslCreds?: ChannelCredentials;
+  databaseId?: string;
 }
 
 export interface KeyToLegacyUrlSafeCallback {

src/request.ts

@@ -44,7 +44,7 @@
   KeyProto,
   ResponseResult,
   Entities,
   ValueProto,
 } from './entity';
 import {
   Query,
@@ -596,9 +596,9 @@
           const results = res.batch.aggregationResults;
           const finalResults = results
             .map(
               (aggregationResult: any) => aggregationResult.aggregateProperties
             )
             .map((aggregateProperties: any) =>
               Object.fromEntries(
                 new Map(
                   Object.keys(aggregateProperties).map(key => [
@@ -933,7 +933,7 @@
     callback?: SaveCallback
   ): void | Promise<CommitResponse> {
     const transaction = this.datastore.transaction();
     transaction.run(async (err: any) => {
       if (err) {
         try {
           await transaction.rollback();
@@ -992,6 +992,10 @@
       }
     }
 
+    if (datastore.options && datastore.options.databaseId) {
+      reqOpts.databaseId = datastore.options.databaseId;
+    }
+
     if (method === 'rollback') {
       reqOpts.transaction = this.id;
     }
@@ -1143,6 +1147,7 @@
   reqOpts?: RequestOptions;
 }
 export interface SharedQueryOptions {
+  databaseId?: string;
   projectId?: string;
   partitionId?: google.datastore.v1.IPartitionId | null;
   readOptions?: {

system-test/datastore.ts

@@ -17,20 +17,33 @@
 import * as path from 'path';
 import {after, before, describe, it} from 'mocha';
 import * as yaml from 'js-yaml';
-import {Datastore, Index} from '../src';
+import {Datastore, DatastoreOptions, Index} from '../src';
 import {google} from '../protos/protos';
 import {Storage} from '@google-cloud/storage';
 import {AggregateField} from '../src/aggregate';
 import {PropertyFilter, and, or} from '../src/filter';
 import {entity} from '../src/entity';
 import KEY_SYMBOL = entity.KEY_SYMBOL;
 
+const async = require('async');
+
+const SECOND_DATABASE_ID = 'multidb-test';
+
+async.each(
+    [
+      {
+        namespace: `${Date.now()}`,
+      },
+      {
+        databaseId: SECOND_DATABASE_ID,
+        namespace: `second-db-${Date.now()}`,
+      },
+    ],
+    (clientOptions: DatastoreOptions) => {
 describe('Datastore', () => {
   let timeBeforeDataCreation: number;
   const testKinds: string[] = [];
-  const datastore = new Datastore({
-    namespace: `${Date.now()}`,
-  });
+  const datastore = new Datastore(clientOptions);
   // Override the Key method so we can track what keys are created during the
   // tests. They are then deleted in the `after` hook.
   const key = datastore.key;
@@ -333,6 +346,110 @@
       await datastore.delete(postKey);
     });
 
+    describe('multi-db support for read and write operations', () => {
+      const namespace = `${Date.now()}`;
+      const keyHierarchy = ['Post', 'post1'];
+      const defaultDatastore = new Datastore({namespace});
+      it('should run a query with another database', async () => {
+        // First verify that a record gets written to datastore
+        const postKey = defaultDatastore.key(keyHierarchy);
+        await defaultDatastore.save({key: postKey, data: post});
+        const query = defaultDatastore.createQuery('Post').hasAncestor(postKey);
+        const [defaultDatastoreResults] = await defaultDatastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 1);
+        const [entity] = await defaultDatastore.get(postKey);
+        assert.strictEqual(entity.author, 'Silvano');
+        // With another database, verify that a query returns no results
+        const otherDatastore = new Datastore({
+          namespace,
+          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 defaultDatastore.delete(postKey);
+      });
+      it('should ensure save works with another database', async () => {
+        // First verify that the default database is empty
+        const postKey = defaultDatastore.key(keyHierarchy);
+        const query = defaultDatastore.createQuery('Post').hasAncestor(postKey);
+        const [defaultDatastoreResults] = await defaultDatastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 0);
+        const [originalSecondaryResults] = await defaultDatastore.runQuery(query);
+        assert.strictEqual(originalSecondaryResults.length, 0);
+        const [entity] = await defaultDatastore.get(postKey);
+        assert(typeof entity === 'undefined');
+        // With another database, verify that saving to the database works
+        const otherDatastore = new Datastore({
+          namespace,
+          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 defaultDatastore.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 () => {
+        interface DatastoreData {
+          key: entity.Key;
+          data: any;
+        }
+        const secondaryDatastore = new Datastore({
+          namespace: `${Date.now()}`,
+          databaseId: SECOND_DATABASE_ID,
+        });
+        // Save all data to the default database
+        const defaultAuthor = 'default database author';
+        const defaultData = Object.assign({}, post);
+        defaultData.author = defaultAuthor;
+        const defaultPostKey = defaultDatastore.key(['Post', 'default post key']);
+        await defaultDatastore.save({key: defaultPostKey, data: defaultData});
+        // Save all data to the secondary database
+        const secondaryIndices = [1, 2, 3];
+        const secondaryData: DatastoreData[] = secondaryIndices.map(number => {
+          const authorName = 'secondary author ' + number.toString();
+          const keyName = 'secondary key ' + number.toString();
+          const postData = Object.assign({}, post);
+          postData.author = authorName;
+          return {
+            key: secondaryDatastore.key(['Post', keyName]),
+            data: postData,
+          };
+        });
+        await Promise.all(
+          secondaryData.map(async datum => secondaryDatastore.save(datum))
+        );
+        // Next, ensure that the default database has the right records
+        const query = defaultDatastore.createQuery('Post').hasAncestor(defaultPostKey);
+        const [defaultDatastoreResults] = await defaultDatastore.runQuery(query);
+        assert.strictEqual(defaultDatastoreResults.length, 1);
+        assert.strictEqual(defaultDatastoreResults[0].author, defaultAuthor);
+        // Next, ensure that the other database has the right records
+        await Promise.all(
+          secondaryData.map(async datum => {
+            const query = secondaryDatastore
+              .createQuery('Post')
+              .hasAncestor(datum.key);
+            const [results] = await secondaryDatastore.runQuery(query);
+            assert.strictEqual(results.length, 1);
+            assert.strictEqual(results[0].author, datum.data.author);
+          })
+        );
+        // Cleanup
+        await defaultDatastore.delete(defaultPostKey);
+        await Promise.all(
+          secondaryData.map(datum => secondaryDatastore.delete(datum.key))
+        );
+      });
+    });
+
     it('should save/get/delete from a snapshot', async () => {
       function sleep(ms: number) {
         return new Promise(resolve => setTimeout(resolve, ms));
@@ -1862,3 +1979,4 @@
     });
   });
 });
+});

test/index.ts

@@ -19,12 +19,13 @@ 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';
 import * as sinon from 'sinon';
 import * as extend from 'extend';
+const async = require('async');
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const v1 = require('../src/v1/index.js');
@@ -91,6 +92,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 {
@@ -139,6 +144,9 @@ function FakeV1() {}
 
 const sandbox = sinon.createSandbox();
 
+async.each(
+    [{namespace: `${Date.now()}`}, {namespace: `second-db-${Date.now()}`, databaseId: SECOND_DATABASE_ID}],
+    (clientOptions: DatastoreOptions) => {
 describe('Datastore', () => {
   let Datastore: typeof ds.Datastore;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -149,7 +157,7 @@ describe('Datastore', () => {
 
   const DATASTORE_PROJECT_ID_CACHED = process.env.DATASTORE_PROJECT_ID;
 
-  const OPTIONS = {
+  const DEFAULT_OPTIONS = {
     projectId: PROJECT_ID,
     apiEndpoint: 'http://localhost',
     credentials: {},
@@ -158,6 +166,8 @@ describe('Datastore', () => {
     namespace: NAMESPACE,
   };
 
+  const OPTIONS = Object.assign(DEFAULT_OPTIONS, clientOptions);
+
   before(() => {
     Object.assign(fakeEntity, fakeEntityInit);
     Datastore = proxyquire('../src', {
@@ -2150,4 +2160,15 @@ 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;
@@ -558,4 +560,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/transaction.ts

@@ -18,11 +18,13 @@ import * as assert from 'assert';
 import {afterEach, beforeEach, before, describe, it} from 'mocha';
 import * as proxyquire from 'proxyquire';
 
-// import {google} from '../proto/datastore';
-import {Datastore, DatastoreRequest, Query, TransactionOptions} from '../src';
+import {Datastore, DatastoreOptions, DatastoreRequest, Query, TransactionOptions} from '../src';
 import {Entity} from '../src/entity';
 import * as tsTypes from '../src/transaction';
 import * as sinon from 'sinon';
+import {RequestConfig} from '../src/request';
+import {SECOND_DATABASE_ID} from './index';
+const async = require('async');
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type Any = any;
@@ -50,19 +52,24 @@ const fakePfy = Object.assign({}, pfy, {
   },
 });
 
+async.each(
+    [{}, {databaseId: SECOND_DATABASE_ID}],
+    (clientOptions: DatastoreOptions) => {
 describe('Transaction', () => {
   let Transaction: typeof tsTypes.Transaction;
   let transaction: tsTypes.Transaction;
   const TRANSACTION_ID = 'transaction-id';
   const PROJECT_ID = 'project-id';
   const NAMESPACE = 'a-namespace';
 
-  const DATASTORE = {
+  const DEFAULT_DATASTORE = {
     request_() {},
     projectId: PROJECT_ID,
     namespace: NAMESPACE,
   } as {} as Datastore;
 
+  const DATASTORE = Object.assign(DEFAULT_DATASTORE, clientOptions);
+
   function key(path: Path) {
     return new entity.Key({path: arrify(path)});
   }
@@ -582,7 +589,7 @@ describe('Transaction', () => {
           },
         } as {} as TransactionOptions;
 
-        transaction.request_ = config => {
+        transaction.request_ = (config: RequestConfig) => {
           assert.deepStrictEqual(config.reqOpts, options);
           done();
         };
@@ -751,3 +758,4 @@ describe('Transaction', () => {
     });
   });
 });
+});