amoghrajesh commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4453045855
##########
airflow-core/docs/core-concepts/task-state-store.rst:
##########
@@ -316,4 +316,10 @@ When ``[state_store] clear_on_success = True``, all task
state store keys for a
``clear_on_success`` clears **task state store only**. Asset store is
scoped to the asset, not the task instance, and is never affected by this
setting. Asset store persists across runs and must be cleared explicitly.
+.. note::
+
+ With a custom ``[workers] state_store_backend``, the worker deletes the
backend's entries only
+ after the success has been reported. If that deletion fails the task stays
successful and the
+ failure is logged; the leftover entries are removed by retention or
overwritten by the next run.
Review Comment:
Corrected in the docs note and the `log.exception` text, which now carries
`dag_id`, `run_id`,
`task_id` and `map_index`. Confirmed all three legs: the CLI skips
non-metastore backends,
`cleanup()` is an unoverridden no-op, and `TaskScope` includes `run_id`.
The docstring was already rewritten for the `finalize()` move.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -7455,6 +7455,66 @@ def execute(self, context):
]
assert ClearTaskStateStore not in sent_types
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_runs_after_success_is_reported(self,
create_runtime_ti, mock_supervisor_comms):
+ """The backend is only cleared once SucceedTask has been reported to
the supervisor."""
+ mock_backend = mock.MagicMock()
+ order = []
+ mock_supervisor_comms.send.side_effect = lambda msg, **kwargs:
order.append(type(msg).__name__)
+ mock_backend.clear.side_effect = lambda *a, **kw: order.append("clear")
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
+
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ ):
+ state, _, _ = run(runtime_ti,
context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ assert state == TaskInstanceState.SUCCESS
+ assert order.index("SucceedTask") < order.index("clear")
+
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_failure_does_not_fail_the_task(self,
create_runtime_ti, mock_supervisor_comms):
+ """A backend that cannot be cleared leaves the task successful instead
of failing it."""
+ mock_backend = mock.MagicMock()
+ mock_backend.clear.side_effect = ConnectionError("object storage is
down")
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
+
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ ):
+ state, msg, error = run(
+ runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock()
+ )
+
+ assert state == TaskInstanceState.SUCCESS
+ assert error is None
+ assert isinstance(msg, SucceedTask)
+ assert runtime_ti._terminal_state_send_failed is False
Review Comment:
Applied all four
--
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]
