kaxil commented on code in PR #67877:
URL: https://github.com/apache/airflow/pull/67877#discussion_r3498389471
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -5597,6 +5600,141 @@ def test_exception_in_context_manager_propagates(self):
raise ValueError("boom")
+class TestRunExecuteCallable:
+ """Tests for ``_run_execute_callable``.
+
+ It runs the task's execute callable inside an isolated contextvars copy
(with
+ the ExecutorSafeguard tracker set), applies the execution timeout when one
is
+ configured, and wraps the call in a ``task.execute`` detail span.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _sampled_carrier_provider(self):
+ """Make new_dagrun_trace_carrier produce a SAMPLED carrier (see
TestDetailSpan)."""
+ provider = TracerProvider()
+ with mock.patch(
+ "airflow._shared.observability.traces.trace.get_tracer_provider",
+ return_value=provider,
+ ):
+ yield
+
+ @staticmethod
+ def _make_task(execution_timeout=None):
+ task = mock.MagicMock(spec=BaseOperator)
+ task.execution_timeout = execution_timeout
+ return task
+
+ def test_runs_in_isolated_context_with_safeguard_tracker_set(self):
+ """The callable runs in an internal context copy that has the
safeguard tracker set and does not leak."""
+ var = contextvars.ContextVar("marker")
+ var.set("outer")
+ task = self._make_task()
+ seen = {}
+
+ def execute(context):
+ var.set("inner")
+ seen["tracker"] = ExecutorSafeguard.tracker.get(None)
+ return context["value"] * 2
+
+ result = _run_execute_callable(context={"value": 21}, execute=execute,
task=task)
+
+ assert result == 42
+ # The safeguard tracker is set to the task inside the copy used to run
execute.
+ assert seen["tracker"] is task
+ # The mutation happened inside the copy, so it does not leak to the
caller's context.
+ assert var.get() == "outer"
Review Comment:
Small gap on the test named for it: the "does not leak" assertion here
(`assert var.get() == "outer"`) covers the separate `marker` var, but nothing
asserts the tracker itself is absent from the caller's real context after the
call. The headline behavior -- `ctx.run(ExecutorSafeguard.tracker.set, task)`
confining the `.set` to the copy -- isn't directly pinned. Would adding
something like `assert ExecutorSafeguard.tracker.get(None) is not task` after
the call be worth it? (No autouse fixture resets the tracker between tests, so
`is not task` is safer than `is None`.)
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -2065,6 +2065,43 @@ def _send_error_email_notification(
log.exception("Failed to send email notification")
+@detail_span("task.execute")
+def _run_execute_callable(
+ context: Context,
+ execute: Callable[..., Any] | functools.partial[Any],
+ task: BaseOperator,
+) -> Any:
+ """
+ Run the task's execute callable, applying the execution timeout if one is
set.
+
+ The contextvars snapshot is taken here, after the ``task.execute`` span is
+ current, so spans the operator emits during ``execute`` nest under it
rather
+ than under the caller. ``ExecutorSafeguard``'s tracker is set into that
copy
+ so the operator's ``execute`` passes the safeguard check, while the copy
keeps
+ the change from leaking into the surrounding context.
+ """
+ ctx = contextvars.copy_context()
Review Comment:
Confirming intent on the snapshot timing: moving `copy_context()` here
(after `pre_execute` / `on_execute_callback` run in `_execute_task`) means
`execute` now inherits any `ContextVar` those callbacks set, whereas before the
snapshot was taken at the top of `_execute_task`, ahead of them. I traced the
SDK contextvars (`_preset_connections` self-resets in
`Connection.test_connection`'s `finally`; `ExecutorSafeguard.tracker` is set
into this copy), so I don't see a real leak today, and sharing context with the
pre-hooks is arguably more consistent. Flagging the delta in case any operator
deliberately relied on `execute` not seeing pre-hook contextvar mutations --
since `on_execute_callback` exceptions are swallowed, a callback that sets a
var and then raises would now be visible to `execute`. Intended?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
