firestore timeout

Author: max-ostapenko

infra/bigquery-export/firestore.js

@@ -3,10 +3,23 @@ import { BigQueryExport } from './bigquery.js'
 
 export class FirestoreBatch {
   constructor () {
-    this.firestore = new Firestore()
+    this.firestore = new Firestore({
+      // Increase timeout to 10 minutes for large batch operations
+      gaxOptions: {
+        grpc: {
+          max_receive_message_length: 100 * 1024 * 1024, // 100MB
+          max_send_message_length: 100 * 1024 * 1024, // 100MB
+          'grpc.max_connection_idle_ms': 5 * 60 * 1000, // 5 minutes
+          'grpc.keepalive_time_ms': 30 * 1000, // 30 seconds
+          'grpc.keepalive_timeout_ms': 60 * 1000, // 1 minute
+          'grpc.keepalive_permit_without_calls': true
+        }
+      }
+    })
     this.bigquery = new BigQueryExport()
-    this.batchSize = 500
-    this.maxConcurrentBatches = 200
+    this.batchSizeDelete = 500
+    this.batchSizeWrite = 400 // Reduced batch size for better performance
+    this.maxConcurrentBatches = 100 // Reduced concurrent batches to avoid overwhelming
   }
 
   queueBatch (operation) {
@@ -28,14 +41,34 @@ export class FirestoreBatch {
 
   async commitBatches () {
     console.log(`Committing ${this.batchPromises.length} batches to ${this.collectionName}`)
+
     await Promise.all(
-      this.batchPromises.map(async (batchPromise) => await batchPromise.commit()
-        .catch((error) => {
-          console.error('Error committing batch:', error)
-          throw error
-        })
-      )
+      this.batchPromises.map(async (batchPromise, index) => {
+        const retryCount = 3
+        let lastError
+
+        for (let attempt = 1; attempt <= retryCount; attempt++) {
+          try {
+            await batchPromise.commit()
+            return
+          } catch (error) {
+            lastError = error
+            console.warn(`Batch ${index} attempt ${attempt} failed:`, error.message)
+
+            if (attempt < retryCount) {
+              // Exponential backoff: 2^attempt seconds
+              const delayMs = Math.pow(2, attempt) * 1000
+              console.log(`Retrying batch ${index} in ${delayMs}ms...`)
+              await new Promise(resolve => setTimeout(resolve, delayMs))
+            }
+          }
+        }
+
+        console.error(`Batch ${index} failed after ${retryCount} attempts:`, lastError)
+        throw lastError
+      })
     )
+
     this.batchPromises = []
   }
 
@@ -71,7 +104,7 @@ export class FirestoreBatch {
     }
 
     while (true) {
-      const snapshot = await collectionQuery.limit(this.batchSize * this.maxConcurrentBatches).get()
+      const snapshot = await collectionQuery.limit(this.batchSizeDelete * this.maxConcurrentBatches).get()
       if (snapshot.empty) {
         break
       }
@@ -127,7 +160,8 @@ export class FirestoreBatch {
 
   async export (query, exportConfig) {
     this.firestore.settings({
-      databaseId: exportConfig.database
+      databaseId: exportConfig.database,
+      timeout: 10 * 60 * 1000 // 10 minutes timeout
     })
     this.collectionName = exportConfig.collection
     this.collectionType = exportConfig.type

infra/tf/.terraform.lock.hcl

@@ -9,7 +9,7 @@ terraform {
   }
 }
 
-resource "google_project_iam_member" "project" {
+resource "google_project_iam_member" "function_identity" {
   for_each = toset(["roles/bigquery.jobUser", "roles/dataform.serviceAgent", "roles/run.invoker", "roles/run.jobsExecutorWithOverrides", "roles/datastore.user", "roles/storage.objectUser"])
 
   project = var.project
@@ -31,8 +31,10 @@ resource "google_bigquery_connection" "remote-functions" {
   cloud_resource {}
 }
 
-resource "google_project_iam_member" "bigquery-remote-functions-connector" {
+resource "google_project_iam_member" "bigquery-connection-remote-functions" {
+  for_each = toset(["roles/run.invoker"])
+
   project = var.project
-  role    = "roles/run.invoker"
+  role    = each.value
   member  = "serviceAccount:${google_bigquery_connection.remote-functions.cloud_resource[0].service_account_id}"
 }

infra/tf/functions/main.tf

@@ -0,0 +1,9 @@
+output "google_bigquery_connection-remote_functions-id" {
+  description = "The connection ID for the remote functions BigQuery connection."
+  value       = google_bigquery_connection.remote-functions.id
+}
+
+output "remote_functions_connection_service_account_id" {
+  description = "The service account ID associated with the remote functions BigQuery connection."
+  value       = google_bigquery_connection.remote-functions.cloud_resource[0].service_account_id
+}

infra/tf/functions/output.tf

@@ -28,16 +28,6 @@ locals {
   function_identity = "cloud-function@httparchive.iam.gserviceaccount.com"
 }
 
-module "dataform_export" {
-  source = "./dataform_export"
-
-  project_number              = local.project_number
-  region                      = local.region
-  function_identity           = local.function_identity
-  function_name               = "dataform-export"
-  remote_functions_connection = google_bigquery_connection.remote-functions.id
-}
-
 module "dataform_trigger" {
   source = "./dataform_trigger"
 
@@ -74,3 +64,14 @@ module "functions" {
   function_identity = local.function_identity
   edit_datasets     = local.edit_datasets
 }
+
+module "dataform_export" {
+  source = "./dataform_export"
+
+  project_number              = local.project_number
+  region                      = local.region
+  function_identity           = local.function_identity
+  function_name               = "dataform-export"
+  remote_functions_connection = module.functions.google_bigquery_connection-remote_functions-id
+  depends_on                  = [module.functions]
+}