Merge pull request #2527 from firebase/@invertase/import-script-alt-firestore

feat(firestore-bigquery-import-script): support non-default dbs

Author: cabljac

firestore-bigquery-export/scripts/import/package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "@firebaseextensions/fs-bq-import-collection",
-  "version": "0.1.26",
+  "version": "0.1.27",
   "description": "Import a Firestore Collection into a BigQuery Changelog Table",
   "main": "./lib/index.js",
   "repository": {

firestore-bigquery-export/scripts/import/package.json

@@ -84,6 +84,13 @@ const questions = [
         PROJECT_ID_MAX_CHARS
       ),
   },
+  {
+    message:
+      "What is your Firestore database instance ID? (Use '(default)' for the default database)",
+    name: "firestoreInstanceId",
+    type: "input",
+    default: "(default)",
+  },
   {
     message: "What is your BigQuery project ID?",
     name: "bigQueryProject",
@@ -271,10 +278,12 @@ export async function parseConfig(): Promise<CliConfig | CliConfigError> {
       cursorPositionFile,
       failedBatchOutput: program.failedBatchOutput,
       transformFunctionUrl: program.transformFunctionUrl,
+      firestoreInstanceId: program.firestoreInstanceId || "(default)",
     };
   }
   const {
     project,
+    firestoreInstanceId,
     sourceCollectionPath,
     bigQueryProject,
     dataset,
@@ -314,6 +323,7 @@ export async function parseConfig(): Promise<CliConfig | CliConfigError> {
     cursorPositionFile,
     failedBatchOutput,
     transformFunctionUrl,
+    firestoreInstanceId: firestoreInstanceId || "(default)",
   };
 }
 

firestore-bigquery-export/scripts/import/src/config.ts

@@ -17,6 +17,7 @@
  */
 import { FirestoreBigQueryEventHistoryTracker } from "@firebaseextensions/firestore-bigquery-change-tracker";
 import * as firebase from "firebase-admin";
+import { getFirestore } from "firebase-admin/firestore";
 import * as fs from "fs";
 import * as util from "util";
 
@@ -50,6 +51,7 @@ const run = async (): Promise<number> => {
     useEmulator,
     cursorPositionFile,
     transformFunctionUrl,
+    firestoreInstanceId,
   } = config;
   if (useEmulator) {
     console.log("Using emulator");
@@ -63,14 +65,23 @@ const run = async (): Promise<number> => {
   // Initialize Firebase
   // This uses applicationDefault to authenticate
   // Please see https://cloud.google.com/docs/authentication/production
+  let app: firebase.app.App;
   if (!firebase.apps.length) {
-    firebase.initializeApp({
+    app = firebase.initializeApp({
       projectId: projectId,
       credential: firebase.credential.applicationDefault(),
       databaseURL: `https://${projectId}.firebaseio.com`,
     });
+  } else {
+    app = firebase.app();
   }
 
+  // Get the Firestore instance for the specified database
+  const db =
+    firestoreInstanceId && firestoreInstanceId !== "(default)"
+      ? getFirestore(app, firestoreInstanceId)
+      : getFirestore(app);
+
   // We pass in the application-level "tableId" here. The tracker determines
   // the name of the raw changelog from this field.
   // TODO: fix this type, it should include clustering apparently
@@ -83,6 +94,7 @@ const run = async (): Promise<number> => {
     useNewSnapshotQuerySyntax,
     bqProjectId: bigQueryProjectId,
     transformFunction: transformFunctionUrl,
+    firestoreInstanceId: firestoreInstanceId,
   });
 
   await initializeDataSink(dataSink, config);
@@ -103,10 +115,10 @@ const run = async (): Promise<number> => {
 
   if (await exists(cursorPositionFile)) {
     let cursorDocumentId = (await read(cursorPositionFile)).toString();
-    cursor = await firebase.firestore().doc(cursorDocumentId).get();
+    cursor = await db.doc(cursorDocumentId).get();
     logs.resumingImport(config, cursorDocumentId);
   }
-  const totalRowsImported = await runSingleThread(dataSink, config, cursor);
+  const totalRowsImported = await runSingleThread(dataSink, config, cursor, db);
   try {
     await unlink(cursorPositionFile);
   } catch (e) {

firestore-bigquery-export/scripts/import/src/index.ts

@@ -65,5 +65,10 @@ export const getCLIOptions = () => {
     .option(
       "-f, --failed-batch-output <file>",
       "Path to the JSON file where failed batches will be recorded."
+    )
+    .option(
+      "--firestore-instance-id <database-id>",
+      "The Firestore database instance ID. Use '(default)' for the default database.",
+      "(default)"
     );
 };

firestore-bigquery-export/scripts/import/src/program.ts

@@ -1,4 +1,5 @@
 import * as firebase from "firebase-admin";
+import { getFirestore } from "firebase-admin/firestore";
 import { cpus } from "os";
 import { pool } from "workerpool";
 import * as logs from "./logs";
@@ -20,19 +21,24 @@ export async function runMultiThread(config: CliConfig): Promise<number> {
         PROJECT_ID: config.projectId,
         GOOGLE_CLOUD_PROJECT: config.projectId,
         GCLOUD_PROJECT: config.projectId,
+        FIRESTORE_INSTANCE_ID: config.firestoreInstanceId,
         FAILED_BATCH_OUTPUT: config.failedBatchOutput || "",
         ...process.env,
       },
     },
   });
 
-  const query = firebase
-    .firestore()
-    .collectionGroup(
-      config.sourceCollectionPath.split("/")[
-        config.sourceCollectionPath.split("/").length - 1
-      ]
-    );
+  const app = firebase.app();
+  const db =
+    config.firestoreInstanceId && config.firestoreInstanceId !== "(default)"
+      ? getFirestore(app, config.firestoreInstanceId)
+      : getFirestore(app);
+
+  const query = db.collectionGroup(
+    config.sourceCollectionPath.split("/")[
+      config.sourceCollectionPath.split("/").length - 1
+    ]
+  );
 
   const partitionsList = query.getPartitions(config.batchSize);