Fix sensor skipping in Airflow 3.x branching operators (#53455)

In Airflow 3.x, sensors inherit from airflow.sdk.BaseOperator instead of
airflow.models.BaseOperator. The _ensure_tasks function in SkipMixin was
only checking for the models BaseOperator, causing sensors to be filtered
out and not properly skipped by branching operators like BranchSQLOperator.

Updated the import logic to use the correct SDK BaseOperator for Airflow 3.x
and added comprehensive tests to verify sensors are properly included in
branching skip operations.

Fixes #52219

Author: kaxil

providers/standard/src/airflow/providers/standard/utils/skipmixin.py

@@ -22,7 +22,7 @@
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_1_PLUS
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
@@ -40,7 +40,7 @@
 
 
 def _ensure_tasks(nodes: Iterable[DAGNode]) -> Sequence[Operator]:
-    if AIRFLOW_V_3_1_PLUS:
+    if AIRFLOW_V_3_0_PLUS:
         from airflow.sdk import BaseOperator
         from airflow.sdk.definitions.mappedoperator import MappedOperator
     else:

providers/standard/tests/unit/standard/operators/test_python.py

@@ -845,6 +845,54 @@ def test_xcom_push_skipped_tasks(self):
             "skipped": ["empty_task"]
         }
 
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Airflow 2 implementation is different")
+    def test_short_circuit_operator_skips_sensors(self):
+        """Test that ShortCircuitOperator properly skips sensors in Airflow 3.x."""
+        from airflow.sdk.bases.sensor import BaseSensorOperator
+
+        # Create a sensor similar to S3FileSensor to reproduce the issue
+        class CustomS3Sensor(BaseSensorOperator):
+            def __init__(self, bucket_name: str, object_key: str, **kwargs):
+                super().__init__(**kwargs)
+                self.bucket_name = bucket_name
+                self.object_key = object_key
+                self.timeout = 0
+                self.poke_interval = 0
+
+            def poke(self, context):
+                # Simulate sensor logic
+                return True
+
+        with self.dag_maker(self.dag_id):
+            # ShortCircuit that evaluates to False (should skip all downstream)
+            short_circuit = ShortCircuitOperator(
+                task_id="check_dis_is_mon_to_fri_not_holiday",
+                python_callable=lambda: False,  # This causes skipping
+            )
+
+            sensor_task = CustomS3Sensor(
+                task_id="wait_for_ticker_to_secid_lookup_s3_file",
+                bucket_name="test-bucket",
+                object_key="ticker_to_secid_lookup.csv",
+            )
+
+            short_circuit >> sensor_task
+
+        dr = self.dag_maker.create_dagrun()
+
+        self.dag_maker.run_ti("check_dis_is_mon_to_fri_not_holiday", dr)
+
+        # Verify the sensor is included in the skip list by checking XCom
+        # (this was the bug - sensors were not being included in skip list)
+        tis = dr.get_task_instances()
+        xcom_data = tis[0].xcom_pull(task_ids="check_dis_is_mon_to_fri_not_holiday", key="skipmixin_key")
+
+        assert xcom_data is not None, "XCom data should exist"
+        skipped_task_ids = set(xcom_data.get("skipped", []))
+        assert "wait_for_ticker_to_secid_lookup_s3_file" in skipped_task_ids, (
+            "Sensor should be skipped by ShortCircuitOperator"
+        )
+
 
 virtualenv_string_args: list[str] = []
 

providers/standard/tests/unit/standard/utils/test_skipmixin.py

