Merge pull request #335 from FirebasePrivate/fz/change-tracker-config-cleanup

fredzqm · web-flow · commit 15193a5f010a · 2019-09-17T11:31:51.000-07:00
Clean up FirestoreBigQueryChangeTracker.
diff --git a/firestore-bigquery-export/firestore-bigquery-change-tracker/src/bigquery/index.ts b/firestore-bigquery-export/firestore-bigquery-change-tracker/src/bigquery/index.ts
@@ -25,136 +25,102 @@ import * as logs from "../logs";
 import { BigQuery } from "@google-cloud/bigquery";
 
 export interface FirestoreBigQueryEventHistoryTrackerConfig {
-  collectionPath: string;
   datasetId: string;
-  initialized: boolean;
-  suppressWarnings: boolean;
+  tableId: string;
 }
 
 /**
- * An interface to BigQuery which handles:
- * - Iniitializing the raw changelog table when the first event gets recorded.
- * - Initializing the latest view over the raw changelog.
- * - Streaming writes into the raw changelog table.
+ * An FirestoreEventHistoryTracker that exports data to BigQuery.
+ *
+ * When the first event is received, it creates necessary BigQuery resources:
+ * - Dataset: {@link FirestoreBigQueryEventHistoryTrackerConfig#datasetId}.
+ * - Table: Raw change log table {@link FirestoreBigQueryEventHistoryTracker#rawChangeLogTableName}.
+ * - View: Latest view {@link FirestoreBigQueryEventHistoryTracker#rawLatestView}.
+ * If any subsequent data export fails, it will attempt to reinitialize.
  */
 export class FirestoreBigQueryEventHistoryTracker implements FirestoreEventHistoryTracker {
   bq: bigquery.BigQuery;
-  tableName: string;
-  initialized: boolean;
-  suppressWarnings: boolean;
+  initialized: boolean = false;
 
   constructor(public config: FirestoreBigQueryEventHistoryTrackerConfig) {
     this.bq = new bigquery.BigQuery();
-    this.initialized = config.initialized;
-    this.tableName = config.collectionPath.replace(/\//g, "_");
-    this.suppressWarnings = config.suppressWarnings;
   }
 
   async record(events: FirestoreDocumentChangeEvent[]) {
-    if (!this.config.initialized) {
-      try {
-        await this.initialize(this.config.datasetId, this.tableName);
-        this.initialized = true;
-      } catch (e) {
-        logs.bigQueryErrorRecordingDocumentChange(e);
-      }
-    }
+    await this.initialize();
+
     const rows = events.map(event => {
-      return this.buildDataRow(
-        // Use the function's event ID to protect against duplicate executions
-        event.eventId,
-        event.operation,
-        event.timestamp,
-        event.documentName,
-        event.data);
+      // This must match firestoreToBQTable().
+      return {
+        timestamp: event.timestamp,
+        event_id: event.eventId,
+        document_name: event.documentName,
+        operation: ChangeType[event.operation],
+        data: JSON.stringify(event.data),
+      };
     });
-    await this.insertData(this.config.datasetId, this.tableName, rows);
+    await this.insertData(rows);
   }
 
   /**
-   * Ensure that the defined Firestore schema exists within BigQuery and
-   * contains the correct information. This is invoked for the first time when
-   * the first document change event is recorded.
+   * Inserts rows of data into the BigQuery raw change log table.
    */
-  async initialize(datasetId: string, tableName: string) {
-    const rawTable = raw(tableName);
-
-    await this.initializeDataset(datasetId);
-    await this.initializeChangelog(datasetId, rawTable);
-    await this.initializeLatestView(datasetId, rawTable);
-  };
-
-  buildDataRow(
-    eventId: string,
-    changeType: ChangeType,
-    timestamp: string,
-    document_name: string,
-    data?: Object
-  ): bigquery.RowMetadata {
-    // This must match firestoreToBQTable().
-    return {
-      timestamp: timestamp,
-      event_id: eventId,
-      document_name: document_name,
-      operation: ChangeType[changeType],
-      data: JSON.stringify(data),
-    };
-  };
+  private async insertData(rows: bigquery.RowMetadata[]) {
+    try {
+      const dataset = this.bq.dataset(this.config.datasetId);
+      const table = dataset.table(this.rawChangeLogTableName());
+      logs.dataInserting(rows.length);
+      await table.insert(rows);
+      logs.dataInserted(rows.length);
+    } catch (e) {
+      // Reinitializing in case the destintation table is modified.
+      this.initialized = false;
+      throw e;
+    }
+  }
 
   /**
-   * Insert a row of data into the BigQuery `raw` data table
+   * Creates the BigQuery resources with the expected schema for {@link FirestoreEventHistoryTracker}.
+   * After the first invokation, it skips initialization assuming these resources are still there.
    */
-  async insertData(
-    datasetId: string,
-    collectionTableName: string,
-    rows: bigquery.RowMetadata[]
-  ) {
-    const name = changeLog(raw(collectionTableName));
-    const dataset = this.bq.dataset(datasetId);
-    const table = dataset.table(name);
-    const rowCount = rows.length;
-
-    logs.dataInserting(rowCount);
-    await table.insert(rows);
-    logs.dataInserted(rowCount);
-  };
+  private async initialize() {
+    if (this.initialized) {
+      return;
+    }
+    await this.initializeDataset();
+    await this.initializeRawChangeLogTable();
+    await this.initializeLatestView();
+    this.initialized = true;
+  }
 
   /**
-   * Check that the specified dataset exists, and create it if it doesn't.
+   * Creates the specified dataset if it doesn't already exists.
    */
-  async initializeDataset(datasetId: string): Promise<bigquery.Dataset> {
-    const dataset = this.bq.dataset(datasetId);
+  private async initializeDataset() {
+    const dataset = this.bq.dataset(this.config.datasetId);
     const [datasetExists] = await dataset.exists();
     if (datasetExists) {
-      if (!this.suppressWarnings) {
-        logs.bigQueryDatasetExists(datasetId);
-      }
+      logs.bigQueryDatasetExists(this.config.datasetId);
     } else {
-      logs.bigQueryDatasetCreating(datasetId);
+      logs.bigQueryDatasetCreating(this.config.datasetId);
       await dataset.create();
-      logs.bigQueryDatasetCreated(datasetId);
+      logs.bigQueryDatasetCreated(this.config.datasetId);
     }
     return dataset;
-  };
+  }
 
   /**
-   * Check that the table exists within the specified dataset, and create it
-   * if it doesn't.  If the table does exist, validate that the BigQuery schema
-   * is correct and add any missing fields.
+   * Creates the raw change log table if it doesn't already exist.
+   * TODO: Validate that the BigQuery schema is correct if the table does exist,
    */
-  async initializeChangelog(
-    datasetId: string,
-    tableName: string,
-  ): Promise<bigquery.Table> {
-    const changelogName = changeLog(tableName);
-    const dataset = this.bq.dataset(datasetId);
-    let table = dataset.table(changelogName);
+  private async initializeRawChangeLogTable() {
+    const changelogName = this.rawChangeLogTableName();
+    const dataset = this.bq.dataset(this.config.datasetId);
+    const table = dataset.table(changelogName);
     const [tableExists] = await table.exists();
 
     if (tableExists) {
-      if (!this.suppressWarnings) {
-        logs.bigQueryTableAlreadyExists(table.id, dataset.id);
-      }
+      logs.bigQueryTableAlreadyExists(table.id, dataset.id);
     } else {
       logs.bigQueryTableCreating(changelogName);
       const options = {
@@ -169,37 +135,35 @@ export class FirestoreBigQueryEventHistoryTracker implements FirestoreEventHisto
   };
 
   /**
-   * Create a view over a table storing a change log of Firestore documents
-   * which contains only latest version of all live documents in the mirrored
-   * collection.
+   * Creates the latest snapshot view, which returns only latest operations
+   * of all existing documents over the raw change log table.
    */
-  async initializeLatestView(
-    datasetId: string,
-    tableName: string
-  ): Promise<bigquery.Table> {
-    let viewName = latest(tableName);
-    const dataset = this.bq.dataset(datasetId);
-    let view = dataset.table(viewName);
+  private async initializeLatestView() {
+    const dataset = this.bq.dataset(this.config.datasetId);
+    const view = dataset.table(this.rawLatestView());
     const [viewExists] = await view.exists();
 
     if (viewExists) {
-      if (!this.suppressWarnings) {
-        logs.bigQueryViewAlreadyExists(view.id, dataset.id);
-      }
+      logs.bigQueryViewAlreadyExists(view.id, dataset.id);
     } else {
-      const latestSnapshot = latestConsistentSnapshotView(datasetId, changeLog(tableName));
-      logs.bigQueryViewCreating(viewName, latestSnapshot.query);
+      const latestSnapshot = latestConsistentSnapshotView(this.config.datasetId, this.rawChangeLogTableName());
+      logs.bigQueryViewCreating(this.rawLatestView(), latestSnapshot.query);
       const options = {
-        friendlyName: viewName,
+        friendlyName: this.rawLatestView(),
         view: latestSnapshot,
       };
       await view.create(options);
-      logs.bigQueryViewCreated(viewName);
+      logs.bigQueryViewCreated(this.rawLatestView());
     }
     return view;
-  };
+  }
+
+  private rawChangeLogTableName(): string {
+    return `${this.config.tableId}_raw_changelog`;
+  }
+
+  private rawLatestView(): string {
+    return `${this.config.tableId}_raw_latest`;
+  }
 }
 
-export function raw(tableName: string): string { return `${tableName}_raw`; };
-export function changeLog(tableName: string): string { return `${tableName}_changelog`; }
-export function latest(tableName: string): string { return `${tableName}_latest`; };
diff --git a/firestore-bigquery-export/firestore-bigquery-change-tracker/src/bigquery/schema.ts b/firestore-bigquery-export/firestore-bigquery-change-tracker/src/bigquery/schema.ts
@@ -17,7 +17,6 @@
 import * as bigquery from "@google-cloud/bigquery";
 import * as errors from "../errors";
 import * as logs from "../logs";
-import { changeLog, raw, latest } from "../bigquery";
 import * as sqlFormatter from "sql-formatter";
 
 export type BigQueryFieldMode = "NULLABLE" | "REPEATED" | "REQUIRED";
diff --git a/firestore-bigquery-export/functions/lib/index.js b/firestore-bigquery-export/functions/lib/index.js
@@ -30,10 +30,8 @@ const firestore_bigquery_change_tracker_1 = require("@firebaseextensions/firesto
 const logs = require("./logs");
 const util_1 = require("./util");
 const eventTracker = new firestore_bigquery_change_tracker_1.FirestoreBigQueryEventHistoryTracker({
-    collectionPath: config_1.default.tableId,
+    tableId: config_1.default.tableId,
     datasetId: config_1.default.datasetId,
-    initialized: false,
-    suppressWarnings: false,
 });
 logs.init();
 exports.fsexportbigquery = functions.handler.firestore.document.onWrite((change, context) => __awaiter(void 0, void 0, void 0, function* () {
diff --git a/firestore-bigquery-export/functions/src/index.ts b/firestore-bigquery-export/functions/src/index.ts
@@ -25,10 +25,8 @@ import * as logs from "./logs";
 import { getChangeType } from "./util";
 
 const eventTracker: FirestoreEventHistoryTracker = new FirestoreBigQueryEventHistoryTracker({
-  collectionPath: config.tableId, // TODO: rename collectionPath.
+  tableId: config.tableId,
   datasetId: config.datasetId,
-  initialized: false, // TODO: Remove this config.
-  suppressWarnings: false, // TODO: Remove this config.
 });
 
 logs.init();
diff --git a/firestore-bigquery-export/scripts/import/src/index.ts b/firestore-bigquery-export/scripts/import/src/index.ts
@@ -126,10 +126,8 @@ const run = async (): Promise<number> => {
   process.env.GOOGLE_CLOUD_PROJECT = projectId;
 
   const dataSink = new FirestoreBigQueryEventHistoryTracker({
-    collectionPath: destinationCollectionPath,
+    tableId: destinationCollectionPath,
     datasetId: datasetId,
-    initialized: false,
-    suppressWarnings: true,
   });
 
   console.log(
