more spark roles

Author: max-ostapenko

infra/bigquery_export_spark/requirements.txt

@@ -1,4 +1,2 @@
-google-cloud-bigquery==3.23
-google-cloud-storage==2.16
 google-cloud-firestore==2.20.1
 # pyspark==3.5.5

infra/bigquery_export_spark/src/firestore.py

@@ -1,26 +1,36 @@
 """This module processes Firestore documents from BigQuery using Spark."""
 
 import json
+import os
 
-from google.cloud import bigquery, firestore  # type: ignore
+from google.cloud import firestore  # type: ignore
 from pyspark.sql import SparkSession  # type: ignore
 
 
+PROJECT = "httparchive"
+
+
+# pylint: disable=too-many-instance-attributes
 class FirestoreBatch:
     """Handles Firestore data batching from BigQuery using Spark."""
 
-    def __init__(self):
+    def __init__(self, export_config):
         """Initialize FirestoreBatch with default settings."""
-        self.firestore = firestore.Client()
-        self.bigquery = bigquery.Client()
+        self.config = {
+            "collection_name": export_config["name"],
+            "date": getattr(export_config, "date", ""),
+            "collection_type": export_config["type"],
+        }
+        self.firestore = firestore.Client(
+            project=PROJECT, database=export_config["database"]
+        )
         self.batch_size = 500
         self.max_concurrent_batches = 200
         self.current_batch = []
         self.batch_promises = []
         self.spark = SparkSession.builder.appName(
             "FirestoreBatchProcessor"
         ).getOrCreate()
-        self.config = {"date": "", "collection_name": "", "collection_type": ""}
 
     def queue_batch(self, operation):
         """Queue a batch commit operation for Firestore."""
@@ -36,7 +46,6 @@ def queue_batch(self, operation):
                 batch.set(doc_ref, doc)
             else:
                 raise ValueError("Invalid operation")
-
         self.batch_promises.append(batch.commit())
         self.current_batch = []
 
@@ -75,16 +84,17 @@ def batch_delete(self):
                 f"Deleting documents from {self.config['collection_name']} "
                 f"for date {self.config['date']}"
             )
-            query = collection_ref.where("date", "==", self.config["date"])
+            collection_query = collection_ref.where("date", "==", self.config["date"])
         elif self.config["collection_type"] == "dict":
             print(f"Deleting documents from {self.config['collection_name']}")
-            query = collection_ref
+            collection_query = collection_ref
         else:
             raise ValueError("Invalid collection type")
-
         while True:
             docs = list(
-                query.limit(self.batch_size * self.max_concurrent_batches).stream()
+                collection_query.limit(
+                    self.batch_size * self.max_concurrent_batches
+                ).stream()
             )
             if not docs:
                 break
@@ -105,13 +115,13 @@ def batch_delete(self):
             f"Time: {duration} seconds"
         )
 
-    def stream_from_bigquery(self, query):
+    def stream_from_bigquery(self, query_str):
         """Stream data from BigQuery to Firestore."""
         print("Starting BigQuery to Firestore transfer...")
         start_time = self.spark.sparkContext.startTime
         total_rows_processed = 0
 
-        df = self.spark.read.format("bigquery").option("query", query).load()
+        df = self.spark.read.format("bigquery").option("query", query_str).load()
 
         for row in df.collect():
             self.current_batch.append(row.asDict())
@@ -132,21 +142,19 @@ def stream_from_bigquery(self, query):
             f"seconds"
         )
 
-    def export(self):
+    def export(self, query_str):
         """Export data from BigQuery to Firestore."""
-        export_config = json.loads(
-            '{"name": "technologies", "type": "dict", "environment": "dev"}'
-        )
-        query = str(json.loads("SELECT * FROM report.tech_report_technologies"))
-
-        self.config["date"] = getattr(export_config, "date", "")
-        self.config["collection_name"] = export_config["name"]
-        self.config["collection_type"] = export_config["type"]
 
         self.batch_delete()
-        self.stream_from_bigquery(query)
+        self.stream_from_bigquery(query_str)
 
 
 if __name__ == "__main__":
-    processor = FirestoreBatch()
-    processor.export()
+    # config_data = json.loads('{"name": "technologies", "type": "dict", "environment": "dev"}')
+    # QUERY_STR = str(json.loads("SELECT * FROM report.tech_report_technologies"))
+
+    config_data = json.loads(os.environ["BIGQUERY_PROC_PARAM.export_config"])
+    QUERY_STR = str(json.loads(os.environ["BIGQUERY_PROC_PARAM.query"]))
+
+    processor = FirestoreBatch(config_data)
+    processor.export(QUERY_STR)

infra/tf/functions.tf

@@ -3,7 +3,7 @@ locals {
 }
 
 resource "google_project_iam_member" "project" {
-  for_each = toset(["roles/bigquery.jobUser", "roles/dataform.serviceAgent", "roles/run.invoker", "roles/run.jobsExecutorWithOverrides"])
+  for_each = toset(["roles/bigquery.user", "roles/dataform.serviceAgent", "roles/run.invoker", "roles/run.jobsExecutorWithOverrides", "roles/datastore.user", "roles/storage.objectAdmin"])
 
   project = local.project
   role    = each.value
@@ -37,7 +37,9 @@ resource "google_project_iam_member" "bigquery-remote-functions-connector" {
 }
 
 resource "google_project_iam_member" "spark-procedures-connector" {
+  for_each = toset(["roles/datastore.user", "roles/artifactregistry.reader", "roles/bigquery.user"])
+
   project = local.project
-  role    = "roles/artifactregistry.reader"
+  role    = each.value
   member  = "serviceAccount:${google_bigquery_connection.spark-procedures.spark[0].service_account_id}"
 }