feat: preferRest option for Firestore

Author: alexander-fenster

etc/firebase-admin.firestore.api.md

@@ -82,6 +82,11 @@ export { Firestore }
 
 export { FirestoreDataConverter }
 
+// @public
+export interface FirestoreSettings {
+    preferRest?: boolean;
+}
+
 export { GeoPoint }
 
 // @public
@@ -94,6 +99,9 @@ export function getFirestore(app: App): Firestore;
 
 export { GrpcStatus }
 
+// @public
+export function initializeFirestore(app: App, settings?: FirestoreSettings): Firestore;
+
 export { NestedUpdateFields }
 
 export { OrderByDirection }

src/firestore/firestore-internal.ts

@@ -23,24 +23,46 @@ import * as validator from '../utils/validator';
 import * as utils from '../utils/index';
 import { App } from '../app';
 
+/**
+ * Settings to pass to the Firestore constructor.
+ * 
+ * @public
+ */
+export interface FirestoreSettings {
+  /**
+   * Use HTTP/1.1 REST transport where possible.
+   * 
+   * @defaultValue `undefined`
+   */
+  preferRest?: boolean;
+}
+
 export class FirestoreService {
 
   private readonly appInternal: App;
   private readonly databases: Map<string, Firestore> = new Map();
+  private readonly firestoreSettings: FirestoreSettings;
 
-  constructor(app: App) {
+  constructor(app: App, firestoreSettings?: FirestoreSettings) {
     this.appInternal = app;
+    this.firestoreSettings = firestoreSettings ?? {};
   }
 
   getDatabase(databaseId: string): Firestore {
     let database = this.databases.get(databaseId);
     if (database === undefined) {
-      database = initFirestore(this.app, databaseId);
+      database = initFirestore(this.app, databaseId, this.firestoreSettings);
       this.databases.set(databaseId, database);
     }
     return database;
   }
 
+  checkIfSameSettings(firestoreSettings: FirestoreSettings): boolean {
+    // If we start passing more settings to Firestore constructor,
+    // replace this with deep equality check.
+    return (firestoreSettings.preferRest === this.firestoreSettings.preferRest);
+  }
+
   /**
    * Returns the app associated with this Storage instance.
    *
@@ -51,7 +73,7 @@ export class FirestoreService {
   }
 }
 
-export function getFirestoreOptions(app: App): Settings {
+export function getFirestoreOptions(app: App, firestoreSettings?: FirestoreSettings): Settings {
   if (!validator.isNonNullObject(app) || !('options' in app)) {
     throw new FirebaseFirestoreError({
       code: 'invalid-argument',
@@ -63,6 +85,7 @@ export function getFirestoreOptions(app: App): Settings {
   const credential = app.options.credential;
   // eslint-disable-next-line @typescript-eslint/no-var-requires
   const { version: firebaseVersion } = require('../../package.json');
+  const preferRest = firestoreSettings?.preferRest;
   if (credential instanceof ServiceAccountCredential) {
     return {
       credentials: {
@@ -73,12 +96,15 @@ export function getFirestoreOptions(app: App): Settings {
       // guaranteed to be available.
       projectId: projectId!,
       firebaseVersion,
+      preferRest,
     };
   } else if (isApplicationDefault(app.options.credential)) {
     // Try to use the Google application default credentials.
     // If an explicit project ID is not available, let Firestore client discover one from the
     // environment. This prevents the users from having to set GOOGLE_CLOUD_PROJECT in GCP runtimes.
-    return validator.isNonEmptyString(projectId) ? { projectId, firebaseVersion } : { firebaseVersion };
+    return validator.isNonEmptyString(projectId)
+      ? { projectId, firebaseVersion, preferRest }
+      : { firebaseVersion, preferRest };
   }
 
   throw new FirebaseFirestoreError({
@@ -89,8 +115,8 @@ export function getFirestoreOptions(app: App): Settings {
   });
 }
 
-function initFirestore(app: App, databaseId: string): Firestore {
-  const options = getFirestoreOptions(app);
+function initFirestore(app: App, databaseId: string, firestoreSettings?: FirestoreSettings): Firestore {
+  const options = getFirestoreOptions(app, firestoreSettings);
   options.databaseId = databaseId;
   let firestoreDatabase: typeof Firestore;
   try {

src/firestore/index.ts

@@ -23,8 +23,9 @@
 import { Firestore } from '@google-cloud/firestore';
 import { App, getApp } from '../app';
 import { FirebaseApp } from '../app/firebase-app';
-import { FirestoreService } from './firestore-internal';
+import { FirestoreService, FirestoreSettings } from './firestore-internal';
 import { DEFAULT_DATABASE_ID } from '@google-cloud/firestore/build/src/path';
+import { FirebaseFirestoreError } from '../utils/error';
 
 export {
   AddPrefixToKeys,
@@ -71,6 +72,8 @@ export {
   setLogFunction,
 } from '@google-cloud/firestore';
 
+export { FirestoreSettings };
+
 /**
  * Gets the {@link https://googleapis.dev/nodejs/firestore/latest/Firestore.html | Firestore}
  * service for the default app.
@@ -139,3 +142,58 @@ export function getFirestore(
     'firestore', (app) => new FirestoreService(app));
   return firestoreService.getDatabase(databaseId);
 }
+
+/**
+ * Gets the {@link https://googleapis.dev/nodejs/firestore/latest/Firestore.html | Firestore}
+ * service for the given app, passing extra parameters to its constructor.
+ *
+ * @example
+ * ```javascript
+ * // Get the Firestore service for a specific app, require HTTP/1.1 REST transport
+ * const otherFirestore = initializeFirestore(app, {preferRest: true});
+ * ```
+ *
+ * @param App - whose `Firestore` service to
+ *   return. If not provided, the default `Firestore` service will be returned.
+ * 
+ * @param settings - Settings object to be passed to the constructor.
+ *
+ * @returns The `Firestore` service associated with the provided app.
+ */
+export function initializeFirestore(app: App, settings?: FirestoreSettings): Firestore;
+
+/**
+ * @param app 
+ * @param settings 
+ * @param databaseId 
+ * @internal
+ */
+export function initializeFirestore(
+  app: App,
+  settings: FirestoreSettings,
+  databaseId: string
+): Firestore;
+
+export function initializeFirestore(
+  app: App,
+  settings?: FirestoreSettings,
+  databaseId?: string
+): Firestore {
+  settings ??= {};
+  databaseId ??= DEFAULT_DATABASE_ID;
+  const firebaseApp: FirebaseApp = app as FirebaseApp;
+  const firestoreService = firebaseApp.getOrInitService(
+    'firestore', (app) => new FirestoreService(app, settings));
+
+  if (!firestoreService.checkIfSameSettings(settings)) {
+    throw new FirebaseFirestoreError({
+      code: 'failed-precondition',
+      message: 'initializeFirestore() has already been called with ' +
+        'different options. To avoid this error, call initializeFirestore() with the ' +
+        'same options as when it was originally called, or call getFirestore() to return the' +
+        ' already initialized instance.'
+    });
+  }
+
+  return firestoreService.getDatabase(databaseId);
+}

test/integration/firestore.spec.ts

@@ -20,7 +20,7 @@ import { clone } from 'lodash';
 import * as admin from '../../lib/index';
 import {
   DocumentReference, DocumentSnapshot, FieldValue, Firestore, FirestoreDataConverter,
-  QueryDocumentSnapshot, Timestamp, getFirestore, setLogFunction,
+  QueryDocumentSnapshot, Timestamp, getFirestore, initializeFirestore, setLogFunction,
 } from '../../lib/firestore/index';
 
 chai.should();
@@ -47,6 +47,11 @@ describe('admin.firestore', () => {
     expect(firestore).to.not.be.undefined;
   });
 
+  it('initializeFirestore returns a Firestore client', () => {
+    const firestore: Firestore = initializeFirestore(admin.app());
+    expect(firestore).to.not.be.undefined;
+  });
+
   it('admin.firestore() returns a Firestore client', () => {
     const firestore: admin.firestore.Firestore = admin.firestore();
     expect(firestore).to.not.be.undefined;

test/unit/firestore/firestore.spec.ts

@@ -200,4 +200,16 @@ describe('Firestore', () => {
       });
     });
   });
