Skip to content

Commit

Permalink
enable overriding the dataproc project for python models
Browse files Browse the repository at this point in the history
  • Loading branch information
matt-winkler committed Oct 8, 2024
1 parent 455c768 commit 1f20dbe
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 1 deletion.
5 changes: 5 additions & 0 deletions dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ class BigQueryCredentials(Credentials):

dataproc_region: Optional[str] = None
dataproc_cluster_name: Optional[str] = None
dataproc_project: Optional[str] = None
gcs_bucket: Optional[str] = None

dataproc_batch: Optional[DataprocBatchConfig] = field(
Expand Down Expand Up @@ -228,6 +229,10 @@ def __pre_deserialize__(cls, d: Dict[Any, Any]) -> Dict[Any, Any]:
# `execution_project` default to dataset/project
if "execution_project" not in d:
d["execution_project"] = d["database"]

if "dataproc_project" not in d:
d["dataproc_project"] = d["execution_project"]

return d


Expand Down
3 changes: 2 additions & 1 deletion dbt/adapters/bigquery/python_submissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ def __init__(self, parsed_model: Dict, credential: BigQueryCredentials) -> None:
self.parsed_model = parsed_model
python_required_configs = [
"dataproc_region",
"dataproc_project",
"gcs_bucket",
]
for required_config in python_required_configs:
Expand Down Expand Up @@ -107,7 +108,7 @@ def _submit_dataproc_job(self) -> dataproc_v1.types.jobs.Job:
}
operation = self.job_client.submit_job_as_operation(
request={
"project_id": self.credential.execution_project,
"project_id": self.credential.dataproc_project,
"region": self.credential.dataproc_region,
"job": job,
}
Expand Down

0 comments on commit 1f20dbe

Please sign in to comment.