Internal implementation of multiDb (#1949)

* Internal implementation of multiDb

* Fixes from PR review

* Fixes from PR review

Author: tom-andersen

etc/firebase-admin.firestore.api.md

@@ -83,10 +83,13 @@ export { FirestoreDataConverter }
 
 export { GeoPoint }
 
+// @public
+export function getFirestore(): Firestore;
+
 // Warning: (ae-forgotten-export) The symbol "App" needs to be exported by the entry point index.d.ts
 //
 // @public
-export function getFirestore(app?: App): Firestore;
+export function getFirestore(app: App): Firestore;
 
 export { GrpcStatus }
 

src/firestore/firestore-internal.ts

@@ -25,14 +25,22 @@ import { App } from '../app';
 
 export class FirestoreService {
 
-  private appInternal: App;
-  private firestoreClient: Firestore;
+  private readonly appInternal: App;
+  private readonly databases: Map<string, Firestore> = new Map();
 
   constructor(app: App) {
-    this.firestoreClient = initFirestore(app);
     this.appInternal = app;
   }
 
+  getDatabase(databaseId: string): Firestore {
+    let database = this.databases.get(databaseId);
+    if (database === undefined) {
+      database = initFirestore(this.app, databaseId);
+      this.databases.set(databaseId, database);
+    }
+    return database;
+  }
+
   /**
    * Returns the app associated with this Storage instance.
    *
@@ -41,10 +49,6 @@ export class FirestoreService {
   get app(): App {
     return this.appInternal;
   }
-
-  get client(): Firestore {
-    return this.firestoreClient;
-  }
 }
 
 export function getFirestoreOptions(app: App): Settings {
@@ -85,8 +89,9 @@ export function getFirestoreOptions(app: App): Settings {
   });
 }
 
-function initFirestore(app: App): Firestore {
+function initFirestore(app: App, databaseId: string): Firestore {
   const options = getFirestoreOptions(app);
+  options.databaseId = databaseId;
   let firestoreDatabase: typeof Firestore;
   try {
     // Lazy-load the Firestore implementation here, which in turns loads gRPC.

src/firestore/index.ts

@@ -24,6 +24,7 @@ import { Firestore } from '@google-cloud/firestore';
 import { App, getApp } from '../app';
 import { FirebaseApp } from '../app/firebase-app';
 import { FirestoreService } from './firestore-internal';
+import { DEFAULT_DATABASE_ID } from '@google-cloud/firestore/build/src/path';
 
 export {
   AddPrefixToKeys,
@@ -71,7 +72,7 @@ export {
 
 /**
  * Gets the {@link https://googleapis.dev/nodejs/firestore/latest/Firestore.html | Firestore}
- * service for the default app or a given app.
+ * service for the default app.
  *
  * `getFirestore()` can be called with no arguments to access the default
  * app's `Firestore` service or as `getFirestore(app)` to access the
@@ -82,6 +83,20 @@ export {
  * // Get the Firestore service for the default app
  * const defaultFirestore = getFirestore();
  * ```
+
+ * @returns The default {@link https://googleapis.dev/nodejs/firestore/latest/Firestore.html | Firestore}
+ *   service if no app is provided or the `Firestore` service associated with the
+ *   provided app.
+ */
+export function getFirestore(): Firestore;
+
+/**
+ * Gets the {@link https://googleapis.dev/nodejs/firestore/latest/Firestore.html | Firestore}
+ * service for the given app.
+ *
+ * `getFirestore()` can be called with no arguments to access the default
+ * app's `Firestore` service or as `getFirestore(app)` to access the
+ * `Firestore` service associated with a specific app.
  *
  * @example
  * ```javascript
@@ -96,13 +111,30 @@ export {
  *   service if no app is provided or the `Firestore` service associated with the
  *   provided app.
  */
-export function getFirestore(app?: App): Firestore {
-  if (typeof app === 'undefined') {
-    app = getApp();
-  }
+export function getFirestore(app: App): Firestore;
+
+/**
+ * @param databaseId
+ * @internal
+ */
+export function getFirestore(databaseId: string): Firestore;
+
+/**
+ * @param app
+ * @param databaseId
+ * @internal
+ */
+export function getFirestore(app: App, databaseId: string): Firestore;
 
+export function getFirestore(
+  appOrDatabaseId?: App | string,
+  optionalDatabaseId?: string
+): Firestore {
+  const app: App = typeof appOrDatabaseId === 'object' ? appOrDatabaseId : getApp();
+  const databaseId =
+    (typeof appOrDatabaseId === 'string' ? appOrDatabaseId : optionalDatabaseId) || DEFAULT_DATABASE_ID;
   const firebaseApp: FirebaseApp = app as FirebaseApp;
   const firestoreService = firebaseApp.getOrInitService(
     'firestore', (app) => new FirestoreService(app));
-  return firestoreService.client;
+  return firestoreService.getDatabase(databaseId);
 }

test/unit/firestore/firestore.spec.ts

@@ -26,6 +26,7 @@ import {
   ComputeEngineCredential, RefreshTokenCredential
 } from '../../../src/app/credential-internal';
 import { FirestoreService, getFirestoreOptions } from '../../../src/firestore/firestore-internal';
+import { DEFAULT_DATABASE_ID } from '@google-cloud/firestore/build/src/path';
 
 describe('Firestore', () => {
   let mockApp: FirebaseApp;
@@ -98,23 +99,25 @@ describe('Firestore', () => {
       it(`should throw given invalid app: ${ JSON.stringify(invalidApp) }`, () => {
         expect(() => {
           const firestoreAny: any = FirestoreService;
-          return new firestoreAny(invalidApp);
+          const firestoreService: FirestoreService = new firestoreAny(invalidApp);
+          return firestoreService.getDatabase(DEFAULT_DATABASE_ID)
         }).to.throw('First argument passed to admin.firestore() must be a valid Firebase app instance.');
       });
     });
 
     it('should throw given no app', () => {
       expect(() => {
         const firestoreAny: any = FirestoreService;
-        return new firestoreAny();
+        const firestoreService: FirestoreService = new firestoreAny();
+        return firestoreService.getDatabase(DEFAULT_DATABASE_ID)
       }).to.throw('First argument passed to admin.firestore() must be a valid Firebase app instance.');
     });
 
     it('should throw given an invalid credential with project ID', () => {
       // Project ID is read from the environment variable, but the credential is unsupported.
       process.env.GOOGLE_CLOUD_PROJECT = 'project_id';
       expect(() => {
-        return new FirestoreService(mockCredentialApp);
+        return new FirestoreService(mockCredentialApp).getDatabase(DEFAULT_DATABASE_ID);
       }).to.throw(invalidCredError);
     });
 
@@ -123,13 +126,13 @@ describe('Firestore', () => {
       delete process.env.GOOGLE_CLOUD_PROJECT;
       delete process.env.GCLOUD_PROJECT;
       expect(() => {
-        return new FirestoreService(mockCredentialApp);
+        return new FirestoreService(mockCredentialApp).getDatabase(DEFAULT_DATABASE_ID);
       }).to.throw(invalidCredError);
     });
 
     it('should not throw given a valid app', () => {
       expect(() => {
-        return new FirestoreService(mockApp);
+        return new FirestoreService(mockApp).getDatabase(DEFAULT_DATABASE_ID);
       }).not.to.throw();
     });
 
@@ -139,7 +142,7 @@ describe('Firestore', () => {
         delete process.env.GOOGLE_CLOUD_PROJECT;
         delete process.env.GCLOUD_PROJECT;
         expect(() => {
-          return new FirestoreService(config.app);
+          return new FirestoreService(config.app).getDatabase(DEFAULT_DATABASE_ID);
         }).not.to.throw();
       });
     });
@@ -158,19 +161,6 @@ describe('Firestore', () => {
     });
   });
 
-  describe('client', () => {
-    it('returns the client from the constructor', () => {
-      // We expect referential equality here
-      expect(firestore.client).to.not.be.null;
-    });
-
-    it('is read-only', () => {
-      expect(() => {
-        firestore.client = mockApp;
-      }).to.throw('Cannot set property client of #<FirestoreService> which has only a getter');
-    });
-  });
-
   describe('options.projectId', () => {
     it('should return a string when project ID is present in credential', () => {
       const options = getFirestoreOptions(mockApp);

test/unit/firestore/index.spec.ts

@@ -24,6 +24,7 @@ import * as chaiAsPromised from 'chai-as-promised';
 import * as mocks from '../../resources/mocks';
 import { App } from '../../../src/app/index';
 import { getFirestore, Firestore } from '../../../src/firestore/index';
+import { DEFAULT_DATABASE_ID } from '@google-cloud/firestore/build/src/path';
 
 chai.should();
 chai.use(sinonChai);
@@ -66,8 +67,23 @@ describe('Firestore', () => {
 
     it('should return the same instance for a given app instance', () => {
       const db1: Firestore = getFirestore(mockApp);
-      const db2: Firestore = getFirestore(mockApp);
+      const db2: Firestore = getFirestore(mockApp, DEFAULT_DATABASE_ID);
       expect(db1).to.equal(db2);
     });
+
+    it('should return the same instance for a given app instance and databaseId', () => {
+      const db1: Firestore = getFirestore(mockApp, 'db');
+      const db2: Firestore = getFirestore(mockApp, 'db');
+      expect(db1).to.equal(db2);
+    });
+
+    it('should return the different instance for given same app instance, but different databaseId', () => {
+      const db0: Firestore = getFirestore(mockApp, DEFAULT_DATABASE_ID);
+      const db1: Firestore = getFirestore(mockApp, 'db1');
+      const db2: Firestore = getFirestore(mockApp, 'db2');
+      expect(db0).to.not.equal(db1);
+      expect(db0).to.not.equal(db2);
+      expect(db1).to.not.equal(db2);
+    });
   });
 });