diff --git a/kits/bigquery-firestore-export/src/helper.ts b/kits/bigquery-firestore-export/src/helper.ts index fe48a03b9..bbbadaafb 100644 --- a/kits/bigquery-firestore-export/src/helper.ts +++ b/kits/bigquery-firestore-export/src/helper.ts @@ -23,10 +23,9 @@ import { Geography, } from "@google-cloud/bigquery"; import { - type DocumentData, - type DocumentReference, type Firestore, Timestamp, + type WriteResult, } from "firebase-admin/firestore"; import type { ResolvedBigqueryFirestoreExportConfig } from "./export-config"; import * as logs from "./logs"; @@ -62,6 +61,17 @@ const TRANSFER_RUN_NAME_REGEX = const TRANSFER_CONFIG_NAME_REGEX = /^projects\/([^/]+)\/locations\/([^/]+)\/transferConfigs\/([^/]+)$/; const FIRESTORE_WRITE_CHUNK_SIZE = 10_000; +const OUTPUT_DOC_ID_LENGTH = 12; + +/** + * Document id for the row at `index` of a run's results. `processMessages` + * deploys with `retry: true`, so a run that dies part way through is redelivered + * and rewritten; keying output documents by row index makes that rewrite an + * overwrite instead of a second copy of the result set. + */ +function outputDocumentId(index: number): string { + return String(index).padStart(OUTPUT_DOC_ID_LENGTH, "0"); +} export function parseTransferRunName(name: string): ParsedTransferRunName { const match = name.match(TRANSFER_RUN_NAME_REGEX); @@ -233,13 +243,17 @@ export async function writeRunResultsToFirestore( let succeededRowCount = 0; for (let i = 0; i < rows.length; i += FIRESTORE_WRITE_CHUNK_SIZE) { - const writes: Array>> = []; + const writes: Array> = []; for ( let j = i; j < i + FIRESTORE_WRITE_CHUNK_SIZE && j < rows.length; j++ ) { - writes.push(collection.add(convertUnsupportedDataTypes(rows[j]))); + writes.push( + collection + .doc(outputDocumentId(j)) + .set(convertUnsupportedDataTypes(rows[j])) + ); } const results = await Promise.allSettled(writes); diff --git a/kits/bigquery-firestore-export/tests/helper.test.ts b/kits/bigquery-firestore-export/tests/helper.test.ts index 74c1ba1e1..5b6c2e454 100644 --- a/kits/bigquery-firestore-export/tests/helper.test.ts +++ b/kits/bigquery-firestore-export/tests/helper.test.ts @@ -14,14 +14,29 @@ * limitations under the License. */ +import type { BigQuery } from "@google-cloud/bigquery"; import { Geography } from "@google-cloud/bigquery"; +import type { Firestore } from "firebase-admin/firestore"; import { Timestamp } from "firebase-admin/firestore"; -import { describe, expect, test } from "vitest"; +import { describe, expect, test, vi } from "vitest"; +import { resolveConfig } from "../src/export-config"; import { convertUnsupportedDataTypes, parseTransferConfigName, parseTransferRunName, + writeRunResultsToFirestore, } from "../src/helper"; +import type { BigQueryRow, TransferRunMessage } from "../src/types"; + +vi.mock("../src/logs", () => ({ + bigqueryJobStarted: vi.fn(), + bigqueryQueryFailed: vi.fn(), + bigqueryResultsRowCount: vi.fn(), + errorWritingToFirestore: vi.fn(), + handlingNonSuccessRun: vi.fn(), + latestDocUpdateSkipped: vi.fn(), + runResultsWrittenToFirestore: vi.fn(), +})); describe("transfer resource parsing", () => { test("parses config and run resource names", () => { @@ -72,3 +87,145 @@ describe("convertUnsupportedDataTypes", () => { expect(converted.nested).toEqual([{ value: true }]); }); }); + +const CONFIG = resolveConfig({ + bigqueryDatasetLocation: "US", + projectId: "test-project", + instanceId: "users-export", + datasetId: "analytics", + tableName: "out", + queryString: "SELECT * FROM source.users", + displayName: "Users export", + schedule: "every 24 hours", + firestoreCollection: "transferConfigs", +}); +const RUN_ID = "run-1"; +const RUNS_PATH = "transferConfigs/config-1/runs"; +const OUTPUT_PATH = `${RUNS_PATH}/${RUN_ID}/output`; +const MESSAGE = { + json: { + name: `projects/test-project/locations/us/transferConfigs/config-1/runs/${RUN_ID}`, + runTime: "2026-08-20T10:05:39Z", + state: "SUCCEEDED", + destinationDatasetId: "analytics", + params: { destination_table_name_template: 'out_{run_time|"%H%M%S"}' }, + }, +} as unknown as TransferRunMessage; + +/** Minimal Firestore double keyed by collection path and document id. */ +function makeDb(rejectDocIds: string[] = []) { + const collections = new Map>(); + const rejected = new Set(rejectDocIds); + + const docsFor = (path: string) => { + const existing = collections.get(path); + if (existing) return existing; + const created = new Map(); + collections.set(path, created); + return created; + }; + const docRef = (path: string, id: string) => ({ + set(data: unknown) { + if (rejected.has(id)) { + return Promise.reject(new Error(`write refused for ${id}`)); + } + docsFor(path).set(id, data); + return Promise.resolve(); + }, + read: () => docsFor(path).get(id), + }); + const db = { + collection: (path: string) => ({ doc: (id: string) => docRef(path, id) }), + runTransaction: ( + fn: (tx: { + get: ( + ref: ReturnType + ) => Promise<{ data: () => unknown }>; + set: (ref: ReturnType, data: unknown) => void; + }) => Promise + ) => + fn({ + get: (ref) => Promise.resolve({ data: () => ref.read() }), + set: (ref, data) => void ref.set(data), + }), + } as unknown as Firestore; + + return { + db, + docs: (path: string) => docsFor(path), + output: () => docsFor(OUTPUT_PATH), + }; +} + +function makeBigquery(rows: BigQueryRow[]) { + return { + createQueryJob: () => + Promise.resolve([ + { id: "job-1", getQueryResults: () => Promise.resolve([rows]) }, + ]), + } as unknown as BigQuery; +} + +function rows(count: number): BigQueryRow[] { + return Array.from({ length: count }, (_unused, index) => ({ + id: index, + label: `row-${index}`, + })); +} + +describe("writeRunResultsToFirestore", () => { + test("keys each output document by its zero-padded row index", async () => { + const { db, output } = makeDb(); + + await writeRunResultsToFirestore( + { db, bigquery: makeBigquery(rows(3)), config: CONFIG }, + MESSAGE + ); + + expect([...output().keys()]).toEqual([ + "000000000000", + "000000000001", + "000000000002", + ]); + expect(output().get("000000000001")).toEqual({ id: 1, label: "row-1" }); + }); + + test("a redelivered run overwrites its output instead of appending", async () => { + const { db, output } = makeDb(); + const ctx = { db, bigquery: makeBigquery(rows(50)), config: CONFIG }; + + await writeRunResultsToFirestore(ctx, MESSAGE); + const firstPass = [...output().keys()]; + await writeRunResultsToFirestore(ctx, MESSAGE); + + expect(output().size).toBe(50); + expect([...output().keys()]).toEqual(firstPass); + }); + + test("continues past a chunk boundary without restarting ids", async () => { + const { db, output } = makeDb(); + + await writeRunResultsToFirestore( + { db, bigquery: makeBigquery(rows(10_001)), config: CONFIG }, + MESSAGE + ); + + expect(output().size).toBe(10_001); + expect(output().has("000000010000")).toBe(true); + }); + + test("counts a rejected row without dropping the others", async () => { + const { db, docs, output } = makeDb(["000000000001"]); + + await writeRunResultsToFirestore( + { db, bigquery: makeBigquery(rows(3)), config: CONFIG }, + MESSAGE + ); + + expect([...output().keys()]).toEqual(["000000000000", "000000000002"]); + expect(docs(RUNS_PATH).get(RUN_ID)).toMatchObject({ + failedRowCount: 1, + totalRowCount: 3, + }); + }); +});