Add batch_id param to profile to pass along with `dataproc_v1.Creat…

…eBatchRequest` (#727) * Add new param to profile batch_id to pass along with google.cloud.dataproc.v1.CreateBatchRequest * add changie * Add to credential fixtures * implement suggestion on PR --------- Co-authored-by: colin-rogers-dbt <111200756+colin-rogers-dbt@users.noreply.github.com> Co-authored-by: Mike Alfare <13974384+mikealfare@users.noreply.github.com>
dbt-labs · Jun 26, 2023 · a42e13d · a42e13d
1 parent 0b08610
commit a42e13d
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 0 deletions.
diff --git a/.changes/unreleased/Features-20230517-092205.yaml b/.changes/unreleased/Features-20230517-092205.yaml
@@ -0,0 +1,6 @@
+kind: Features
+body: Add batch_id param to profile
+time: 2023-05-17T09:22:05.264368+01:00
+custom:
+  Author: nickozilla
+  Issue: "671"
diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py
@@ -133,6 +133,7 @@ class BigQueryCredentials(Credentials):
     dataproc_region: Optional[str] = None
     dataproc_cluster_name: Optional[str] = None
     gcs_bucket: Optional[str] = None
+    batch_id: Optional[str] = None
 
     dataproc_batch: Optional[DataprocBatchConfig] = field(
         metadata={

diff --git a/dbt/adapters/bigquery/python_submissions.py b/dbt/adapters/bigquery/python_submissions.py
@@ -122,6 +122,7 @@ def _submit_dataproc_job(self) -> dataproc_v1.types.jobs.Job:
         request = dataproc_v1.CreateBatchRequest(
             parent=parent,
             batch=batch,
+            batch_id=self.credential.batch_id,
         )
         # make the request
         operation = self.job_client.create_batch(request=request)  # type: ignore

diff --git a/test.env.example b/test.env.example
@@ -15,3 +15,4 @@ DBT_TEST_USER_3="serviceAccount:dbt-integration-test-user@dbt-test-env.iam.gserv
 DATAPROC_REGION=us-
 DATAPROC_CLUSTER_NAME=
 GCS_BUCKET=
+BATCH_ID=
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -33,6 +33,7 @@ def oauth_target():
         "dataproc_region": os.getenv("DATAPROC_REGION"),
         "dataproc_cluster_name": os.getenv("DATAPROC_CLUSTER_NAME"),
         "gcs_bucket": os.getenv("GCS_BUCKET"),
+        "batch_id": os.getenv("BATCH_ID"),
     }
 
 
@@ -53,4 +54,5 @@ def service_account_target():
             "DATAPROC_CLUSTER_NAME"
         ),  # only needed for cluster submission method
         "gcs_bucket": os.getenv("GCS_BUCKET"),
+        "batch_id": os.getenv("BATCH_ID"),
     }