Fix to read location parameter properly in BigQueryToBigQueryOperator (#27661)

VladaZakharova · web-flow · commit bc0063af9962 · 2022-11-16T23:36:04.000+01:00
diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py b/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py
@@ -58,7 +58,11 @@ class BigQueryToBigQueryOperator(BaseOperator):
             encryption_configuration = {
                 "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key"
             }
-    :param location: The location used for the operation.
+    :param location: The geographic location of the job. You must specify the location to run the job if
+        the location to run a job is not in the US or the EU multi-regional location or
+        the location is in a single region (for example, us-central1).
+        For more details check:
+        https://cloud.google.com/bigquery/docs/locations#specifying_your_location
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
         of the last account in the list, which will be impersonated in the request.
@@ -131,7 +135,7 @@ def execute(self, context: Context) -> None:
                 encryption_configuration=self.encryption_configuration,
             )
 
-            job = hook.get_job(job_id=job_id).to_api_repr()
+            job = hook.get_job(job_id=job_id, location=self.location).to_api_repr()
             conf = job["configuration"]["copy"]["destinationTable"]
             BigQueryTableLink.persist(
                 context=context,
diff --git a/tests/providers/google/cloud/transfers/test_bigquery_to_bigquery.py b/tests/providers/google/cloud/transfers/test_bigquery_to_bigquery.py
@@ -22,14 +22,15 @@
 
 from airflow.providers.google.cloud.transfers.bigquery_to_bigquery import BigQueryToBigQueryOperator
 
+BQ_HOOK_PATH = "airflow.providers.google.cloud.transfers.bigquery_to_bigquery.BigQueryHook"
 TASK_ID = "test-bq-create-table-operator"
 TEST_DATASET = "test-dataset"
 TEST_TABLE_ID = "test-table-id"
 
 
 class TestBigQueryToBigQueryOperator(unittest.TestCase):
-    @mock.patch("airflow.providers.google.cloud.transfers.bigquery_to_bigquery.BigQueryHook")
-    def test_execute(self, mock_hook):
+    @mock.patch(BQ_HOOK_PATH)
+    def test_execute_without_location_should_execute_successfully(self, mock_hook):
         source_project_dataset_tables = f"{TEST_DATASET}.{TEST_TABLE_ID}"
         destination_project_dataset_table = f"{TEST_DATASET + '_new'}.{TEST_TABLE_ID}"
         write_disposition = "WRITE_EMPTY"
@@ -56,3 +57,31 @@ def test_execute(self, mock_hook):
             labels=labels,
             encryption_configuration=encryption_configuration,
         )
+
+    @mock.patch(BQ_HOOK_PATH)
+    def test_execute_single_regional_location_should_execute_successfully(self, mock_hook):
+        source_project_dataset_tables = f"{TEST_DATASET}.{TEST_TABLE_ID}"
+        destination_project_dataset_table = f"{TEST_DATASET + '_new'}.{TEST_TABLE_ID}"
+        write_disposition = "WRITE_EMPTY"
+        create_disposition = "CREATE_IF_NEEDED"
+        labels = {"k1": "v1"}
+        location = "us-central1"
+        encryption_configuration = {"key": "kk"}
+        mock_hook.return_value.run_copy.return_value = "job-id"
+
+        operator = BigQueryToBigQueryOperator(
+            task_id=TASK_ID,
+            source_project_dataset_tables=source_project_dataset_tables,
+            destination_project_dataset_table=destination_project_dataset_table,
+            write_disposition=write_disposition,
+            create_disposition=create_disposition,
+            labels=labels,
+            encryption_configuration=encryption_configuration,
+            location=location,
+        )
+
+        operator.execute(context=mock.MagicMock())
+        mock_hook.return_value.get_job.assert_called_once_with(
+            job_id="job-id",
+            location=location,
+        )
diff --git a/tests/system/providers/google/cloud/bigquery/example_bigquery_to_bigquery.py b/tests/system/providers/google/cloud/bigquery/example_bigquery_to_bigquery.py
@@ -38,6 +38,7 @@
 DATASET_NAME = f"dataset_{DAG_ID}_{ENV_ID}"
 ORIGIN = "origin"
 TARGET = "target"
+LOCATION = "US"
 
 
 with models.DAG(
@@ -47,7 +48,11 @@
     catchup=False,
     tags=["example", "bigquery"],
 ) as dag:
-    create_dataset = BigQueryCreateEmptyDatasetOperator(task_id="create_dataset", dataset_id=DATASET_NAME)
+    create_dataset = BigQueryCreateEmptyDatasetOperator(
+        task_id="create_dataset",
+        dataset_id=DATASET_NAME,
+        location=LOCATION,
+    )
 
     create_origin_table = BigQueryCreateEmptyTableOperator(
         task_id="create_origin_table",
