AIP-72: Fix recursion bug with XComArg

airflow/models/baseoperator.py

@@ -25,7 +25,6 @@
 
 import collections.abc
 import contextlib
-import copy
 import functools
 import logging
 from collections.abc import Collection, Iterable, Sequence
@@ -703,12 +702,6 @@ def get_outlet_defs(self):
     extended/overridden by subclasses.
     """
 
-    def prepare_for_execution(self) -> BaseOperator:
-        """Lock task for execution to disable custom action in ``__setattr__`` and return a copy."""
-        other = copy.copy(self)
-        other._lock_for_execution = True
-        return other
-
     @prepare_lineage
     def pre_execute(self, context: Any):
         """Execute right before self.execute() is called."""

task_sdk/src/airflow/sdk/definitions/baseoperator.py

@@ -974,7 +974,7 @@ def __deepcopy__(self, memo: dict[int, Any]):
 
     def __getstate__(self):
         state = dict(self.__dict__)
-        if self._log:
+        if "_log" in state:
             del state["_log"]
 
         return state
@@ -1219,6 +1219,12 @@ def get_serialized_fields(cls):
 
         return cls.__serialized_fields
 
+    def prepare_for_execution(self) -> BaseOperator:
+        """Lock task for execution to disable custom action in ``__setattr__`` and return a copy."""
+        other = copy.copy(self)
+        other._lock_for_execution = True
+        return other
+
     def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
         """Serialize; required by DAGNode."""
         from airflow.serialization.enums import DagAttributeTypes

task_sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -271,6 +271,7 @@ def run(ti: RuntimeTaskInstance, log: Logger):
         # TODO: pre execute etc.
         # TODO next_method to support resuming from deferred
         # TODO: Get a real context object
+        ti.task = ti.task.prepare_for_execution()
         context = ti.get_template_context()
         ti.task.execute(context)  # type: ignore[attr-defined]
         msg = TaskState(state=TerminalTIState.SUCCESS, end_date=datetime.now(tz=timezone.utc))

task_sdk/tests/defintions/test_baseoperator.py

@@ -211,18 +211,12 @@ def test_warnings_are_properly_propagated(self):
             assert warning.filename == __file__
 
     def test_setattr_performs_no_custom_action_at_execute_time(self, spy_agency):
-        from airflow.models.xcom_arg import XComArg
-
         op = MockOperator(task_id="test_task")
 
-        op._lock_for_execution = True
-        # TODO: Task-SDK
-        # op_copy = op.prepare_for_execution()
-        op_copy = op
-
-        spy_agency.spy_on(XComArg.apply_upstream_relationship, call_original=False)
+        op_copy = op.prepare_for_execution()
+        spy_agency.spy_on(op._set_xcomargs_dependency, call_original=False)
         op_copy.arg1 = "b"
-        assert XComArg.apply_upstream_relationship.called is False
+        assert op._set_xcomargs_dependency.called is False
 
     def test_upstream_is_set_when_template_field_is_xcomarg(self):
         with DAG("xcomargs_test", schedule=None):