fix(bigquery-firestore-export): write run output to deterministic ids

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<Promise<DocumentReference<DocumentData>>> = [];
+    const writes: Array<Promise<WriteResult>> = [];
     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);

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<string, Map<string, unknown>>();
+  const rejected = new Set(rejectDocIds);
+
+  const docsFor = (path: string) => {
+    const existing = collections.get(path);
+    if (existing) return existing;
+    const created = new Map<string, unknown>();
+    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<typeof docRef>
+        ) => Promise<{ data: () => unknown }>;
+        set: (ref: ReturnType<typeof docRef>, data: unknown) => void;
+      }) => Promise<void>
+    ) =>
+      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,
+    });
+  });
+});