This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a0d7ae6070b Only log task state cleanup for configured worker state
store backends (#68878)
a0d7ae6070b is described below
commit a0d7ae6070ba2854dc889383a20f20486b634efe
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Jun 23 13:30:53 2026 +0530
Only log task state cleanup for configured worker state store backends
(#68878)
* Only log task state cleanup for configured worker state store backends
* fixing CI and adding test for wei
---
.../src/airflow/sdk/execution_time/task_runner.py | 3 +-
.../task_sdk/execution_time/test_task_runner.py | 33 ++++++++++++++++++++--
2 files changed, 33 insertions(+), 3 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 0c5cd7e8746..9160432ce0d 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -130,6 +130,7 @@ from airflow.sdk.execution_time.context import (
TaskStateStoreAccessor,
TriggeringAssetEventsAccessor,
VariableAccessor,
+ _get_worker_state_store_backend,
context_get_outlet_events,
context_to_airflow_vars,
get_previous_dagrun_success,
@@ -1733,7 +1734,7 @@ def _handle_current_task_success(
task_outlets = list(_build_asset_profiles(ti.task.outlets))
outlet_events = list(_serialize_outlet_events(context["outlet_events"]))
- if conf.getboolean("state_store", "clear_on_success"):
+ if conf.getboolean("state_store", "clear_on_success") and
_get_worker_state_store_backend() is not None:
log.info(
"Clearing task state from custom backend as clear_on_success is
enabled. The database references will be cleared by the API server."
)
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index b257b9ec800..4fd0692425b 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -5983,6 +5983,28 @@ class TestTaskInstanceStateOperations:
)
)
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_skips_backend_clear_when_no_custom_backend(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ """clear_on_success does not call backend.clear() when no custom
backend is configured."""
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ task = MyOperator(task_id="t")
+ runtime_ti = create_runtime_ti(task=task)
+
+ # no need to patch because default of
_get_worker_state_store_backend() is None
+ run(runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock())
+
+ sent_types = [
+ type(call.kwargs.get("msg") or (call.args[0] if call.args else
None))
+ for call in mock_supervisor_comms.send.call_args_list
+ ]
+ assert ClearTaskStateStore not in sent_types
+
@conf_vars({("state_store", "clear_on_success"): "True"})
def test_clear_on_success_clears_backend_without_comms_roundtrip(
self, create_runtime_ti, mock_supervisor_comms
@@ -5997,8 +6019,15 @@ class TestTaskInstanceStateOperations:
task = MyOperator(task_id="t")
runtime_ti = create_runtime_ti(task=task)
- with mock.patch(
-
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
return_value=mock_backend
+ 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,
+ ),
):
run(runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock())
