This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f0bae33ee1 Add missing location param to
`BigQueryUpdateTableSchemaOperator` (#40237)
f0bae33ee1 is described below
commit f0bae33ee1011e9d6120a40c7f44bf1ac40d46aa
Author: Shahar Epstein <[email protected]>
AuthorDate: Fri Jun 14 16:53:17 2024 +0300
Add missing location param to `BigQueryUpdateTableSchemaOperator` (#40237)
---
airflow/providers/google/cloud/operators/bigquery.py | 5 +++--
tests/providers/google/cloud/operators/test_bigquery.py | 7 +++++++
2 files changed, 10 insertions(+), 2 deletions(-)
diff --git a/airflow/providers/google/cloud/operators/bigquery.py
b/airflow/providers/google/cloud/operators/bigquery.py
index fbce7ca768..cb63955286 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -2719,6 +2719,7 @@ class
BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
project_id: str = PROVIDE_PROJECT_ID,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
+ location: str | None = None,
**kwargs,
) -> None:
self.schema_fields_updates = schema_fields_updates
@@ -2728,12 +2729,12 @@ class
BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
self.project_id = project_id
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
+ self.location = location
super().__init__(**kwargs)
def execute(self, context: Context):
bq_hook = BigQueryHook(
- gcp_conn_id=self.gcp_conn_id,
- impersonation_chain=self.impersonation_chain,
+ gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain, location=self.location
)
table = bq_hook.update_table_schema(
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py
b/tests/providers/google/cloud/operators/test_bigquery.py
index 6a92fe8089..6f4e76baa7 100644
--- a/tests/providers/google/cloud/operators/test_bigquery.py
+++ b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -458,9 +458,16 @@ class TestBigQueryUpdateTableSchemaOperator:
dataset_id=TEST_DATASET,
table_id=TEST_TABLE_ID,
project_id=TEST_GCP_PROJECT_ID,
+ location=TEST_DATASET_LOCATION,
+
impersonation_chain=["[email protected]"],
)
operator.execute(context=MagicMock())
+ mock_hook.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+
impersonation_chain=["[email protected]"],
+ location=TEST_DATASET_LOCATION,
+ )
mock_hook.return_value.update_table_schema.assert_called_once_with(
schema_fields_updates=schema_field_updates,
include_policy_tags=False,
