Merge branch 'smart-otter' into smart-otter

Author: max-ostapenko

definitions/output/reports/tech_report_categories.js

@@ -73,9 +73,6 @@ technology_stats AS (
     category_obj AS categories,
     origins.mobile AS mobile_origins
   FROM ${ctx.ref('reports', 'tech_report_technologies')}
-  GROUP BY
-    technology,
-    categories
 )
 
 SELECT

infra/bigquery-export/firestore.js

@@ -3,13 +3,67 @@ import { BigQueryExport } from './bigquery.js'
 
 export class FirestoreBatch {
   constructor () {
-    this.firestore = new Firestore()
+    this.firestore = new Firestore({
+      gaxOptions: {
+        grpc: {
+          max_receive_message_length: 500 * 1024 * 1024, // 500MB
+          max_send_message_length: 500 * 1024 * 1024, // 500MB
+          '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
+
+    // Configuration constants
+    this.config = {
+      batchSize: {
+        delete: 500,
+        write: 400
+      },
+      maxConcurrentBatches: 200,
+      retryCount: 5,
+      timeout: 10 * 60 * 1000 // 10 minutes
+    }
+
+    this.reset()
   }
 
-  queueBatch (operation) {
+  reset () {
+    this.currentBatch = []
+    this.batchPromises = []
+  }
+
+  getCurrentBatchSize (operation) {
+    return this.config.batchSize[operation === 'delete' ? 'delete' : 'write']
+  }
+
+  async commitWithRetry (batch, index) {
+    let lastError
+
+    for (let attempt = 1; attempt <= this.config.retryCount; attempt++) {
+      try {
+        await batch.commit()
+        return
+      } catch (error) {
+        lastError = error
+        console.warn(`Batch ${index} attempt ${attempt} failed:`, error.message)
+
+        if (attempt < this.config.retryCount) {
+          const delayMs = Math.pow(2, attempt) * 500
+          console.log(`Retrying batch ${index} in ${delayMs}ms...`)
+          await new Promise(resolve => setTimeout(resolve, delayMs))
+        }
+      }
+    }
+
+    console.error(`Batch ${index} failed after ${this.config.retryCount} attempts:`, lastError)
+    throw lastError
+  }
+
+  createBatch (operation) {
     const batch = this.firestore.batch()
 
     this.currentBatch.forEach((doc) => {
@@ -19,119 +73,126 @@ export class FirestoreBatch {
         const docRef = this.firestore.collection(this.collectionName).doc()
         batch.set(docRef, doc)
       } else {
-        throw new Error('Invalid operation')
+        throw new Error(`Invalid operation: ${operation}`)
       }
     })
+
+    return batch
+  }
+
+  queueBatch (operation) {
+    const batch = this.createBatch(operation)
     this.batchPromises.push(batch)
     this.currentBatch = []
   }
 
   async commitBatches () {
+    if (this.batchPromises.length === 0) return
+
     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((batch, index) =>
+        this.commitWithRetry(batch, index)
       )
     )
+
     this.batchPromises = []
   }
 
-  async finalFlush (operation) {
-    if (this.currentBatch.length > 0) {
+  async processInBatches (operation, shouldFlush = false) {
+    const batchSize = this.getCurrentBatchSize(operation)
+
+    if (this.currentBatch.length >= batchSize || shouldFlush) {
       this.queueBatch(operation)
     }
 
-    if (this.batchPromises.length > 0) {
+    if (this.batchPromises.length >= this.config.maxConcurrentBatches || shouldFlush) {
       await this.commitBatches()
     }
   }
 
+  buildQuery (collectionRef) {
+    const queryMap = {
+      report: () => {
+        console.info(`Deleting documents from ${this.collectionName} for date ${this.date}`)
+        return collectionRef.where('date', '==', this.date)
+      },
+      dict: () => {
+        console.info(`Deleting documents from ${this.collectionName}`)
+        return collectionRef
+      }
+    }
+
+    const queryBuilder = queryMap[this.collectionType]
+    if (!queryBuilder) {
+      throw new Error(`Invalid collection type: ${this.collectionType}`)
+    }
+
+    return queryBuilder()
+  }
+
   async batchDelete () {
     console.info('Starting batch deletion...')
     const startTime = Date.now()
-    this.currentBatch = []
-    this.batchPromises = []
+    this.reset()
 
     let totalDocsDeleted = 0
     const collectionRef = this.firestore.collection(this.collectionName)
-
-    let collectionQuery
-    if (this.collectionType === 'report') {
-      console.info('Deleting documents from ' + this.collectionName + ' for date ' + this.date)
-      // Query to fetch monthly documents
-      collectionQuery = collectionRef.where('date', '==', this.date)
-    } else if (this.collectionType === 'dict') {
-      console.info('Deleting documents from ' + this.collectionName)
-      collectionQuery = collectionRef
-    } else {
-      throw new Error('Invalid collection type')
-    }
+    const collectionQuery = this.buildQuery(collectionRef)
+    const batchSize = this.getCurrentBatchSize('delete')
 
     while (true) {
-      const snapshot = await collectionQuery.limit(this.batchSize * this.maxConcurrentBatches).get()
-      if (snapshot.empty) {
-        break
-      }
+      const snapshot = await collectionQuery.limit(batchSize * this.config.maxConcurrentBatches).get()
+      if (snapshot.empty) break
 
-      for await (const doc of snapshot.docs) {
+      for (const doc of snapshot.docs) {
         this.currentBatch.push(doc)
-
-        if (this.currentBatch.length >= this.batchSize) {
-          this.queueBatch('delete')
-        }
-        if (this.batchPromises.length >= this.maxConcurrentBatches) {
-          await this.commitBatches()
-        }
+        await this.processInBatches('delete')
         totalDocsDeleted++
       }
     }
-    await this.finalFlush('delete')
+
+    // Final flush
+    await this.processInBatches('delete', true)
 
     const duration = (Date.now() - startTime) / 1000
     console.info(`Deletion complete. Total docs deleted: ${totalDocsDeleted}. Time: ${duration} seconds`)
   }
 
-  /**
-   * Streams BigQuery query results into a Firestore collection using batch commits.
-   * @param {string} query - The BigQuery SQL query.
-   */
   async streamFromBigQuery (rowStream) {
     console.info('Starting BigQuery to Firestore transfer...')
     const startTime = Date.now()
     let totalRowsProcessed = 0
 
-    this.currentBatch = []
-    this.batchPromises = []
+    this.reset()
 
     for await (const row of rowStream) {
       this.currentBatch.push(row)
-
-      // Write batch when it reaches specified size
-      if (this.currentBatch.length >= this.batchSize) {
-        this.queueBatch('set')
-      }
-
-      if (this.batchPromises.length >= this.maxConcurrentBatches) {
-        await this.commitBatches()
-      }
+      await this.processInBatches('set')
       totalRowsProcessed++
     }
-    await this.finalFlush('set')
+
+    // Final flush
+    await this.processInBatches('set', true)
 
     const duration = (Date.now() - startTime) / 1000
     console.info(`Transfer to ${this.collectionName} complete. Total rows processed: ${totalRowsProcessed}. Time: ${duration} seconds`)
   }
 
   async export (query, exportConfig) {
+    // Configure Firestore settings
     this.firestore.settings({
-      databaseId: exportConfig.database
+      databaseId: exportConfig.database,
+      timeout: this.config.timeout
+    })
+
+    // Set instance properties
+    Object.assign(this, {
+      collectionName: exportConfig.collection,
+      collectionType: exportConfig.type,
+      date: exportConfig.date
     })
-    this.collectionName = exportConfig.collection
-    this.collectionType = exportConfig.type
-    this.date = exportConfig.date
 
     await this.batchDelete()
 

infra/tf/.terraform.lock.hcl

@@ -35,7 +35,7 @@ resource "google_cloud_run_v2_job" "bigquery_export" {
           value = ""
         }
       }
-      timeout         = "3600s"
+      timeout         = "7200s"
       service_account = var.function_identity
       max_retries     = 1
     }

infra/tf/bigquery_export/main.tf

@@ -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
+}