Merge branch 'main' into unusual-mosquito

Author: max-ostapenko

infra/bigquery-export/index.js

@@ -1,27 +1,15 @@
-import { StorageUpload } from './storage.js'
 import { FirestoreBatch } from './firestore.js'
 
 async function main () {
   const { query, destination, config } = process.env.EXPORT_CONFIG && JSON.parse(process.env.EXPORT_CONFIG)
-  if (!destination) {
-    throw new Error('No destination found')
-  }
-
-  if (destination === 'cloud_storage') {
-    console.info('Cloud Storage export')
-    console.log(query, config)
 
-    const storage = new StorageUpload(config.bucket)
-    await storage.exportToJson(query, config.name)
-  } else if (destination === 'firestore') {
-    console.info('Firestore export')
+  if (destination === 'firestore') {
     console.log(query, config)
 
     const firestore = new FirestoreBatch()
     await firestore.export(query, config)
-  } else {
-    throw new Error('Bad Request: destination unknown')
   }
+
   console.info('Export finished successfully')
   return 'OK'
 }

infra/bigquery-export/package-lock.json

@@ -9,8 +9,7 @@
   "type": "module",
   "dependencies": {
     "@google-cloud/bigquery": "8.1.1",
-    "@google-cloud/firestore": "7.11.3",
-    "@google-cloud/storage": "7.16.0"
+    "@google-cloud/firestore": "7.11.3"
   },
   "author": "@max-ostapenko"
 }

infra/bigquery-export/package.json

@@ -0,0 +1,34 @@
+import { BigQuery } from '@google-cloud/bigquery'
+
+export class BigQueryExport {
+  constructor (options = {}) {
+    this.bigquery = new BigQuery(options)
+  }
+
+  async queryResults (query) {
+    const options = {
+      query,
+      projectId: this.projectId,
+      location: this.location
+    }
+
+    const [job] = await this.bigquery.createQueryJob(options)
+    console.info(`Running BigQuery query: ${job.id}`)
+    const [rows] = await job.getQueryResults()
+    console.log('Fetching query results completed')
+    return rows
+  }
+
+  async queryResultsStream (query) {
+    const options = {
+      query,
+      projectId: this.projectId,
+      location: this.location
+    }
+
+    const [job] = await this.bigquery.createQueryJob(options)
+    console.info(`Running BigQuery query: ${job.id}`)
+    const rows = job.getQueryResultsStream()
+    return rows
+  }
+}

infra/dataform-service/bigquery.js

@@ -1,13 +1,19 @@
 import functions from '@google-cloud/functions-framework'
-import { BigQuery } from '@google-cloud/bigquery'
 
+import { BigQueryExport } from './bigquery.js'
 import { callRunJob } from './cloud_run.js'
 import { getCompilationResults, runWorkflow } from './dataform.js'
+import { StorageUpload } from './storage.js'
 
 const projectId = 'httparchive'
 const location = 'us-central1'
 const jobId = 'bigquery-export'
 
+const bigquery = new BigQueryExport({
+  projectId,
+  location: 'US'
+})
+
 const TRIGGERS = {
   crux_ready: {
     type: 'poller',
@@ -83,8 +89,22 @@ async function handleExport (req, res) {
       return
     }
 
-    const jobName = `projects/${projectId}/locations/${location}/jobs/${jobId}`
-    await callRunJob(jobName, payload)
+    const { query, destination, config } = payload
+
+    if (destination === 'cloud_storage') {
+      console.info('Cloud Storage export')
+      console.log(query, config)
+
+      const data = await bigquery.queryResults(query)
+      const storage = new StorageUpload(config.bucket)
+      await storage.exportToJson(data, config.name)
+    } else if (destination === 'firestore') {
+      console.info('Firestore export')
+      const jobName = `projects/${projectId}/locations/${location}/jobs/${jobId}`
+      await callRunJob(jobName, payload)
+    } else {
+      throw new Error('Bad Request: destination unknown')
+    }
 
     res.status(200).json({
       replies: [200],
@@ -136,7 +156,9 @@ async function handleTrigger (req, res) {
       const trigger = TRIGGERS[eventName]
       if (trigger.type === 'poller') {
         console.info(`Poller action ${eventName}`)
-        const result = await runQuery(trigger.query)
+
+        const rows = await bigquery.queryResults(trigger.query)
+        const result = rows.length > 0 && rows[0][Object.keys(rows[0])[0]] === true
         console.info(`Query result: ${result}`)
         if (result) {
           await executeAction(trigger.action, trigger.actionArgs)
@@ -160,22 +182,6 @@ async function handleTrigger (req, res) {
   }
 }
 
-/**
- * Run BigQuery poll query.
- *
- * @param {string} query Polling query.
- * @returns {boolean} Query result.
- */
-async function runQuery (query) {
-  const bigquery = new BigQuery()
-
-  const [job] = await bigquery.createQueryJob({ query })
-  console.info(`Query job ${job.id} started.`)
-
-  const [rows] = await job.getQueryResults()
-  return rows.length > 0 && rows[0][Object.keys(rows[0])[0]] === true
-}
-
 /**
  * Execute action based on the trigger configuration.
  *