From 27d8258abe0addeb9151acbea31196c4873105bc Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Mon, 7 Sep 2026 15:19:37 +0100 Subject: [PATCH] feat(firestore-bigquery-export): add the sync task enqueue module Adds enqueueSyncTask and syncQueuePath ahead of wiring them into the write path. The queue path is left unprefixed so firebase-admin resolves the kit instance prefix, and the region comes from FUNCTION_REGION with DATABASE_REGION as the local fallback. Enqueue retries in-process with bounded backoff, treats task-already-exists as success, and clamps the attempt budget to at least one. SerializedDocumentChange is exported as the task payload type. --- .../firestore-bigquery-export/src/handlers.ts | 2 +- kits/firestore-bigquery-export/src/lib.ts | 1 + kits/firestore-bigquery-export/src/tasks.ts | 126 ++++++++++ .../tests/tasks.emulator.test.ts | 131 ++++++++++ .../tests/tasks.test.ts | 229 ++++++++++++++++++ 5 files changed, 488 insertions(+), 1 deletion(-) create mode 100644 kits/firestore-bigquery-export/src/tasks.ts create mode 100644 kits/firestore-bigquery-export/tests/tasks.emulator.test.ts create mode 100644 kits/firestore-bigquery-export/tests/tasks.test.ts diff --git a/kits/firestore-bigquery-export/src/handlers.ts b/kits/firestore-bigquery-export/src/handlers.ts index 7510f3c8ee..daf7cd254e 100644 --- a/kits/firestore-bigquery-export/src/handlers.ts +++ b/kits/firestore-bigquery-export/src/handlers.ts @@ -30,7 +30,7 @@ import * as logs from "./logs"; import { getChangeType, getDocumentId } from "./util"; /** Serialized Firestore change ready to write to BigQuery. */ -interface SerializedDocumentChange { +export interface SerializedDocumentChange { timestamp: string; eventId: string; fullResourceName: string; diff --git a/kits/firestore-bigquery-export/src/lib.ts b/kits/firestore-bigquery-export/src/lib.ts index c17127998b..3c83dd304c 100644 --- a/kits/firestore-bigquery-export/src/lib.ts +++ b/kits/firestore-bigquery-export/src/lib.ts @@ -51,5 +51,6 @@ export { export { type DocumentWriteEvent, type HandlerContext, + type SerializedDocumentChange, handleDocumentWrite, } from "./handlers"; diff --git a/kits/firestore-bigquery-export/src/tasks.ts b/kits/firestore-bigquery-export/src/tasks.ts new file mode 100644 index 0000000000..48c4f3fdf6 --- /dev/null +++ b/kits/firestore-bigquery-export/src/tasks.ts @@ -0,0 +1,126 @@ +/* + * Copyright 2019 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { createHash } from "node:crypto"; +import { getFunctions } from "firebase-admin/functions"; +import type { SerializedDocumentChange } from "./handlers"; +import { firestoreLocationToFunctionRegion } from "./region"; + +/** Export name of the write-buffer task function. */ +export const SYNC_BIGQUERY_FUNCTION = "syncBigQuery"; + +const MAX_BACKOFF_MS = 5000; +const BACKOFF_BASE_MS = 100; +const JITTER_MS = 100; + +// Hashed rather than sanitized: a lossy mapping could collapse two event ids +// into one task id, and Cloud Tasks wants ids uniformly distributed. +function taskIdFor(change: SerializedDocumentChange): string { + return createHash("sha256").update(change.eventId).digest("hex"); +} + +/** + * Resolves the queue resource path for a task function of this kit instance. + * + * The name is deliberately unprefixed: firebase-admin >= 14.2.0 resolves the + * deployed `kit--` prefix itself from the + * `FIREBASE_KIT_INSTANCE_ID` env var, which the CLI sets on every deployed kit + * function. All functions of a kit instance deploy to one region, so the + * enqueuing function's own region is the queue's region. + * + * The CLI-set `FUNCTION_REGION` wins because it is the region the function + * was actually deployed to; a `DATABASE_REGION`-derived region can disagree + * with it on a first deploy or when the variable is unset, and is only the + * fallback for local runs where the CLI has not populated the environment. + * With neither set, the bare function name is returned and the Admin SDK + * applies its default location, `us-central1`, which is also where the CLI + * places functions that declare no region. + * + * @param functionName - The export name of the task function. + * @returns The queue resource path, `locations//functions/`, + * or the bare `` when no region is known. + */ +export function syncQueuePath( + functionName: string = SYNC_BIGQUERY_FUNCTION +): string { + const region = + process.env.FUNCTION_REGION || + firestoreLocationToFunctionRegion(process.env.DATABASE_REGION); + + return region + ? `locations/${region}/functions/${functionName}` + : functionName; +} + +function backoffMs(attempt: number, jitter: number): number { + return ( + Math.min(Math.pow(2, attempt) * BACKOFF_BASE_MS, MAX_BACKOFF_MS) + jitter + ); +} + +/** + * Enqueues a payload onto the `syncBigQuery` queue, retrying transient enqueue + * failures in-process with exponential backoff and jitter. + * + * The task id is derived from the event id, so a retried enqueue of an event + * that already reached Cloud Tasks is rejected rather than buffered twice. + * + * @param payload - The serialized change to enqueue. + * @param maxAttempts - How many enqueue attempts to make before giving up. + * Anything but a positive integer means a single attempt: resolving without + * an enqueue would report success for an event that was never buffered. + * @throws The last enqueue error, once every attempt has failed. + */ +export async function enqueueSyncTask( + payload: SerializedDocumentChange, + maxAttempts: number +): Promise { + const queue = getFunctions().taskQueue(syncQueuePath()); + const id = taskIdFor(payload); + + // Math.max(1, NaN) is NaN and would skip the loop entirely. + const attemptBudget = + Number.isInteger(maxAttempts) && maxAttempts >= 1 ? maxAttempts : 1; + const jitter = Math.random() * JITTER_MS; + let attempts = 0; + + while (attempts < attemptBudget) { + if (attempts > 0) { + await new Promise((resolve) => + setTimeout(resolve, backoffMs(attempts, jitter)) + ); + } + + attempts++; + try { + await queue.enqueue(payload, { id }); + return; + } catch (enqueueErr) { + // The event is already buffered; a second task would double-write the row. + // firebase-admin prefixes its codes: `functions/task-already-exists`. + if ( + (enqueueErr as { code?: string })?.code === + "functions/task-already-exists" + ) { + return; + } + + if (attempts >= attemptBudget) { + throw enqueueErr; + } + } + } +} diff --git a/kits/firestore-bigquery-export/tests/tasks.emulator.test.ts b/kits/firestore-bigquery-export/tests/tasks.emulator.test.ts new file mode 100644 index 0000000000..c28570d992 --- /dev/null +++ b/kits/firestore-bigquery-export/tests/tasks.emulator.test.ts @@ -0,0 +1,131 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { createHash } from "node:crypto"; +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, +} from "vitest"; +import type { SerializedDocumentChange } from "../src/handlers"; + +// Real firebase-admin against a local Cloud Tasks emulator host: the kit +// prefix and the default location are resolved by the SDK, which the unit +// suite mocks away. +const INSTANCE_ID = "test-instance"; +const REGION_KEYS = ["DATABASE_REGION", "FUNCTION_REGION"] as const; + +let server: Server; +let requests: { path: string; body: string }[] = []; +const originalEnv: Record = {}; + +beforeAll(async () => { + server = createServer((request, response) => { + let body = ""; + request.on("data", (chunk) => (body += chunk)); + request.on("end", () => { + requests.push({ path: request.url ?? "", body }); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as AddressInfo; + process.env.CLOUD_TASKS_EMULATOR_HOST = `127.0.0.1:${port}`; + process.env.FIREBASE_KIT_INSTANCE_ID = INSTANCE_ID; + const { initializeApp } = await import("firebase-admin/app"); + initializeApp({ + projectId: "test-project", + serviceAccountId: "tasks@test-project.iam.gserviceaccount.com", + credential: { + getAccessToken: async () => ({ + access_token: "owner", + expires_in: 3600, + }), + }, + }); +}); + +afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); +}); + +beforeEach(() => { + requests = []; + for (const key of REGION_KEYS) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } +}); + +afterEach(() => { + for (const key of REGION_KEYS) { + if (originalEnv[key] === undefined) { + delete process.env[key]; + } else { + process.env[key] = originalEnv[key]; + } + } +}); + +function change(eventId: string): SerializedDocumentChange { + return { + timestamp: "2026-01-01T00:00:00.000Z", + eventId, + fullResourceName: "projects/p/databases/(default)/documents/c/d", + changeType: "CREATE", + documentId: "d", + params: null, + data: { a: 1 }, + oldData: undefined, + } as SerializedDocumentChange; +} + +function queueUrl(region: string): string { + return `/projects/test-project/locations/${region}/queues/kit-${INSTANCE_ID}-syncBigQuery/tasks`; +} + +describe("enqueueSyncTask against the Admin SDK", () => { + test("targets the kit-prefixed queue in FUNCTION_REGION", async () => { + process.env.FUNCTION_REGION = "europe-west2"; + const { enqueueSyncTask } = await import("../src/tasks"); + await enqueueSyncTask(change("evt-1"), 1); + expect(requests.map((r) => r.path)).toEqual([queueUrl("europe-west2")]); + }); + + test("falls back to the SDK default location with no region variables", async () => { + const { enqueueSyncTask } = await import("../src/tasks"); + await enqueueSyncTask(change("evt-1"), 1); + expect(requests.map((r) => r.path)).toEqual([queueUrl("us-central1")]); + }); + + test("names the task by the hashed event id", async () => { + process.env.FUNCTION_REGION = "us-central1"; + const { enqueueSyncTask } = await import("../src/tasks"); + await enqueueSyncTask(change("a/b:c d"), 1); + const task = JSON.parse(requests[0].body).task as { name: string }; + expect(task.name).toBe( + `projects/test-project/locations/us-central1/queues/kit-${INSTANCE_ID}-syncBigQuery/tasks/` + + createHash("sha256").update("a/b:c d").digest("hex") + ); + }); +}); diff --git a/kits/firestore-bigquery-export/tests/tasks.test.ts b/kits/firestore-bigquery-export/tests/tasks.test.ts new file mode 100644 index 0000000000..1d7c2e4c3b --- /dev/null +++ b/kits/firestore-bigquery-export/tests/tasks.test.ts @@ -0,0 +1,229 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { createHash } from "node:crypto"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +vi.mock("firebase-admin/functions", () => ({ + getFunctions: vi.fn(), +})); + +import { getFunctions } from "firebase-admin/functions"; +import type { SerializedDocumentChange } from "../src/handlers"; +import { enqueueSyncTask, syncQueuePath } from "../src/tasks"; + +function sha256(value: string): string { + return createHash("sha256").update(value).digest("hex"); +} + +function makeChange( + overrides: Partial = {} +): SerializedDocumentChange { + return { + timestamp: "2026-01-01T00:00:00.000Z", + eventId: "evt-1", + fullResourceName: "projects/p/databases/(default)/documents/c/d", + changeType: "CREATE", + documentId: "d", + params: null, + data: { a: 1 }, + oldData: undefined, + ...overrides, + } as SerializedDocumentChange; +} + +const ENV_KEYS = ["DATABASE_REGION", "FUNCTION_REGION"] as const; +const originalEnv: Record = {}; + +beforeEach(() => { + for (const key of ENV_KEYS) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } +}); + +afterEach(() => { + for (const key of ENV_KEYS) { + if (originalEnv[key] === undefined) { + delete process.env[key]; + } else { + process.env[key] = originalEnv[key]; + } + } + vi.useRealTimers(); + vi.clearAllMocks(); +}); + +describe("syncQueuePath", () => { + test("derives the region from DATABASE_REGION, mapping multi-regions", () => { + process.env.DATABASE_REGION = "nam5"; + expect(syncQueuePath()).toBe( + "locations/us-central1/functions/syncBigQuery" + ); + }); + + test("passes a regional DATABASE_REGION through", () => { + process.env.DATABASE_REGION = "europe-west2"; + expect(syncQueuePath()).toBe( + "locations/europe-west2/functions/syncBigQuery" + ); + }); + + test("prefers the CLI-set FUNCTION_REGION, the region the function is deployed in", () => { + // On a first interactive deploy the functions land in us-central1 while + // DATABASE_REGION already names the target region; the queue is where + // the functions are. + process.env.FUNCTION_REGION = "us-central1"; + process.env.DATABASE_REGION = "europe-west2"; + expect(syncQueuePath()).toBe( + "locations/us-central1/functions/syncBigQuery" + ); + }); + + test("falls back to DATABASE_REGION when FUNCTION_REGION is unset", () => { + process.env.DATABASE_REGION = "europe-west2"; + expect(syncQueuePath()).toBe( + "locations/europe-west2/functions/syncBigQuery" + ); + }); + + test("returns the bare name when neither region variable is set, leaving the location to the SDK default", () => { + expect(syncQueuePath()).toBe("syncBigQuery"); + }); + + test("treats empty region variables as unset", () => { + process.env.DATABASE_REGION = ""; + process.env.FUNCTION_REGION = ""; + expect(syncQueuePath()).toBe("syncBigQuery"); + }); +}); + +describe("enqueueSyncTask", () => { + function mockQueue(enqueue: ReturnType) { + const taskQueue = vi.fn(() => ({ enqueue })); + vi.mocked(getFunctions).mockReturnValue({ + taskQueue, + } as unknown as ReturnType); + return taskQueue; + } + + beforeEach(() => { + process.env.DATABASE_REGION = "us-central1"; + }); + + test("targets the bare function name; the admin SDK adds the kit prefix", async () => { + const enqueue = vi.fn().mockResolvedValue(undefined); + const taskQueue = mockQueue(enqueue); + + const change = makeChange(); + await enqueueSyncTask(change, 3); + + expect(taskQueue).toHaveBeenCalledWith( + "locations/us-central1/functions/syncBigQuery" + ); + expect(enqueue).toHaveBeenCalledTimes(1); + expect(enqueue).toHaveBeenCalledWith(change, { id: sha256("evt-1") }); + }); + + test("hashes the event id into the task id so a retry cannot double-buffer", async () => { + const enqueue = vi.fn().mockResolvedValue(undefined); + mockQueue(enqueue); + + await enqueueSyncTask(makeChange({ eventId: "a/b:c d" }), 3); + + expect(enqueue).toHaveBeenCalledWith(expect.anything(), { + id: sha256("a/b:c d"), + }); + }); + + test("event ids that differ only in punctuation get distinct task ids", async () => { + const enqueue = vi.fn().mockResolvedValue(undefined); + mockQueue(enqueue); + + await enqueueSyncTask(makeChange({ eventId: "a.b" }), 3); + await enqueueSyncTask(makeChange({ eventId: "a/b" }), 3); + + const ids = enqueue.mock.calls.map((call) => call[1].id); + expect(ids[0]).not.toBe(ids[1]); + for (const id of ids) expect(id).toMatch(/^[a-f0-9]{64}$/); + }); + + test("treats an already-enqueued task as success", async () => { + const enqueue = vi.fn().mockRejectedValue( + // Shaped like firebase-admin's PrefixedFirebaseError: `/`. + Object.assign(new Error("exists"), { + code: "functions/task-already-exists", + }) + ); + mockQueue(enqueue); + + await expect(enqueueSyncTask(makeChange(), 3)).resolves.toBeUndefined(); + expect(enqueue).toHaveBeenCalledTimes(1); + }); + + test("a non-positive attempt budget still enqueues once", async () => { + const enqueue = vi.fn().mockResolvedValue(undefined); + mockQueue(enqueue); + + await enqueueSyncTask(makeChange(), 0); + + expect(enqueue).toHaveBeenCalledTimes(1); + }); + + test("a NaN or non-integer attempt budget still enqueues once", async () => { + const enqueue = vi.fn().mockResolvedValue(undefined); + mockQueue(enqueue); + + await enqueueSyncTask(makeChange(), NaN); + await enqueueSyncTask(makeChange(), Infinity); + + expect(enqueue).toHaveBeenCalledTimes(2); + }); + + test("retries a failed enqueue after a backoff and then succeeds", async () => { + vi.useFakeTimers(); + const enqueue = vi + .fn() + .mockRejectedValueOnce(new Error("blip")) + .mockResolvedValueOnce(undefined); + mockQueue(enqueue); + + const pending = enqueueSyncTask(makeChange(), 3); + await vi.runAllTimersAsync(); + await pending; + + expect(enqueue).toHaveBeenCalledTimes(2); + }); + + test("throws the last error once every attempt fails", async () => { + vi.useFakeTimers(); + const enqueue = vi + .fn() + .mockRejectedValueOnce(new Error("first")) + .mockRejectedValueOnce(new Error("second")) + .mockRejectedValue(new Error("last")); + mockQueue(enqueue); + + const pending = enqueueSyncTask(makeChange(), 3); + // Attach the rejection expectation before advancing timers so the + // rejection is never unhandled. + const assertion = expect(pending).rejects.toThrow("last"); + await vi.runAllTimersAsync(); + await assertion; + + expect(enqueue).toHaveBeenCalledTimes(3); + }); +});