fix: Resolving issues with multi-DB support (#1864)

MarkDuckworth · web-flow · commit 1af49c1a852c · 2023-07-20T10:57:19.000-04:00
* Make named DB work.

* override of gax.routingHeader.fromParams
diff --git a/dev/src/index.ts b/dev/src/index.ts
@@ -84,6 +84,7 @@ import {
   RECURSIVE_DELETE_MIN_PENDING_OPS,
   RecursiveDelete,
 } from './recursive-delete';
+import {stringify} from 'querystring';
 
 export {
   CollectionReference,
@@ -590,6 +591,21 @@ export class Firestore implements firestore.Firestore {
           }
         }
 
+        // TODO (multi-db) Revert this override of gax.routingHeader.fromParams
+        // after a permanent fix is applied. See b/292075646
+        // This override of the routingHeader.fromParams does not
+        // encode forward slash characters. This is a temporary fix for b/291780066
+        gax.routingHeader.fromParams = params => {
+          return stringify(params, undefined, undefined, {
+            encodeURIComponent: (val: string) => {
+              return val
+                .split('/')
+                .map(component => encodeURIComponent(component))
+                .join('/');
+            },
+          });
+        };
+
         if (this._settings.ssl === false) {
           const grpcModule = this._settings.grpc ?? require('google-gax').grpc;
           const sslCreds = grpcModule.credentials.createInsecure();
diff --git a/dev/src/path.ts b/dev/src/path.ts
@@ -308,17 +308,14 @@ export class ResourcePath extends Path<ResourcePath> {
    *
    * @private
    * @internal
-   * @param projectIdIfMissing The project ID of the current Firestore project.
-   * The project ID is only used if it's not provided as part of this
-   * ResourcePath.
+   * @param projectId The project ID of the current Firestore project.
    * @return A fully-qualified resource path pointing to the same element.
    */
-  toQualifiedResourcePath(projectIdIfMissing: string): QualifiedResourcePath {
-    return new QualifiedResourcePath(
-      projectIdIfMissing,
-      DEFAULT_DATABASE_ID,
-      ...this.segments
-    );
+  toQualifiedResourcePath(
+    projectId: string,
+    databaseId: string
+  ): QualifiedResourcePath {
+    return new QualifiedResourcePath(projectId, databaseId, ...this.segments);
   }
 }
 
diff --git a/dev/src/reference.ts b/dev/src/reference.ts
@@ -155,7 +155,9 @@ export class DocumentReference<T = firestore.DocumentData>
    */
   get formattedName(): string {
     const projectId = this.firestore.projectId;
-    return this._path.toQualifiedResourcePath(projectId).formattedName;
+    const databaseId = this.firestore.databaseId;
+    return this._path.toQualifiedResourcePath(projectId, databaseId)
+      .formattedName;
   }
 
   /**
@@ -2328,8 +2330,11 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
     transactionIdOrReadTime?: Uint8Array | Timestamp
   ): api.IRunQueryRequest {
     const projectId = this.firestore.projectId;
-    const parentPath =
-      this._queryOptions.parentPath.toQualifiedResourcePath(projectId);
+    const databaseId = this.firestore.databaseId;
+    const parentPath = this._queryOptions.parentPath.toQualifiedResourcePath(
+      projectId,
+      databaseId
+    );
 
     const structuredQuery = this.toStructuredQuery();
 
@@ -2387,8 +2392,11 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
    */
   _toBundledQuery(): protos.firestore.IBundledQuery {
     const projectId = this.firestore.projectId;
-    const parentPath =
-      this._queryOptions.parentPath.toQualifiedResourcePath(projectId);
+    const databaseId = this.firestore.databaseId;
+    const parentPath = this._queryOptions.parentPath.toQualifiedResourcePath(
+      projectId,
+      databaseId
+    );
     const structuredQuery = this.toStructuredQuery();
 
     const bundledQuery: protos.firestore.IBundledQuery = {
@@ -2858,7 +2866,8 @@ export class CollectionReference<T = firestore.DocumentData>
     const tag = requestTag();
     return this.firestore.initializeIfNeeded(tag).then(() => {
       const parentPath = this._queryOptions.parentPath.toQualifiedResourcePath(
-        this.firestore.projectId
+        this.firestore.projectId,
+        this.firestore.databaseId
       );
 
       const request: api.IListDocumentsRequest = {
diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts
@@ -82,8 +82,14 @@ if (process.env.NODE_ENV === 'DEBUG') {
 }
 
 function getTestRoot(settings: Settings = {}) {
+  const internalSettings: Settings = {};
+  if (process.env.FIRESTORE_NAMED_DATABASE) {
+    internalSettings.databaseId = process.env.FIRESTORE_NAMED_DATABASE;
+  }
+
   const firestore = new Firestore({
-    ...settings,
+    ...internalSettings,
+    ...settings, // caller settings take precedent over internal settings
   });
   return firestore.collection(`node_${version}_${autoId()}`);
 }