@@ -359,3 +359,84 @@ def test_raise_exception_on_not_valid_branch_task_ids(self, dag_maker, branch_ta
         error_message = r"'branch_task_ids' must contain only valid task_ids. Invalid tasks found: .*"
         with pytest.raises(AirflowException, match=error_message):
             SkipMixin().skip_all_except(ti=ti1, branch_task_ids=branch_task_ids)
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Issue only exists in Airflow 3.x")
+    def test_ensure_tasks_includes_sensors_airflow_3x(self, dag_maker):
+        """Test that sensors (inheriting from airflow.sdk.BaseOperator) are properly handled by _ensure_tasks."""
+        from airflow.providers.standard.utils.skipmixin import _ensure_tasks
+        from airflow.sdk import BaseOperator as SDKBaseOperator
+        from airflow.sdk.bases.sensor import BaseSensorOperator
+
+        class DummySensor(BaseSensorOperator):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.timeout = 0
+                self.poke_interval = 0
+
+            def poke(self, context):
+                return True
+
+        with dag_maker("dag_test_sensor_skipping") as dag:
+            regular_task = EmptyOperator(task_id="regular_task")
+            sensor_task = DummySensor(task_id="sensor_task")
+            downstream_task = EmptyOperator(task_id="downstream_task")
+
+            regular_task >> [sensor_task, downstream_task]
+
+        dag_maker.create_dagrun(run_id=DEFAULT_DAG_RUN_ID)
+
+        downstream_nodes = dag.get_task("regular_task").downstream_list
+        task_list = _ensure_tasks(downstream_nodes)
+
+        # Verify both the regular operator and sensor are included
+        task_ids = [t.task_id for t in task_list]
+        assert "sensor_task" in task_ids, "Sensor should be included in task list"
+        assert "downstream_task" in task_ids, "Regular task should be included in task list"
+        assert len(task_list) == 2, "Both tasks should be included"
+
+        # Also verify that the sensor is actually an instance of the correct BaseOperator
+        sensor_in_list = next((t for t in task_list if t.task_id == "sensor_task"), None)
+        assert sensor_in_list is not None, "Sensor task should be found in list"
+        assert isinstance(sensor_in_list, SDKBaseOperator), "Sensor should be instance of SDK BaseOperator"
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Integration test for Airflow 3.x sensor skipping")
+    def test_skip_sensor_in_branching_scenario(self, dag_maker):
+        """Integration test: verify sensors are properly skipped by branching operators in Airflow 3.x."""
+        from airflow.sdk.bases.sensor import BaseSensorOperator
+
+        # Create a dummy sensor for testing
+        class DummySensor(BaseSensorOperator):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.timeout = 0
+                self.poke_interval = 0
+
+            def poke(self, context):
+                return True
+
+        with dag_maker("dag_test_branch_sensor_skipping"):
+            branch_task = EmptyOperator(task_id="branch_task")
+            regular_task = EmptyOperator(task_id="regular_task")
+            sensor_task = DummySensor(task_id="sensor_task")
+            branch_task >> [regular_task, sensor_task]
+
+        dag_maker.create_dagrun(run_id=DEFAULT_DAG_RUN_ID)
+
+        dag_version = DagVersion.get_latest_version(branch_task.dag_id)
+        ti_branch = TI(branch_task, run_id=DEFAULT_DAG_RUN_ID, dag_version_id=dag_version.id)
+
+        # Test skipping the sensor (follow regular_task branch)
+        with pytest.raises(DownstreamTasksSkipped) as exc_info:
+            SkipMixin().skip_all_except(ti=ti_branch, branch_task_ids="regular_task")
+
+        # Verify that the sensor task is properly marked for skipping
+        skipped_tasks = set(exc_info.value.tasks)
+        assert ("sensor_task", -1) in skipped_tasks, "Sensor task should be marked for skipping"
+
+        # Test skipping the regular task (follow sensor_task branch)
+        with pytest.raises(DownstreamTasksSkipped) as exc_info:
+            SkipMixin().skip_all_except(ti=ti_branch, branch_task_ids="sensor_task")
+
+        # Verify that the regular task is properly marked for skipping
+        skipped_tasks = set(exc_info.value.tasks)
+        assert ("regular_task", -1) in skipped_tasks, "Regular task should be marked for skipping"