fix(firstore-bigquery-export): comment out unused role for now and use logging

Author: cabljac

firestore-bigquery-export/extension.yaml

@@ -48,8 +48,8 @@ roles:
   - role: datastore.user
     reason: Allows the extension to write updates to the database.
 
-  - role: storage.objectAdmin
-    reason: Allows the extension to create objects in the storage bucket.
+  # - role: storage.objectAdmin
+  #   reason: Allows the extension to create objects in the storage bucket.
 
 resources:
   - name: fsexportbigquery
@@ -207,7 +207,7 @@ params:
         value: yes
       - label: No
         value: no
-    required: true
+    default: yes
 
   - param: WILDCARD_IDS
     label: Enable Wildcard Column field with Parent Firestore Document IDs

firestore-bigquery-export/functions/src/index.ts

@@ -310,7 +310,7 @@ async function attemptToEnqueue(
     await events.recordErrorEvent(enqueueErr as Error);
 
     // Log the error if it has not been logged already.
-    if (!enqueueErr.logged) {
+    if (!enqueueErr.logged && config.logFailedExportData) {
       logs.error(
         true,
         "Failed to enqueue task to syncBigQuery",

firestore-bigquery-export/functions/stress_test/main.js

@@ -1,29 +1,57 @@
 const { Worker } = require("worker_threads");
 const { performance } = require("perf_hooks");
 const path = require("path");
+const admin = require("firebase-admin");
 
-const totalDocs = 10000000; // Total number of documents to write
+// Initialize Firebase Admin SDK
+admin.initializeApp({
+  projectId: "vertex-testing-1efc3",
+});
+
+// Get a reference to the Firestore service
+const db = admin.firestore();
+
+const totalDocs = 1000000; // Total number of documents to write
 const maxThreads = 20; // Maximum number of worker threads
 const batchSize = 500; // Documents per batch
-const rampUpDelay = 2000; // 5 seconds delay between ramp-ups
-const rampUps = 20; // Number of ramp-ups (planned)
-
+const targetRate = 500; // Target docs per second
+const rampUpDelay = 1000; // Delay between ramp-ups
+const rampUps = 5; // Number of ramp-ups
 const docsPerRampUp = Math.ceil(totalDocs / rampUps); // Documents per ramp-up
 
+// Calculate the delay needed to meet the target rate (in milliseconds)
+const delayBetweenBatches = Math.max(1000 / (targetRate / batchSize), 0); // Delay between batches in ms
+
+// Hardcoded collection paths with the form: A/{aid}/B/{bid}/C/{cid}/D/{did}/E/{eid}/F/{fid}/G
+const collectionPaths = [
+  "A/aid1/B/bid1/C/cid1/D/did1/E/eid1/F/fid1/G",
+  "A/aid2/B/bid2/C/cid2/D/did2/E/eid2/F/fid2/G",
+  "A/aid3/B/bid3/C/cid3/D/did3/E/eid3/F/fid3/G",
+  "A/aid4/B/bid4/C/cid4/D/did4/E/eid4/F/fid4/G",
+  "A/aid5/B/bid5/C/cid5/D/did5/E/eid5/F/fid5/G",
+];
+
 // Start measuring total execution time
 const totalStartTime = performance.now();
 
 const workerJsPath = path.resolve(__dirname, "worker.js");
 
 // Function to spawn worker threads for a specific ramp-up
-const spawnWorkers = async (activeThreads, startDoc, docsPerRampUp) => {
-  console.log(`Spawning ${activeThreads} worker(s)...`);
+const spawnWorkers = async (
+  activeThreads,
+  startDoc,
+  docsPerRampUp,
+  collectionPath
+) => {
+  console.log(
+    `Spawning ${activeThreads} worker(s) for collection ${collectionPath}...`
+  );
   let promises = [];
   const docsPerThread = Math.ceil(docsPerRampUp / activeThreads);
 
   for (let i = 0; i < activeThreads; i++) {
     const docsForThisThread = Math.min(docsPerThread, docsPerRampUp);
-    const start = startDoc + i * docsPerThread;
+    const start = startDoc + i * docsForThisThread;
     const end = Math.min(start + docsForThisThread, startDoc + docsPerRampUp);
 
     promises.push(
@@ -33,6 +61,8 @@ const spawnWorkers = async (activeThreads, startDoc, docsPerRampUp) => {
             start,
             end,
             batchSize,
+            collectionPath, // Pass the collection path to the worker
+            delayBetweenBatches, // Pass the delay to the worker
           },
         });
 
@@ -64,13 +94,44 @@ const spawnWorkers = async (activeThreads, startDoc, docsPerRampUp) => {
   }
 };
 
+// Function to query Firestore for the total document count using count() aggregation
+const getCollectionCounts = async () => {
+  let counts = {};
+
+  for (const collectionPath of collectionPaths) {
+    const collectionRef = db.collection(collectionPath);
+    const snapshot = await collectionRef.count().get(); // Use the count aggregation query
+    const count = snapshot.data().count;
+    counts[collectionPath] = count;
+    console.log(`Collection ${collectionPath} has ${count} documents.`);
+  }
+
+  return counts;
+};
+
+// Function to calculate the difference between two count objects
+const calculateCountDifference = (beforeCounts, afterCounts) => {
+  let totalDifference = 0;
+
+  for (const collectionPath in beforeCounts) {
+    const beforeCount = beforeCounts[collectionPath] || 0;
+    const afterCount = afterCounts[collectionPath] || 0;
+    const difference = afterCount - beforeCount;
+    console.log(`Collection ${collectionPath} difference: ${difference}`);
+    totalDifference += difference;
+  }
+
+  return totalDifference;
+};
+
 // Function to execute ramp-ups
 const executeRampUps = async () => {
   let activeThreads = 1;
   let startDoc = 0;
 
   for (let i = 0; i < rampUps; i++) {
-    await spawnWorkers(activeThreads, startDoc, docsPerRampUp);
+    const collectionPath = collectionPaths[i % collectionPaths.length]; // Rotate through collections
+    await spawnWorkers(activeThreads, startDoc, docsPerRampUp, collectionPath);
     startDoc += docsPerRampUp;
 
     if (activeThreads < maxThreads) {
@@ -88,17 +149,38 @@ const executeRampUps = async () => {
   }
 };
 
-// Run the ramp-ups
-executeRampUps()
-  .then(() => {
+// Main execution flow
+const main = async () => {
+  try {
+    // Count documents before writing
+    console.log("Counting documents before the operation...");
+    const beforeCounts = await getCollectionCounts();
+
+    // Perform the writing operation
+    await executeRampUps();
+
+    // Count documents after writing
+    console.log("Counting documents after the operation...");
+    const afterCounts = await getCollectionCounts();
+
+    // Calculate and log the difference
+    const totalDocsWritten = calculateCountDifference(
+      beforeCounts,
+      afterCounts
+    );
+    console.log(`Total documents written: ${totalDocsWritten}`);
+
     const totalEndTime = performance.now();
     const totalDuration = (totalEndTime - totalStartTime) / 1000; // Convert to seconds
     console.log(
-      `Successfully written ${totalDocs} documents to the collection in ${totalDuration.toFixed(
+      `Successfully written ${totalDocsWritten} documents in ${totalDuration.toFixed(
         2
       )} seconds.`
     );
-  })
-  .catch((error) => {
-    console.error("Error in worker threads: ", error);
-  });
+  } catch (error) {
+    console.error("Error during execution: ", error);
+  }
+};
+
+// Run the main function
+main();

firestore-bigquery-export/functions/stress_test/worker.js

@@ -10,29 +10,48 @@ admin.initializeApp({
 
 // Get a reference to the Firestore service
 const db = admin.firestore();
-const collectionName = "posts_2";
 
-// Generate a random document
+// Generate a large random document closer to 1MB
 const generateRandomDocument = () => {
+  // const largeString = "x".repeat(300000); // A string of 300,000 characters (~300 KB)
+  // const largeArray = new Array(5000).fill().map((_, i) => ({
+  //   index: i,
+  //   value: `Value_${Math.random().toString(36).substring(7)}`,
+  // }));
+
   return {
     id: uuidv4(),
     name: `Name_${Math.random().toString(36).substring(7)}`,
     age: Math.floor(Math.random() * 60) + 18, // Random age between 18 and 78
     email: `user_${Math.random().toString(36).substring(7)}@example.com`,
     isActive: Math.random() > 0.5, // Random boolean value
     createdAt: admin.firestore.Timestamp.now(),
+    // largeString, // Large string field
+    // largeArray, // Large array field
   };
 };
 
-// Write a batch of documents to Firestore
-const writeBatch = async (start, end, batchSize) => {
+// Delay function for rate control
+const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+
+// Write a batch of documents to a specific collection in Firestore
+const writeBatch = async (
+  start,
+  end,
+  batchSize,
+  collectionPath,
+  delayBetweenBatches
+) => {
   let count = start;
   while (count < end) {
     const batchStartTime = performance.now();
 
     let batch = db.batch();
-    for (let i = 0; i < batchSize && count < end; i++) {
-      let docRef = db.collection(collectionName).doc();
+    const remainingDocs = end - count;
+    const adjustedBatchSize = Math.min(batchSize, remainingDocs); // Adjust batch size if remaining docs < batchSize
+
+    for (let i = 0; i < adjustedBatchSize && count < end; i++) {
+      let docRef = db.collection(collectionPath).doc();
       batch.set(docRef, generateRandomDocument());
       count++;
     }
@@ -42,15 +61,24 @@ const writeBatch = async (start, end, batchSize) => {
     const batchEndTime = performance.now();
     const batchDuration = (batchEndTime - batchStartTime) / 1000; // Convert to seconds
     parentPort.postMessage(
-      `Batch of ${batchSize} documents written in ${batchDuration.toFixed(
+      `Batch of ${adjustedBatchSize} documents written in ${batchDuration.toFixed(
         2
-      )} seconds.`
+      )} seconds to ${collectionPath}.`
     );
+
+    // Introduce delay between batches to meet target rate
+    await delay(delayBetweenBatches);
   }
 };
 
 // Start writing in batches
-writeBatch(workerData.start, workerData.end, workerData.batchSize)
+writeBatch(
+  workerData.start,
+  workerData.end,
+  workerData.batchSize,
+  workerData.collectionPath,
+  workerData.delayBetweenBatches // Pass the delay for rate control
+)
   .then(() => {
     parentPort.postMessage("Completed writing documents.");
   })