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 9287e0c1d47 Move CloudComposerExternalTaskSensor task-id wrapping out
of __init__ (#73103)
9287e0c1d47 is described below
commit 9287e0c1d477a290bbd1afb9a4d7c45ae9fd7fe2
Author: Kailash Nelson <[email protected]>
AuthorDate: Tue Sep 22 23:56:45 2026 -0600
Move CloudComposerExternalTaskSensor task-id wrapping out of __init__
(#73103)
Template fields are rendered after the constructor, so wrapping a task id
into a list or treating an empty list as unset inspects the un-rendered
Jinja expression and keeps the class on the exemption list.
Reuse self._total_states instead of recalculating, and call
_normalize_external_task_ids() at the start of execute_complete() so
deferrable resume on a fresh instance raises ExternalTaskFailedError.
---
.../google/cloud/sensors/cloud_composer.py | 89 +++++----
.../google/cloud/sensors/test_cloud_composer.py | 202 +++++++++++++++++++++
.../ci/prek/validate_operators_init_exemptions.txt | 1 -
3 files changed, 254 insertions(+), 38 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py
b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py
index 7698f3af4f2..8a307e7be75 100644
---
a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py
+++
b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py
@@ -398,18 +398,56 @@ class CloudComposerExternalTaskSensor(BaseSensorOperator):
self.skipped_states = list(skipped_states) if skipped_states else []
self.failed_states = list(failed_states) if failed_states else []
- total_states = set(self.allowed_states + self.skipped_states +
self.failed_states)
+ self._total_states = set(self.allowed_states + self.skipped_states +
self.failed_states)
- if len(total_states) != len(self.allowed_states) +
len(self.skipped_states) + len(self.failed_states):
+ if len(self._total_states) != len(self.allowed_states) +
len(self.skipped_states) + len(
+ self.failed_states
+ ):
raise DuplicateStateError(
"Duplicate values provided across allowed_states,
skipped_states and failed_states."
)
- # convert [] to None
+ self.execution_range = execution_range
+ self.composer_external_dag_id = composer_external_dag_id
+ self.composer_external_task_id = composer_external_task_id
+ self.composer_external_task_ids = composer_external_task_ids
+ self.composer_external_task_group_id = composer_external_task_group_id
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+ self.deferrable = deferrable
+ self.poll_interval = poll_interval
+ self._external_task_ids_normalized = False
+
+ def _get_logical_dates(self, context) -> tuple[datetime, datetime]:
+ logical_date = context.get("logical_date", None)
+ if logical_date is None:
+ raise RuntimeError(
+ "logical_date is None. Please make sure the sensor is not used
in an asset-triggered Dag. "
+ "CloudComposerDAGRunSensor was designed to be used in
time-based scheduled Dags only, "
+ "and asset-triggered Dags do not have logical_date. "
+ )
+ if isinstance(self.execution_range, timedelta):
+ if self.execution_range < timedelta(0):
+ return logical_date, logical_date - self.execution_range
+ return logical_date - self.execution_range, logical_date
+ if isinstance(self.execution_range, list) and
len(self.execution_range) > 0:
+ return self.execution_range[0], self.execution_range[1] if len(
+ self.execution_range
+ ) > 1 else logical_date
+ return logical_date - timedelta(1), logical_date
+
+ def _normalize_external_task_ids(self) -> None:
+ """Wrap a single task id and validate the target type after templates
have been rendered."""
+ if self._external_task_ids_normalized:
+ return
+
+ composer_external_task_id = self.composer_external_task_id
+ composer_external_task_ids = self.composer_external_task_ids
+ composer_external_task_group_id = self.composer_external_task_group_id
+
if not composer_external_task_ids:
composer_external_task_ids = None
- # can't set both single task id and a list of task ids
if composer_external_task_id is not None and
composer_external_task_ids is not None:
raise ValueError(
"Only one of `composer_external_task_id` or
`composer_external_task_ids` may "
@@ -417,7 +455,6 @@ class CloudComposerExternalTaskSensor(BaseSensorOperator):
"use `composer_external_task_id` or
`composer_external_task_ids` or `composer_external_task_group_id`."
)
- # since both not set, convert the single id to a 1-elt list - from
here on, we only consider the list
if composer_external_task_id is not None:
composer_external_task_ids = [composer_external_task_id]
@@ -428,49 +465,25 @@ class CloudComposerExternalTaskSensor(BaseSensorOperator):
"use `composer_external_task_id` or
`composer_external_task_ids` or `composer_external_task_group_id`."
)
- # check the requested states are all valid states for the target type,
be it dag or task
if composer_external_task_ids or composer_external_task_group_id:
- if not total_states <= set(State.task_states):
+ if not self._total_states <= set(State.task_states):
raise ValueError(
"Valid values for `allowed_states`, `skipped_states` and
`failed_states` "
- "when `composer_external_task_id` or
`composer_external_task_ids` or `composer_external_task_group_id` "
- f"is not `None`: {State.task_states}"
+ "when `composer_external_task_id` or
`composer_external_task_ids` or "
+ f"`composer_external_task_group_id` is not `None`:
{State.task_states}"
)
- elif not total_states <= set(State.dag_states):
+ elif not self._total_states <= set(State.dag_states):
raise ValueError(
"Valid values for `allowed_states`, `skipped_states` and
`failed_states` "
- f"when `composer_external_task_id` and
`composer_external_task_group_id` is `None`: {State.dag_states}"
+ f"when `composer_external_task_id` and
`composer_external_task_group_id` is `None`: "
+ f"{State.dag_states}"
)
- self.execution_range = execution_range
- self.composer_external_dag_id = composer_external_dag_id
- self.composer_external_task_id = composer_external_task_id
self.composer_external_task_ids = composer_external_task_ids
- self.composer_external_task_group_id = composer_external_task_group_id
- self.gcp_conn_id = gcp_conn_id
- self.impersonation_chain = impersonation_chain
- self.deferrable = deferrable
- self.poll_interval = poll_interval
-
- def _get_logical_dates(self, context) -> tuple[datetime, datetime]:
- logical_date = context.get("logical_date", None)
- if logical_date is None:
- raise RuntimeError(
- "logical_date is None. Please make sure the sensor is not used
in an asset-triggered Dag. "
- "CloudComposerDAGRunSensor was designed to be used in
time-based scheduled Dags only, "
- "and asset-triggered Dags do not have logical_date. "
- )
- if isinstance(self.execution_range, timedelta):
- if self.execution_range < timedelta(0):
- return logical_date, logical_date - self.execution_range
- return logical_date - self.execution_range, logical_date
- if isinstance(self.execution_range, list) and
len(self.execution_range) > 0:
- return self.execution_range[0], self.execution_range[1] if len(
- self.execution_range
- ) > 1 else logical_date
- return logical_date - timedelta(1), logical_date
+ self._external_task_ids_normalized = True
def poke(self, context: Context) -> bool:
+ self._normalize_external_task_ids()
start_date, end_date = self._get_logical_dates(context)
task_instances = self._get_task_instances(
@@ -634,6 +647,7 @@ class CloudComposerExternalTaskSensor(BaseSensorOperator):
)
def execute(self, context: Context) -> None:
+ self._normalize_external_task_ids()
self._composer_airflow_version = self._get_composer_airflow_version()
if self.composer_external_task_ids and
len(self.composer_external_task_ids) > len(
@@ -694,6 +708,7 @@ class CloudComposerExternalTaskSensor(BaseSensorOperator):
super().execute(context)
def execute_complete(self, context: Context, event: dict):
+ self._normalize_external_task_ids()
if event and event["status"] == "error":
raise AirflowException(event["message"])
if event and event["status"] == "failed":
diff --git
a/providers/google/tests/unit/google/cloud/sensors/test_cloud_composer.py
b/providers/google/tests/unit/google/cloud/sensors/test_cloud_composer.py
index 21a545a9dc3..f97b46151b2 100644
--- a/providers/google/tests/unit/google/cloud/sensors/test_cloud_composer.py
+++ b/providers/google/tests/unit/google/cloud/sensors/test_cloud_composer.py
@@ -28,6 +28,7 @@ from airflow.providers.google.cloud.sensors.cloud_composer
import (
CloudComposerDAGRunSensor,
CloudComposerExternalTaskSensor,
)
+from airflow.providers.standard.exceptions import ExternalTaskFailedError
TEST_PROJECT_ID = "test_project_id"
TEST_OPERATION_NAME = "test_operation_name"
@@ -350,3 +351,204 @@ class TestCloudComposerExternalTaskSensor:
task._composer_airflow_version = composer_airflow_version
assert not task.poke(context={"logical_date": datetime(2024, 5, 23, 0,
0, 0)})
+
+ def test_init_does_not_wrap_task_id_into_task_ids(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id=TEST_COMPOSER_EXTERNAL_TASK_ID,
+ )
+
+ assert task.composer_external_task_id == TEST_COMPOSER_EXTERNAL_TASK_ID
+ assert task.composer_external_task_ids is None
+
+
@mock.patch("airflow.providers.google.cloud.sensors.cloud_composer.CloudComposerHook")
+ def test_poke_wraps_task_id_into_task_ids(self, mock_hook):
+ mock_hook.return_value.get_task_instances.return_value =
TEST_GET_TASK_INSTANCES_RESULT(
+ "success",
+ "logical_date",
+ TEST_COMPOSER_EXTERNAL_TASK_ID,
+ )
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id=TEST_COMPOSER_EXTERNAL_TASK_ID,
+ allowed_states=["success"],
+ )
+ task._composer_airflow_version = 3
+
+ context = {"logical_date": datetime(2024, 5, 23, 0, 0, 0)}
+ assert task.poke(context=context)
+ assert task.composer_external_task_ids ==
[TEST_COMPOSER_EXTERNAL_TASK_ID]
+ assert task.poke(context=context)
+ assert task.composer_external_task_ids ==
[TEST_COMPOSER_EXTERNAL_TASK_ID]
+
+ def test_init_keeps_empty_task_ids_until_poke(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_ids=[],
+ )
+
+ assert task.composer_external_task_ids == []
+
+
@mock.patch("airflow.providers.google.cloud.sensors.cloud_composer.CloudComposerHook")
+ def test_poke_treats_empty_task_ids_as_unset(self, mock_hook):
+ mock_hook.return_value.get_task_instances.return_value =
TEST_GET_TASK_INSTANCES_RESULT(
+ "success",
+ "logical_date",
+ TEST_COMPOSER_EXTERNAL_TASK_ID,
+ )
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_ids=[],
+ allowed_states=["success"],
+ )
+ task._composer_airflow_version = 3
+
+ task.poke(context={"logical_date": datetime(2024, 5, 23, 0, 0, 0)})
+ assert task.composer_external_task_ids is None
+
+ def test_task_id_and_task_ids_are_exclusive_after_render(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id=TEST_COMPOSER_EXTERNAL_TASK_ID,
+ composer_external_task_ids=[TEST_COMPOSER_EXTERNAL_TASK_ID],
+ )
+
+ with pytest.raises(ValueError, match="Only one of
`composer_external_task_id`"):
+ task.poke(context={"logical_date": datetime(2024, 5, 23, 0, 0, 0)})
+
+ def test_empty_task_ids_list_normalizes_before_exclusivity_check(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+
composer_external_task_group_id=TEST_COMPOSER_EXTERNAL_TASK_GROUP_ID,
+ composer_external_task_ids=[],
+ )
+
+ assert task.composer_external_task_ids == []
+
+ with (
+ mock.patch.object(task, "_get_composer_airflow_version",
return_value=3),
+ mock.patch.object(task, "poke", return_value=True),
+ ):
+ task.execute(context={})
+
+ assert task.composer_external_task_ids is None
+ assert task.composer_external_task_group_id ==
TEST_COMPOSER_EXTERNAL_TASK_GROUP_ID
+
+ def test_empty_task_ids_use_dag_state_validation_after_normalize(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_ids=[],
+ allowed_states=["up_for_retry"],
+ )
+
+ assert task.composer_external_task_ids == []
+ with pytest.raises(ValueError, match="composer_external_task_id.*is
`None`"):
+ task.poke(context={"logical_date": datetime(2024, 5, 23, 0, 0, 0)})
+
+ def test_templated_task_id_is_normalized_after_rendering(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id="{{ target_task_id }}",
+ )
+
+ assert task.composer_external_task_ids is None
+ task.render_template_fields(context={"target_task_id":
TEST_COMPOSER_EXTERNAL_TASK_ID})
+ assert task.composer_external_task_id == TEST_COMPOSER_EXTERNAL_TASK_ID
+ assert task.composer_external_task_ids is None
+
+ with (
+ mock.patch.object(task, "_get_composer_airflow_version",
return_value=3),
+ mock.patch.object(task, "poke", return_value=True),
+ ):
+ task.execute(context={})
+
+ assert task.composer_external_task_ids ==
[TEST_COMPOSER_EXTERNAL_TASK_ID]
+
+ def test_state_validation_uses_rendered_target(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id=lambda *, context, jinja_env: None,
+ allowed_states=["up_for_retry"],
+ )
+
+ task.render_template_fields(context={})
+ assert task.composer_external_task_id is None
+
+ with pytest.raises(ValueError, match="composer_external_task_id.*is
`None`"):
+ task.execute(context={})
+
+ def test_templated_empty_task_ids_do_not_conflict_with_task_group(self):
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+
composer_external_task_group_id=TEST_COMPOSER_EXTERNAL_TASK_GROUP_ID,
+ composer_external_task_ids=lambda *, context, jinja_env: [],
+ )
+
+ task.render_template_fields(context={})
+ assert task.composer_external_task_ids == []
+
+ with (
+ mock.patch.object(task, "_get_composer_airflow_version",
return_value=3),
+ mock.patch.object(task, "poke", return_value=True),
+ ):
+ task.execute(context={})
+
+ assert task.composer_external_task_ids is None
+ assert task.composer_external_task_group_id ==
TEST_COMPOSER_EXTERNAL_TASK_GROUP_ID
+
+ def test_execute_complete_normalizes_single_task_id_for_failed_event(self):
+ """Deferrable resume hits execute_complete on a fresh instance;
normalize first."""
+ task = CloudComposerExternalTaskSensor(
+ task_id="task-id",
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ environment_id=TEST_ENVIRONMENT_ID,
+ composer_external_dag_id="test_dag_id",
+ composer_external_task_id="task_a",
+ deferrable=True,
+ )
+
+ assert task.composer_external_task_ids is None
+ with pytest.raises(ExternalTaskFailedError):
+ task.execute_complete(context={}, event={"status": "failed"})
+ assert task.composer_external_task_ids == ["task_a"]
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index c49a448fcd3..675813b6031 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -11,4 +11,3 @@
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/cloud_composer.py::CloudComposerExternalTaskSensor
