GCSToBigQueryOperator Resolve max_id_key job retrieval and xcom return (#26285)

patricker · web-flow · commit 07fe356de074 · 2022-09-18T22:12:10.000+02:00
* Resolve `max_id_key` job retrieval and xcom return

* Fixing Flake8 Issue

* Fixing Unit Test failure

* Code fixes requested during review
diff --git a/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py b/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py
@@ -316,9 +316,10 @@ def execute(self, context: Context):
                 warnings.simplefilter("ignore", DeprecationWarning)
                 job_id = bq_hook.run_query(
                     sql=select_command,
+                    location=self.location,
                     use_legacy_sql=False,
                 )
-            row = list(bq_hook.get_job(job_id).result())
+            row = list(bq_hook.get_job(job_id=job_id, location=self.location).result())
             if row:
                 max_id = row[0] if row[0] else 0
                 self.log.info(
@@ -327,5 +328,6 @@ def execute(self, context: Context):
                     self.max_id_key,
                     max_id,
                 )
+                return max_id
             else:
                 raise RuntimeError(f"The {select_command} returned no rows!")
diff --git a/tests/providers/google/cloud/transfers/test_gcs_to_bigquery.py b/tests/providers/google/cloud/transfers/test_gcs_to_bigquery.py
@@ -45,10 +45,13 @@ def test_execute_explicit_project(self, bq_hook):
 
         bq_hook.return_value.get_job.return_value.result.return_value = ('1',)
 
-        operator.execute(None)
+        result = operator.execute(None)
+
+        assert result == '1'
 
         bq_hook.return_value.run_query.assert_called_once_with(
             sql="SELECT MAX(id) FROM `test-project.dataset.table`",
+            location=None,
             use_legacy_sql=False,
         )
 
