@firebase-function-kits/bigquery-firestore-export 0.0.1 → 0.0.2-rc.1

package/src/handlers.ts

@@ -0,0 +1,141 @@
+/*
+ * 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 type { BigQuery } from "@google-cloud/bigquery";
+import type { PubSub } from "@google-cloud/pubsub";
+import type { Firestore } from "firebase-admin/firestore";
+import type { CloudEvent } from "firebase-functions/v2";
+import type { MessagePublishedData } from "firebase-functions/v2/pubsub";
+import {
+  createTransferConfig,
+  type DataTransferClient,
+  getTransferConfig,
+  updateTransferConfig,
+} from "./dts";
+import type { ResolvedBigqueryFirestoreExportConfig } from "./export-config";
+import { handleTransferRunMessage, parseTransferConfigName } from "./helper";
+import * as logs from "./logs";
+import type { TransferRunPayload } from "./types";
+
+export type TransferRunEvent = CloudEvent<
+  MessagePublishedData<TransferRunPayload>
+>;
+
+/** External services used by both handlers, injected for testability. */
+export interface HandlerContext {
+  db: Firestore;
+  bigquery: BigQuery;
+  dataTransfer: DataTransferClient;
+  pubsub: PubSub;
+  config: ResolvedBigqueryFirestoreExportConfig;
+}
+
+async function ensureNotificationTopic(ctx: HandlerContext): Promise<void> {
+  const topic = ctx.pubsub.topic(ctx.config.pubSubTopic);
+  const [exists] = await topic.exists();
+  if (exists) return;
+
+  try {
+    await ctx.pubsub.createTopic(ctx.config.pubSubTopic);
+    logs.topicCreated(ctx.config.pubSubTopic);
+  } catch (err) {
+    // A concurrent init can create the same topic between exists() and create().
+    if (
+      typeof err !== "object" ||
+      err === null ||
+      !("code" in err) ||
+      err.code !== 6
+    ) {
+      throw err;
+    }
+  }
+}
+
+async function storeTransferConfig(
+  ctx: HandlerContext,
+  transferConfig: Awaited<ReturnType<typeof getTransferConfig>>
+): Promise<void> {
+  if (!transferConfig?.name) {
+    throw new Error("BigQuery transfer config is missing its resource name");
+  }
+  const { transferConfigId } = parseTransferConfigName(transferConfig.name);
+  await ctx.db
+    .collection(ctx.config.firestoreCollection)
+    .doc(transferConfigId)
+    .set({ extInstanceId: ctx.config.instanceId, ...transferConfig });
+}
+
+/** Handles a v2 Pub/Sub completion notification from BigQuery DTS. */
+export async function handleMessagePublished(
+  event: TransferRunEvent,
+  ctx: HandlerContext
+): Promise<void> {
+  logs.start();
+  try {
+    await handleTransferRunMessage(ctx, { json: event.data.message.json });
+    logs.complete();
+  } catch (err) {
+    logs.error(err);
+    throw err;
+  }
+}
+
+/** Idempotently creates, links, or updates this deployment's DTS config. */
+export async function handleUpsertTransferConfig(
+  ctx: HandlerContext
+): Promise<void> {
+  await ensureNotificationTopic(ctx);
+
+  if (ctx.config.transferConfigName) {
+    const linked = await getTransferConfig(
+      ctx.dataTransfer,
+      ctx.config.transferConfigName
+    );
+    if (!linked) {
+      throw new Error(
+        `Transfer config not found: ${ctx.config.transferConfigName}`
+      );
+    }
+    await storeTransferConfig(ctx, linked);
+    return;
+  }
+
+  const existing = await ctx.db
+    .collection(ctx.config.firestoreCollection)
+    .where("extInstanceId", "==", ctx.config.instanceId)
+    .limit(1)
+    .get();
+
+  if (existing.empty) {
+    const created = await createTransferConfig(ctx.dataTransfer, ctx.config);
+    await storeTransferConfig(ctx, created);
+    return;
+  }
+
+  const transferConfigName = existing.docs[0].data().name;
+  if (typeof transferConfigName !== "string" || !transferConfigName) {
+    throw new Error(
+      `Existing transfer config document in ${ctx.config.firestoreCollection} is missing required 'name' field.`
+    );
+  }
+
+  const updated = await updateTransferConfig(
+    ctx.dataTransfer,
+    transferConfigName,
+    ctx.config
+  );
+  await storeTransferConfig(ctx, updated);
+}

package/src/helper.ts

@@ -0,0 +1,320 @@
+/*
+ * Copyright 2025 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 {
+  type BigQuery,
+  BigQueryDate,
+  BigQueryDatetime,
+  BigQueryTime,
+  BigQueryTimestamp,
+  Geography,
+} from "@google-cloud/bigquery";
+import {
+  type DocumentData,
+  type DocumentReference,
+  type Firestore,
+  Timestamp,
+} from "firebase-admin/firestore";
+import type { ResolvedBigqueryFirestoreExportConfig } from "./export-config";
+import * as logs from "./logs";
+import type {
+  BigQueryRow,
+  BigQueryRowValue,
+  FirestoreRow,
+  FirestoreRowValue,
+  TransferRunMessage,
+} from "./types";
+
+export interface ParsedTransferRunName {
+  projectId: string;
+  location: string;
+  transferConfigId: string;
+  runId: string;
+}
+
+export interface ParsedTransferConfigName {
+  projectId: string;
+  location: string;
+  transferConfigId: string;
+}
+
+export interface ResultHandlerContext {
+  db: Firestore;
+  bigquery: BigQuery;
+  config: ResolvedBigqueryFirestoreExportConfig;
+}
+
+const TRANSFER_RUN_NAME_REGEX =
+  /^projects\/([^/]+)\/locations\/([^/]+)\/transferConfigs\/([^/]+)\/runs\/([^/]+)$/;
+const TRANSFER_CONFIG_NAME_REGEX =
+  /^projects\/([^/]+)\/locations\/([^/]+)\/transferConfigs\/([^/]+)$/;
+const FIRESTORE_WRITE_CHUNK_SIZE = 10_000;
+
+export function parseTransferRunName(name: string): ParsedTransferRunName {
+  const match = name.match(TRANSFER_RUN_NAME_REGEX);
+  if (!match) {
+    throw new Error(
+      `Invalid transfer run name format: "${name}". Expected format: projects/{projectId}/locations/{location}/transferConfigs/{configId}/runs/{runId}`
+    );
+  }
+
+  return {
+    projectId: match[1],
+    location: match[2],
+    transferConfigId: match[3],
+    runId: match[4],
+  };
+}
+
+export function parseTransferConfigName(
+  name: string
+): ParsedTransferConfigName {
+  const match = name.match(TRANSFER_CONFIG_NAME_REGEX);
+  if (!match) {
+    throw new Error(
+      `Invalid transfer config name format: "${name}". Expected format: projects/{projectId}/locations/{location}/transferConfigs/{configId}`
+    );
+  }
+
+  return {
+    projectId: match[1],
+    location: match[2],
+    transferConfigId: match[3],
+  };
+}
+
+export async function updateLatestRunDocument(
+  db: Firestore,
+  config: ResolvedBigqueryFirestoreExportConfig,
+  transferConfigId: string,
+  runId: string,
+  message: TransferRunMessage,
+  rowCounts: { failedRowCount: number; totalRowCount: number }
+): Promise<void> {
+  const latestRef = db
+    .collection(`${config.firestoreCollection}/${transferConfigId}/runs`)
+    .doc("latest");
+  const runTime = new Date(message.json.runTime);
+  const docUpdate = {
+    runMetadata: message.json,
+    latestRunId: runId,
+    ...rowCounts,
+  };
+
+  await db.runTransaction(async (transaction) => {
+    const latest = await transaction.get(latestRef);
+    const latestData = latest.data();
+    const existingRunTime = latestData?.runMetadata?.runTime;
+    const existingRunId = latestData?.latestRunId;
+    const shouldUpdate =
+      !latestData ||
+      !existingRunTime ||
+      new Date(existingRunTime) < runTime ||
+      existingRunId === runId;
+
+    if (shouldUpdate) {
+      transaction.set(latestRef, docUpdate);
+      return;
+    }
+
+    logs.latestDocUpdateSkipped(
+      transferConfigId,
+      runId,
+      `existing run is newer (${existingRunTime} >= ${message.json.runTime})`
+    );
+  });
+}
+
+export async function getBigqueryResults(
+  bigquery: BigQuery,
+  config: ResolvedBigqueryFirestoreExportConfig,
+  transferConfigId: string,
+  runId: string,
+  datasetId: string,
+  tableName: string
+): Promise<BigQueryRow[]> {
+  const query = `SELECT * FROM \`${config.projectId}.${datasetId}.${tableName}\``;
+
+  try {
+    const [job] = await bigquery.createQueryJob({
+      query,
+      location: config.bigqueryDatasetLocation,
+    });
+    logs.bigqueryJobStarted(job.id);
+    const [rows] = await job.getQueryResults();
+    logs.bigqueryResultsRowCount(transferConfigId, runId, rows.length);
+    return rows as BigQueryRow[];
+  } catch (err) {
+    logs.bigqueryQueryFailed(transferConfigId, runId, tableName, err);
+    throw err;
+  }
+}
+
+export function convertUnsupportedDataTypes(row: null): null;
+export function convertUnsupportedDataTypes(row: string): string;
+export function convertUnsupportedDataTypes(row: number): number;
+export function convertUnsupportedDataTypes(row: boolean): boolean;
+export function convertUnsupportedDataTypes(row: BigQueryRow): FirestoreRow;
+export function convertUnsupportedDataTypes(
+  row: BigQueryRowValue
+): FirestoreRowValue;
+export function convertUnsupportedDataTypes(
+  row: BigQueryRowValue
+): FirestoreRowValue {
+  if (row === null || typeof row !== "object") {
+    return row as FirestoreRowValue;
+  }
+
+  if (
+    row instanceof BigQueryTimestamp ||
+    row instanceof BigQueryDate ||
+    row instanceof BigQueryTime ||
+    row instanceof BigQueryDatetime
+  ) {
+    return Timestamp.fromDate(new Date(row.value));
+  }
+  if (row instanceof Date) return Timestamp.fromDate(row);
+  if (row instanceof Buffer) return new Uint8Array(row);
+  if (row instanceof Geography) return row.value;
+  if (Array.isArray(row)) {
+    return row.map((value) => convertUnsupportedDataTypes(value));
+  }
+
+  return Object.fromEntries(
+    Object.entries(row).map(([key, value]) => [
+      key,
+      convertUnsupportedDataTypes(value),
+    ])
+  ) as FirestoreRow;
+}
+
+export async function writeRunResultsToFirestore(
+  ctx: ResultHandlerContext,
+  message: TransferRunMessage
+): Promise<void> {
+  const { db, bigquery, config } = ctx;
+  const { transferConfigId, runId } = parseTransferRunName(message.json.name);
+  const runTime = new Date(message.json.runTime);
+  const runTimeSuffix = [
+    runTime.getUTCHours(),
+    runTime.getUTCMinutes(),
+    runTime.getUTCSeconds(),
+  ]
+    .map((part) => String(part).padStart(2, "0"))
+    .join("");
+  const tableName = message.json.params.destination_table_name_template.replace(
+    '{run_time|"%H%M%S"}',
+    runTimeSuffix
+  );
+  const rows = await getBigqueryResults(
+    bigquery,
+    config,
+    transferConfigId,
+    runId,
+    message.json.destinationDatasetId,
+    tableName
+  );
+  const collection = db.collection(
+    `${config.firestoreCollection}/${transferConfigId}/runs/${runId}/output`
+  );
+  let succeededRowCount = 0;
+
+  for (let i = 0; i < rows.length; i += FIRESTORE_WRITE_CHUNK_SIZE) {
+    const writes: Array<Promise<DocumentReference<DocumentData>>> = [];
+    for (
+      let j = i;
+      j < i + FIRESTORE_WRITE_CHUNK_SIZE && j < rows.length;
+      j++
+    ) {
+      writes.push(collection.add(convertUnsupportedDataTypes(rows[j])));
+    }
+
+    const results = await Promise.allSettled(writes);
+    for (const result of results) {
+      if (result.status === "fulfilled") succeededRowCount++;
+      else logs.errorWritingToFirestore(result.reason);
+    }
+  }
+
+  const rowCounts = {
+    failedRowCount: rows.length - succeededRowCount,
+    totalRowCount: rows.length,
+  };
+  logs.runResultsWrittenToFirestore(runId, succeededRowCount, rows.length);
+
+  await db
+    .collection(`${config.firestoreCollection}/${transferConfigId}/runs`)
+    .doc(runId)
+    .set({ runMetadata: message.json, ...rowCounts });
+  await updateLatestRunDocument(
+    db,
+    config,
+    transferConfigId,
+    runId,
+    message,
+    rowCounts
+  );
+}
+
+export async function transferConfigAssociatedWithInstance(
+  db: Firestore,
+  config: ResolvedBigqueryFirestoreExportConfig,
+  transferConfigId: string
+): Promise<boolean> {
+  const results = await db
+    .collection(config.firestoreCollection)
+    .where("extInstanceId", "==", config.instanceId)
+    .get();
+  return results.docs.some((doc) => doc.id === transferConfigId);
+}
+
+export async function handleTransferRunMessage(
+  ctx: ResultHandlerContext,
+  message: TransferRunMessage
+): Promise<void> {
+  const { transferConfigId, runId } = parseTransferRunName(message.json.name);
+  const associated = await transferConfigAssociatedWithInstance(
+    ctx.db,
+    ctx.config,
+    transferConfigId
+  );
+
+  if (!associated) {
+    throw new Error(
+      `Skipping handling pubsub message because transferConfig '${transferConfigId}' is not associated with extension instance '${ctx.config.instanceId}'.`
+    );
+  }
+
+  if (message.json.state === "SUCCEEDED") {
+    await writeRunResultsToFirestore(ctx, message);
+    return;
+  }
+
+  logs.handlingNonSuccessRun(transferConfigId, runId, message.json.state);
+  const rowCounts = { failedRowCount: 0, totalRowCount: 0 };
+  await ctx.db
+    .collection(`${ctx.config.firestoreCollection}/${transferConfigId}/runs`)
+    .doc(runId)
+    .set({ runMetadata: message.json, ...rowCounts });
+  await updateLatestRunDocument(
+    ctx.db,
+    ctx.config,
+    transferConfigId,
+    runId,
+    message,
+    rowCounts
+  );
+}

package/src/index.ts

@@ -0,0 +1,126 @@
+/*
+ * 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.
+ */
+
+/**
+ * Main entry point. Exports the wired Pub/Sub and lifecycle task functions,
+ * while resolving concrete config and clients lazily at runtime. Import from
+ * `./lib` for the side-effect-free handlers, helpers, and config types.
+ */
+
+import { BigQuery } from "@google-cloud/bigquery";
+import { v1 as bigqueryDataTransfer } from "@google-cloud/bigquery-data-transfer";
+import { PubSub } from "@google-cloud/pubsub";
+import { getApps, initializeApp } from "firebase-admin/app";
+import { getFirestore } from "firebase-admin/firestore";
+import { onMessagePublished } from "firebase-functions/v2/pubsub";
+import { onTaskDispatched } from "firebase-functions/v2/tasks";
+import type { Role } from "firebase-functions/v2";
+import { requiresAPI, requiresRole } from "firebase-functions/v2";
+import {
+  afterFirstDeploy,
+  afterRedeploy,
+} from "firebase-functions/v2/lifecycle";
+import { CONFIG_EXPRESSIONS, configFromEnv } from "./config";
+import type { ResolvedBigqueryFirestoreExportConfig } from "./export-config";
+import { resolveConfig } from "./export-config";
+import {
+  type HandlerContext,
+  handleMessagePublished,
+  handleUpsertTransferConfig,
+} from "./handlers";
+import * as logs from "./logs";
+import type { TransferRunPayload } from "./types";
+
+export * from "./lib";
+
+const UPSERT_FUNCTION = "upsertTransferConfig";
+const REQUIRED_ROLES: ReadonlyArray<Role> = [
+  "roles/datastore.user",
+  "roles/bigquery.admin",
+  "roles/pubsub.admin",
+  "roles/eventarc.eventReceiver",
+  "roles/run.invoker",
+];
+const REQUIRED_APIS = [
+  {
+    api: "bigquery.googleapis.com",
+    reason: "Runs scheduled queries and reads their destination tables.",
+  },
+  {
+    api: "bigquerydatatransfer.googleapis.com",
+    reason: "Creates and reconciles the scheduled-query transfer config.",
+  },
+  {
+    api: "pubsub.googleapis.com",
+    reason: "Delivers BigQuery transfer-completion notifications.",
+  },
+] as const;
+
+for (const role of REQUIRED_ROLES) {
+  requiresRole(role);
+}
+
+for (const { api, reason } of REQUIRED_APIS) {
+  requiresAPI(api, reason);
+}
+
+afterFirstDeploy({
+  task: { function: UPSERT_FUNCTION, body: { data: {} } },
+});
+afterRedeploy({
+  task: { function: UPSERT_FUNCTION, body: { data: {} } },
+});
+
+let ctx: HandlerContext | undefined;
+
+function getContext(): HandlerContext {
+  if (ctx) return ctx;
+
+  const config: ResolvedBigqueryFirestoreExportConfig = resolveConfig(
+    configFromEnv()
+  );
+  if (getApps().length === 0) initializeApp({ projectId: config.projectId });
+  logs.init(config);
+
+  ctx = {
+    db: getFirestore(),
+    bigquery: new BigQuery({ projectId: config.projectId }),
+    dataTransfer: new bigqueryDataTransfer.DataTransferServiceClient({
+      projectId: config.projectId,
+    }),
+    pubsub: new PubSub({ projectId: config.projectId }),
+    config,
+  };
+  return ctx;
+}
+
+/** Consumes BigQuery Data Transfer completion notifications. */
+export const processMessages = onMessagePublished<TransferRunPayload>(
+  {
+    topic: CONFIG_EXPRESSIONS.pubSubTopic,
+    retry: false,
+  },
+  (event) => handleMessagePublished(event, getContext())
+);
+
+/** Creates, links, or reconciles this deployment's scheduled query. */
+export const upsertTransferConfig = onTaskDispatched(
+  {
+    memory: "1GiB",
+    retryConfig: { maxAttempts: 5, minBackoffSeconds: 30 },
+  },
+  () => handleUpsertTransferConfig(getContext())
+);

