Add location support to BigQueryDataTransferServiceTransferRunSensor.

apache · Apr 25, 2022 · 967140e · 967140e
1 parent 692a089
commit 967140e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/sensors/bigquery_dts.py b/airflow/providers/google/cloud/sensors/bigquery_dts.py
@@ -84,6 +84,7 @@ def __init__(
  retry: Union[Retry, _MethodDefault] = DEFAULT,
  request_timeout: Optional[float] = None,
  metadata: Sequence[Tuple[str, str]] = (),
+ location: Optional[str] = None,
  impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
  **kwargs,
  ) -> None:
@@ -97,6 +98,7 @@ def __init__(
  self.project_id = project_id
  self.gcp_cloud_conn_id = gcp_conn_id
  self.impersonation_chain = impersonation_chain
+ self.location = location
 
  def _normalize_state_list(self, states) -> Set[TransferState]:
  states = {states} if isinstance(states, (str, TransferState, int)) else states
@@ -122,6 +124,7 @@ def poke(self, context: 'Context') -> bool:
  hook = BiqQueryDataTransferServiceHook(
  gcp_conn_id=self.gcp_cloud_conn_id,
  impersonation_chain=self.impersonation_chain,
+ location=self.location,
  )
  run = hook.get_transfer_run(
  run_id=self.run_id,

diff --git a/tests/providers/google/cloud/sensors/test_bigquery_dts.py b/tests/providers/google/cloud/sensors/test_bigquery_dts.py
@@ -30,6 +30,8 @@
 TRANSFER_CONFIG_ID = "config_id"
 RUN_ID = "run_id"
 PROJECT_ID = "project_id"
+LOCATION = "europe"
+GCP_CONN_ID = "google_cloud_default"
 
 
 class TestBigQueryDataTransferServiceTransferRunSensor(unittest.TestCase):
@@ -48,6 +50,8 @@ def test_poke_returns_false(self, mock_hook):
 
  with pytest.raises(AirflowException, match="Transfer"):
  op.poke({})
+
+ mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=None, location=None)
  mock_hook.return_value.get_transfer_run.assert_called_once_with(
  transfer_config_id=TRANSFER_CONFIG_ID,
  run_id=RUN_ID,
@@ -68,10 +72,15 @@ def test_poke_returns_true(self, mock_hook):
  task_id="id",
  project_id=PROJECT_ID,
  expected_statuses={"SUCCEEDED"},
+ location=LOCATION,
  )
  result = op.poke({})
 
  assert result is True
+
+ mock_hook.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID, impersonation_chain=None, location=LOCATION
+ )
  mock_hook.return_value.get_transfer_run.assert_called_once_with(
  transfer_config_id=TRANSFER_CONFIG_ID,
  run_id=RUN_ID,