diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts index 25fbfaa82..4ddfbfa2e 100644 --- a/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts +++ b/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts @@ -26,6 +26,7 @@ export interface CliConfig { bigQueryProjectId: string; datasetId: string; tableNamePrefix: string; + databaseId?: string; schemas: { [schemaName: string]: FirestoreSchema }; useGemini?: boolean; geminiAnalyzeCollectionPath?: string; @@ -48,6 +49,7 @@ export async function parseConfig(): Promise { bigQueryProjectId: program.bigQueryProject || program.project, datasetId: program.dataset, tableNamePrefix: program.tableNamePrefix, + databaseId: program.databaseId || "(default)", useGemini: !!program.useGemini, schemas: !program.useGemini ? readSchemas(program.schemaFiles) : {}, geminiAnalyzeCollectionPath: program.useGemini, @@ -62,6 +64,7 @@ export async function parseConfig(): Promise { bigQueryProjectId, datasetId, tableNamePrefix, + databaseId, schemaFiles, useGemini, geminiAnalyzeCollectionPath, @@ -75,6 +78,7 @@ export async function parseConfig(): Promise { bigQueryProjectId, datasetId, tableNamePrefix, + databaseId, schemas: !useGemini ? readSchemas([schemaFiles]) : {}, useGemini, geminiAnalyzeCollectionPath, diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts index 24104ebb5..6aa3165b5 100644 --- a/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts +++ b/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts @@ -62,6 +62,13 @@ export const questions = [ validate: (value) => validateInput(value, "table name prefix", BIGQUERY_VALID_CHARACTERS), }, + { + message: + "What is your Firestore database ID? (Leave as default '(default)' for the default database)", + name: "databaseId", + type: "input", + default: "(default)", + }, { message: "Would you like to use a Gemini to automatically analyze your data and generate a draft schema?", diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/config/non-interactive.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/config/non-interactive.ts index c6e659638..67158511e 100644 --- a/firestore-bigquery-export/scripts/gen-schema-view/src/config/non-interactive.ts +++ b/firestore-bigquery-export/scripts/gen-schema-view/src/config/non-interactive.ts @@ -54,6 +54,10 @@ export const configureProgram = () => { "-t, --table-name-prefix ", "A common prefix for the names of all views generated by this script." ) + .option( + "-D, --database-id ", + "Firestore database ID (default: '(default)')" + ) .option( "-f, --schema-files ", "A collection of files from which to read schemas.", diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts index 31259bb9f..296e6bca7 100644 --- a/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts +++ b/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts @@ -15,7 +15,7 @@ */ import type { CliConfig } from "../config"; -import firebase = require("firebase-admin"); +import { getFirestore, GeoPoint, DocumentReference } from "firebase-admin/firestore"; import { genkit, z } from "genkit"; import { googleAI } from "@genkit-ai/googleai"; import * as fs from "fs"; @@ -24,9 +24,12 @@ import inquirer from "inquirer"; export async function sampleFirestoreDocuments( collectionPath: string, - sampleSize: number + sampleSize: number, + databaseId?: string ): Promise { - const db = firebase.firestore(); + const db = (databaseId && databaseId !== "(default)") + ? getFirestore(databaseId) + : getFirestore(); try { const snapshot = await db @@ -54,7 +57,7 @@ export function serializeDocument(data: any): any { return { _type: "timestamp", value: data.toISOString() }; } - if (data instanceof firebase.firestore.GeoPoint) { + if (data instanceof GeoPoint) { return { _type: "geopoint", latitude: data.latitude, @@ -62,7 +65,7 @@ export function serializeDocument(data: any): any { }; } - if (data instanceof firebase.firestore.DocumentReference) { + if (data instanceof DocumentReference) { return { _type: "reference", path: data.path }; } @@ -197,7 +200,8 @@ export const generateSchemaFilesWithGemini = async (config: CliConfig) => { // get sample data from Firestore const sampleData = await sampleFirestoreDocuments( config.geminiAnalyzeCollectionPath!, - config.agentSampleSize! + config.agentSampleSize!, + config.databaseId ); if (sampleData.length === 0) {