+
+  describe('options.preferRest', () => {
+    it('should not enable preferRest by default', () => {
+      const options = getFirestoreOptions(mockApp);
+      expect(options.preferRest).to.be.undefined;
+    });
+
+    it('should enable preferRest if provided', () => {
+      const options = getFirestoreOptions(mockApp, { preferRest: true });
+      expect(options.preferRest).to.be.true;
+    });
+  });
 });

test/unit/firestore/index.spec.ts

@@ -23,7 +23,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 { getFirestore, initializeFirestore, Firestore } from '../../../src/firestore/index';
 import { DEFAULT_DATABASE_ID } from '@google-cloud/firestore/build/src/path';
 
 chai.should();
@@ -86,4 +86,59 @@ describe('Firestore', () => {
       expect(db1).to.not.equal(db2);
     });
   });
+
+  describe('initializeFirestore()', () => {
+    it('should throw when default app is not available', () => {
+      expect(() => {
+        initializeFirestore(mockApp);
+      }).to.throw('The default Firebase app does not exist.');
+    });
+
+    it('should reject given an invalid credential without project ID', () => {
+      // Project ID not set in the environment.
+      delete process.env.GOOGLE_CLOUD_PROJECT;
+      delete process.env.GCLOUD_PROJECT;
+      expect(() => initializeFirestore(mockCredentialApp)).to.throw(noProjectIdError);
+    });
+
+    it('should not throw given a valid app', () => {
+      expect(() => {
+        return initializeFirestore(mockApp);
+      }).not.to.throw();
+    });
+
+    it('should return the same instance for a given app instance', () => {
+      const db1: Firestore = initializeFirestore(mockApp);
+      const db2: Firestore = initializeFirestore(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 = initializeFirestore(mockApp, {}, 'db');
+      const db2: Firestore = initializeFirestore(mockApp, {}, 'db');
+      expect(db1).to.equal(db2);
+    });
+
+    it('should return the different instance for given same app instance, but different databaseId', () => {
+      const db0: Firestore = initializeFirestore(mockApp, {}, DEFAULT_DATABASE_ID);
+      const db1: Firestore = initializeFirestore(mockApp, {}, 'db1');
+      const db2: Firestore = initializeFirestore(mockApp, {}, 'db2');
+      expect(db0).to.not.equal(db1);
+      expect(db0).to.not.equal(db2);
+      expect(db1).to.not.equal(db2);
+    });
+
+    it('getFirestore should return the same instance as initializeFirestore returned earlier', () => {
+      const db1: Firestore = initializeFirestore(mockApp, {}, 'db');
+      const db2: Firestore = getFirestore(mockApp, 'db');
+      expect(db1).to.equal(db2);
+    });
+
+    it('initializeFirestore should not allow create an instance with different settings', () => {
+      initializeFirestore(mockApp, {}, 'db');
+      expect(() => {
+        return initializeFirestore(mockApp, { preferRest: true }, 'db');
+      }).to.throw(/has already been called with different options/);
+    });
+  });
 });