package/src/lib.ts

@@ -0,0 +1,69 @@
+/*
+ * 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.
+ */
+
+/**
+ * Side-effect-free library surface. Importing this module registers no
+ * functions, lifecycle hooks, roles, or APIs and reads no environment params.
+ */
+
+export {
+  constructUpdateTransferConfigRequest,
+  createTransferConfig,
+  createTransferConfigRequest,
+  type DataTransferClient,
+  getTransferConfig,
+  PARTITIONING_FIELD_REMOVAL_ERROR,
+  PARTITIONING_FIELD_REMOVAL_ERROR_PREFIX,
+  type TransferConfig,
+  updateTransferConfig,
+} from "./dts";
+export {
+  type BigqueryFirestoreExportConfig,
+  type DeployTimeOptions,
+  type LogLevel,
+  type ResolvedBigqueryFirestoreExportConfig,
+  resolveConfig,
+} from "./export-config";
+export {
+  type HandlerContext,
+  handleMessagePublished,
+  handleUpsertTransferConfig,
+  type TransferRunEvent,
+} from "./handlers";
+export {
+  convertUnsupportedDataTypes,
+  getBigqueryResults,
+  handleTransferRunMessage,
+  type ParsedTransferConfigName,
+  type ParsedTransferRunName,
+  parseTransferConfigName,
+  parseTransferRunName,
+  type ResultHandlerContext,
+  transferConfigAssociatedWithInstance,
+  updateLatestRunDocument,
+  writeRunResultsToFirestore,
+} from "./helper";
+export { metadata } from "./metadata";
+export type {
+  BigQueryRow,
+  BigQueryRowValue,
+  FirestoreRow,
+  FirestoreRowValue,
+  TransferRunMessage,
+  TransferRunParams,
+  TransferRunPayload,
+  TransferRunState,
+} from "./types";