This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 64002a40580 Explicitly rollback on task state update exception (#71076)
64002a40580 is described below
commit 64002a40580f3990affb97a869e8d55ff79cc637
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Tue Aug 4 19:17:13 2026 +0800
Explicitly rollback on task state update exception (#71076)
---
.../execution_api/routes/task_instances.py | 1 +
.../versions/head/test_task_instances.py | 42 ++++++++++++++++++++++
2 files changed, 43 insertions(+)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index b45aa4e2655..41ecf49b053 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -469,6 +469,7 @@ def ti_update_state(
"Error updating Task Instance state. Setting the task to failed.",
payload=ti_patch_payload,
)
+ session.rollback()
ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
if session.bind is not None:
query = TI.duration_expression_update(timezone.utcnow(), query,
session.bind)
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index a744db8bf10..bb3c0f7e5a7 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -1331,6 +1331,48 @@ class TestTIUpdateState:
assert asset_listener.changed[0].uri == "s3://bucket/my-task"
assert len(asset_listener.emitted) == 1
+ def test_ti_update_state_to_success_rolls_back_partial_asset_registration(
+ self, client, session, create_task_instance
+ ):
+ """A failure partway through asset registration rolls back the partial
writes.
+
+ The endpoint's exception handler marks the TI failed and commits; the
explicit rollback
+ ensures an asset event flushed before the failure is not committed
alongside it.
+ """
+ asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task",
group="asset", extra={})
+ session.add_all([asset, AssetActive.for_asset(asset)])
+
+ ti = create_task_instance(
+
task_id="test_ti_update_state_to_success_rolls_back_partial_asset_registration",
+ start_date=DEFAULT_START_DATE,
+ state=State.RUNNING,
+ )
+ session.commit()
+
+ def _partial_then_fail(ti, task_outlets, outlet_events, *, session):
+ # Simulate a half-way registration: an asset event is flushed,
then registration
+ # fails before the endpoint commits.
+ session.add(AssetEvent(asset_id=asset.id))
+ session.flush()
+ raise RuntimeError("boom partway through outlets")
+
+ with mock.patch.object(TaskInstance, "register_asset_changes_in_db",
side_effect=_partial_then_fail):
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": "success",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "task_outlets": [{"name": "my-task", "uri":
"s3://bucket/my-task", "type": "Asset"}],
+ "outlet_events": [],
+ },
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ # The partially-written asset event was rolled back, and the TI is
marked failed.
+ assert session.scalars(select(AssetEvent)).all() == []
+ assert session.get(TaskInstance, ti.id).state == State.FAILED
+
@pytest.mark.parametrize(
("outlet_events", "expected_extra"),
[
