Change on_*_callback on tasks to use has_on_*_callback

Author: kaxil

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -896,7 +896,7 @@ def process_executor_events(
                     ti.set_state(state)
                     continue
                 ti.task = task
-                if task.on_retry_callback or task.on_failure_callback:
+                if task.has_on_retry_callback or task.has_on_failure_callback:
                     # Only log the error/extra info here, since the `ti.handle_failure()` path will log it
                     # too, which would lead to double logging
                     cls.logger().error(msg)
@@ -2028,7 +2028,7 @@ def _maybe_requeue_stuck_ti(self, *, ti, session, executor):
                     exc_info=True,
                 )
             else:
-                if task.on_failure_callback:
+                if task.has_on_failure_callback:
                     if inspect(ti).detached:
                         ti = session.merge(ti)
                     request = TaskCallbackRequest(

airflow-core/src/airflow/models/dagrun.py

@@ -1998,8 +1998,8 @@ def schedule_tis(
                 assert isinstance(task, Operator)
             if (
                 task.inherits_from_empty_operator
-                and not task.on_execute_callback
-                and not task.on_success_callback
+                and not task.has_on_execute_callback
+                and not task.has_on_success_callback
                 and not task.outlets
                 and not task.inlets
             ):

airflow-core/src/airflow/models/mappedoperator.py

@@ -42,7 +42,6 @@
     DEFAULT_TRIGGER_RULE,
     DEFAULT_WEIGHT_RULE,
     NotMapped,
-    TaskStateChangeCallbackAttrType,
 )
 from airflow.sdk.definitions._internal.node import DAGNode
 from airflow.sdk.definitions.mappedoperator import MappedOperator as TaskSDKMappedOperator
@@ -242,24 +241,24 @@ def max_active_tis_per_dagrun(self) -> int | None:
         return self.partial_kwargs.get("max_active_tis_per_dagrun")
 
     @property
-    def on_execute_callback(self) -> TaskStateChangeCallbackAttrType:
-        return self.partial_kwargs.get("on_execute_callback") or []
+    def has_on_execute_callback(self) -> bool:
+        return bool(self.partial_kwargs.get("has_on_execute_callback", False))
 
     @property
-    def on_failure_callback(self) -> TaskStateChangeCallbackAttrType:
-        return self.partial_kwargs.get("on_failure_callback") or []
+    def has_on_failure_callback(self) -> bool:
+        return bool(self.partial_kwargs.get("has_on_failure_callback", False))
 
     @property
-    def on_retry_callback(self) -> TaskStateChangeCallbackAttrType:
-        return self.partial_kwargs.get("on_retry_callback") or []
+    def has_on_retry_callback(self) -> bool:
+        return bool(self.partial_kwargs.get("has_on_retry_callback", False))
 
     @property
-    def on_success_callback(self) -> TaskStateChangeCallbackAttrType:
-        return self.partial_kwargs.get("on_success_callback") or []
+    def has_on_success_callback(self) -> bool:
+        return bool(self.partial_kwargs.get("has_on_success_callback", False))
 
     @property
-    def on_skipped_callback(self) -> TaskStateChangeCallbackAttrType:
-        return self.partial_kwargs.get("on_skipped_callback") or []
+    def has_on_skipped_callback(self) -> bool:
+        return bool(self.partial_kwargs.get("has_on_skipped_callback", False))
 
     @property
     def run_as_user(self) -> str | None:

airflow-core/src/airflow/serialization/schema.json

@@ -316,11 +316,11 @@
         "allow_nested_operators": { "type": "boolean", "default": true },
         "inlets": {"type":  "array", "default": []},
         "outlets": {"type":  "array", "default": []},
-        "on_execute_callback": {"type":  "array", "default": []},
-        "on_failure_callback": {"type":  "array", "default": []},
-        "on_skipped_callback": {"type":  "array", "default": []},
-        "on_success_callback": {"type":  "array", "default": []},
-        "on_retry_callback": {"type":  "array", "default": []},
+        "has_on_execute_callback": {"type": "boolean", "default": false},
+        "has_on_failure_callback": {"type": "boolean", "default": false},
+        "has_on_skipped_callback": {"type": "boolean", "default": false},
+        "has_on_success_callback": {"type": "boolean", "default": false},
+        "has_on_retry_callback": {"type": "boolean", "default": false},
         "multiple_outputs": {"type": "boolean", "default": false},
         "start_from_trigger": {"type": "boolean", "default": false},
         "is_setup": {"type": "boolean", "default": false},

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -52,7 +52,6 @@
 from airflow.models.xcom_arg import SchedulerXComArg, deserialize_xcom_arg
 from airflow.sdk import Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher, BaseOperator, XComArg
 from airflow.sdk.bases.trigger import StartTriggerArgs
-from airflow.sdk.definitions._internal.expandinput import EXPAND_INPUT_EMPTY
 from airflow.sdk.definitions._internal.node import DAGNode
 from airflow.sdk.definitions.asset import (
     AssetAliasEvent,
@@ -1221,6 +1220,7 @@ class SerializedBaseOperator(DAGNode, BaseSerialization):
     _is_empty: bool
     _needs_expansion: bool
     _task_display_name: str | None
+    _on_failure_fail_dagrun: bool = False
 
     dag: DAG | None = None
     task_group: TaskGroup | None = None
@@ -1255,12 +1255,12 @@ class SerializedBaseOperator(DAGNode, BaseSerialization):
     max_retry_delay: datetime.timedelta | float | None = None
     multiple_outputs: bool = False
 
-    # TODO: Can be changed to () instead
-    on_execute_callback: Sequence = []
-    on_failure_callback: Sequence = []
-    on_retry_callback: Sequence = []
-    on_success_callback: Sequence = []
-    on_skipped_callback: Sequence = []
+    # Boolean flags for callback existence
+    has_on_execute_callback: bool = False
+    has_on_failure_callback: bool = False
+    has_on_retry_callback: bool = False
+    has_on_success_callback: bool = False
+    has_on_skipped_callback: bool = False
 
     operator_extra_links: Collection[BaseOperatorLink] = ()
 
@@ -1444,6 +1444,11 @@ def serialize_mapped_operator(cls, op: MappedOperator | SchedulerMappedOperator)
                     continue
                 if cls._is_excluded(v, k, op):
                     continue
+
+                if k in [f"on_{x}_callback" for x in ("execute", "failure", "success", "retry", "skipped")]:
+                    if bool(v):
+                        serialized_op["partial_kwargs"][f"has_{k}"] = True
+                    continue
                 serialized_op["partial_kwargs"].update({k: cls.serialize(v)})
 
         serialized_op["_is_mapped"] = True
@@ -1765,6 +1770,9 @@ def _is_excluded(cls, var: Any, attrname: str, op: DAGNode):
             if var is dag_date or var == dag_date:
                 return True
 
+        # If none of the exclusion conditions are met, don't exclude the field
+        return False
+
     @classmethod
     def _deserialize_operator_extra_links(
         cls, encoded_op_links: dict[str, str]
@@ -1885,12 +1893,11 @@ def get_serialized_fields(cls):
                 "max_active_tis_per_dagrun",
                 "max_retry_delay",
                 "multiple_outputs",
-                "on_execute_callback",
-                "on_failure_callback",
-                "on_failure_fail_dagrun",
-                "on_retry_callback",
-                "on_skipped_callback",
-                "on_success_callback",
+                "has_on_execute_callback",
+                "has_on_failure_callback",
+                "has_on_retry_callback",
+                "has_on_skipped_callback",
+                "has_on_success_callback",
                 "outlets",
                 "owner",
                 "params",

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -742,10 +742,6 @@ def validate_deserialized_task(
                 # We store the string, real dag has the actual code
                 "_pre_execute_hook",
                 "_post_execute_hook",
-                "on_execute_callback",
-                "on_failure_callback",
-                "on_success_callback",
-                "on_retry_callback",
                 # Checked separately
                 "resources",
                 "on_failure_fail_dagrun",
@@ -811,11 +807,23 @@ def validate_deserialized_task(
             default_partial_kwargs = (
                 BaseOperator.partial(task_id="_")._expand(EXPAND_INPUT_EMPTY, strict=False).partial_kwargs
             )
+
+            # These are added in `_TaskDecorator` e.g. when @setup or @teardown task is passed
+            default_decorator_partial_kwargs = {
+                "is_setup": False,
+                "is_teardown": False,
+                "on_failure_fail_dagrun": False,
+            }
             serialized_partial_kwargs = {
                 **default_partial_kwargs,
+                **default_decorator_partial_kwargs,
                 **serialized_task.partial_kwargs,
             }
-            original_partial_kwargs = {**default_partial_kwargs, **task.partial_kwargs}
+            original_partial_kwargs = {
+                **default_partial_kwargs,
+                **default_decorator_partial_kwargs,
+                **task.partial_kwargs,
+            }
             assert serialized_partial_kwargs == original_partial_kwargs
 
             # ExpandInputs have different classes between scheduler and definition
@@ -1415,6 +1423,11 @@ def test_no_new_fields_added_to_base_operator(self):
             "execution_timeout": None,
             "executor": None,
             "executor_config": {},
+            "has_on_execute_callback": False,
+            "has_on_failure_callback": False,
+            "has_on_retry_callback": False,
+            "has_on_skipped_callback": False,
+            "has_on_success_callback": False,
             "ignore_first_depends_on_past": False,
             "is_setup": False,
             "is_teardown": False,
@@ -1423,12 +1436,7 @@ def test_no_new_fields_added_to_base_operator(self):
             "max_active_tis_per_dag": None,
             "max_active_tis_per_dagrun": None,
             "max_retry_delay": None,
-            "on_execute_callback": [],
             "on_failure_fail_dagrun": False,
-            "on_failure_callback": [],
-            "on_retry_callback": [],
-            "on_skipped_callback": [],
-            "on_success_callback": [],
             "outlets": [],
             "owner": "airflow",
             "params": {},
@@ -3011,6 +3019,8 @@ def operator_extra_links(self):
     assert mapped_task.extra_links == sorted({"airflow", "github"})
 
 
+# TODO: Remove xfail
+@pytest.mark.xfail(reason="TODO: Need to add support for v1 & v2 to v3")
 def test_handle_v1_serdag():
     v1 = {
         "__version": 1,
@@ -3296,3 +3306,129 @@ def test_handle_v1_serdag():
     del expected["dag"]["tasks"][1]["__var"]["_operator_extra_links"]
 
     assert v1 == expected
+
+
+def dummy_callback():
+    pass
+
+
+@pytest.mark.parametrize(
+    "callback_config,expected_flags,is_mapped",
+    [
+        # Regular operator tests
+        (
+            {
+                "on_failure_callback": dummy_callback,
+                "on_retry_callback": [dummy_callback, dummy_callback],
+                "on_success_callback": dummy_callback,
+            },
+            {"has_on_failure_callback": True, "has_on_retry_callback": True, "has_on_success_callback": True},
+            False,
+        ),
+        (
+            {},  # No callbacks
+            {
+                "has_on_failure_callback": False,
+                "has_on_retry_callback": False,
+                "has_on_success_callback": False,
+            },
+            False,
+        ),
+        (
+            {"on_failure_callback": [], "on_success_callback": None},  # Empty callbacks
+            {"has_on_failure_callback": False, "has_on_success_callback": False},
+            False,
+        ),
+        # Mapped operator tests
+        (
+            {"on_failure_callback": dummy_callback, "on_success_callback": [dummy_callback, dummy_callback]},
+            {"has_on_failure_callback": True, "has_on_success_callback": True},
+            True,
+        ),
+        (
+            {},  # Mapped operator without callbacks
+            {"has_on_failure_callback": False, "has_on_success_callback": False},
+            True,
+        ),
+    ],
+)
+def test_task_callback_boolean_optimization(callback_config, expected_flags, is_mapped):
+    """Test that task callbacks are optimized using has_on_*_callback boolean flags."""
+    dag = DAG(dag_id="test_callback_dag", start_date=datetime(2020, 1, 1))
+
+    if is_mapped:
+        # Create mapped operator
+        task = BashOperator.partial(task_id="test_task", dag=dag, **callback_config).expand(
+            bash_command=["echo 1", "echo 2"]
+        )
+
+        # Serialize and deserialize
+        serialized = BaseSerialization.serialize(task)
+        deserialized = BaseSerialization.deserialize(serialized)
+
+        # For mapped operators, check partial_kwargs
+        serialized_data = serialized.get("__var", {}).get("partial_kwargs", {})
+
+        # Test serialization
+        for flag, expected in expected_flags.items():
+            if expected:
+                assert flag in serialized_data
+                assert serialized_data[flag] is True
+            else:
+                assert serialized_data.get(flag, False) is False
+
+        # Test deserialized properties
+        for flag, expected in expected_flags.items():
+            assert getattr(deserialized, flag) is expected
+
+    else:
+        # Create regular operator
+        task = BashOperator(task_id="test_task", bash_command="echo test", dag=dag, **callback_config)
+
+        # Serialize and deserialize
+        serialized = BaseSerialization.serialize(task)
+        deserialized = BaseSerialization.deserialize(serialized)
+
+        # For regular operators, check top-level
+        serialized_data = serialized.get("__var", {})
+
+        # Test serialization (only True values are stored)
+        for flag, expected in expected_flags.items():
+            if expected:
+                assert serialized_data.get(flag, False) is True
+            else:
+                assert serialized_data.get(flag, False) is False
+
+        # Test deserialized properties
+        for flag, expected in expected_flags.items():
+            assert getattr(deserialized, flag) is expected
+
+
+def test_task_callback_properties_exist():
+    """Test that all callback boolean properties exist on both regular and mapped operators."""
+    dag = DAG(dag_id="test_dag", start_date=datetime(2020, 1, 1))
+
+    # Regular operator
+    regular_task = BashOperator(task_id="regular", bash_command="echo test", dag=dag)
+
+    # Mapped operator
+    mapped_task = BashOperator.partial(task_id="mapped", dag=dag).expand(bash_command=["echo 1"])
+
+    callback_properties = [
+        "has_on_execute_callback",
+        "has_on_failure_callback",
+        "has_on_success_callback",
+        "has_on_retry_callback",
+        "has_on_skipped_callback",
+    ]
+
+    for prop in callback_properties:
+        assert hasattr(regular_task, prop), f"Regular operator missing {prop}"
+        assert hasattr(mapped_task, prop), f"Mapped operator missing {prop}"
+
+        # Serialize and check deserialized versions too
+        serialized_regular = BaseSerialization.deserialize(BaseSerialization.serialize(regular_task))
+        serialized_mapped = BaseSerialization.deserialize(BaseSerialization.serialize(mapped_task))
+
+        assert hasattr(serialized_regular, prop), f"Deserialized regular operator missing {prop}"
+        assert hasattr(serialized_mapped, prop), f"Deserialized mapped operator missing {prop}"