GCSToBigQueryOperator Resolve max_id_key job retrieval and xcom ret…

…urn (#26285) * Resolve `max_id_key` job retrieval and xcom return * Fixing Flake8 Issue * Fixing Unit Test failure * Code fixes requested during review
apache · Sep 18, 2022 · 07fe356 · 07fe356
1 parent b4f8a06
commit 07fe356
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
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,
  )