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 b45aa4e265560..41ecf49b053fb 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 a744db8bf10fe..bb3c0f7e5a785 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 @@ def test_ti_update_state_to_success_runs_deferred_asset_listener_callbacks( 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"), [