This is an automated email from the ASF dual-hosted git repository.
shahar1 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 9a1273c52b8 Normalize BigQuery DTS sensor expected statuses after
rendering (#70528)
9a1273c52b8 is described below
commit 9a1273c52b8bcfad135037e7d2280c7d388f11a5
Author: Dr Alex Mitre <[email protected]>
AuthorDate: Tue Sep 22 09:28:53 2026 -0600
Normalize BigQuery DTS sensor expected statuses after rendering (#70528)
Co-authored-by: Shahar Epstein <[email protected]>
---
.../providers/google/cloud/sensors/bigquery_dts.py | 23 ++++++++++-----
.../unit/google/cloud/sensors/test_bigquery_dts.py | 34 ++++++++++++++++++++++
.../ci/prek/validate_operators_init_exemptions.txt | 1 -
3 files changed, 49 insertions(+), 9 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py
b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py
index 6f12f92f36d..05f79d38145 100644
---
a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py
+++
b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py
@@ -43,13 +43,14 @@ class
BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):
For more information on how to use this sensor, take a look at the
guide:
:ref:`howto/operator:BigQueryDataTransferServiceTransferRunSensor`
- :param expected_statuses: The expected state of the operation.
+ :param expected_statuses: The expected state of the operation. (templated)
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
- :param run_id: ID of the transfer run.
- :param transfer_config_id: ID of transfer config to be used.
+ :param run_id: ID of the transfer run. (templated)
+ :param transfer_config_id: ID of transfer config to be used. (templated)
:param project_id: The BigQuery project id where the transfer
configuration should be
created. If set to None or missing, the default project_id from the
Google Cloud connection is used.
+ (templated)
:param retry: A retry object used to retry requests. If `None` is
specified, requests will not be retried.
:param request_timeout: The amount of time, in seconds, to wait for the
request to
@@ -99,7 +100,7 @@ class
BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):
self.retry = retry
self.request_timeout = request_timeout
self.metadata = metadata
- self.expected_statuses = self._normalize_state_list(expected_statuses)
+ self.expected_statuses = expected_statuses
self.project_id = project_id
self.gcp_cloud_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
@@ -110,9 +111,15 @@ class
BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):
result = set()
for state in states:
if isinstance(state, str):
- # The proto.Enum type is indexable (via MetaClass and aliased)
but MyPy is not able to
- # infer this https://github.com/python/mypy/issues/8968
- result.add(TransferState[state.upper()]) # type: ignore[misc]
+ try:
+ # The proto.Enum type is indexable (via MetaClass and
aliased) but MyPy is not able to
+ # infer this https://github.com/python/mypy/issues/8968
+ result.add(TransferState[state.upper()]) # type:
ignore[misc]
+ except KeyError:
+ raise ValueError(
+ f"Invalid expected status {state!r}. "
+ f"Valid statuses: {sorted(TransferState.__members__)}"
+ ) from None
elif isinstance(state, int):
result.add(TransferState(state))
elif isinstance(state, TransferState):
@@ -144,4 +151,4 @@ class
BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):
if run.state in (TransferState.FAILED, TransferState.CANCELLED):
message = f"Transfer {self.run_id} did not succeed"
raise AirflowException(message)
- return run.state in self.expected_statuses
+ return run.state in self._normalize_state_list(self.expected_statuses)
diff --git
a/providers/google/tests/unit/google/cloud/sensors/test_bigquery_dts.py
b/providers/google/tests/unit/google/cloud/sensors/test_bigquery_dts.py
index a0f65811172..514432dd105 100644
--- a/providers/google/tests/unit/google/cloud/sensors/test_bigquery_dts.py
+++ b/providers/google/tests/unit/google/cloud/sensors/test_bigquery_dts.py
@@ -89,3 +89,37 @@ class TestBigQueryDataTransferServiceTransferRunSensor:
retry=DEFAULT,
timeout=None,
)
+
+ @mock.patch(
+
"airflow.providers.google.cloud.sensors.bigquery_dts.BiqQueryDataTransferServiceHook",
+
return_value=MM(get_transfer_run=MM(return_value=MM(state=TransferState.SUCCEEDED))),
+ )
+ def test_templated_expected_statuses_rendered_before_poke(self, mock_hook):
+ op = BigQueryDataTransferServiceTransferRunSensor(
+ transfer_config_id=TRANSFER_CONFIG_ID,
+ run_id=RUN_ID,
+ task_id="id",
+ project_id=PROJECT_ID,
+ expected_statuses="{{ expected }}",
+ )
+ assert op.expected_statuses == "{{ expected }}"
+
+ op.render_template_fields({"expected": "succeeded"})
+
+ assert op.poke({}) is True
+
+ @mock.patch(
+
"airflow.providers.google.cloud.sensors.bigquery_dts.BiqQueryDataTransferServiceHook",
+
return_value=MM(get_transfer_run=MM(return_value=MM(state=TransferState.SUCCEEDED))),
+ )
+ def test_poke_raises_value_error_for_invalid_expected_status(self,
mock_hook):
+ op = BigQueryDataTransferServiceTransferRunSensor(
+ transfer_config_id=TRANSFER_CONFIG_ID,
+ run_id=RUN_ID,
+ task_id="id",
+ project_id=PROJECT_ID,
+ expected_statuses="SUCCESS",
+ )
+
+ with pytest.raises(ValueError, match="SUCCESS.*Valid
statuses.*SUCCEEDED"):
+ op.poke({})
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index 45b4d646fc3..c49a448fcd3 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -11,5 +11,4 @@
providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py::Cl
providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py::CloudDataTransferServiceCreateJobOperator
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator
providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator
-providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor
providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor
