This is an automated email from the ASF dual-hosted git repository.
weilee 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 6cc2232a79d fix(hitl): handle hitl details when task instance is
retried (#53824)
6cc2232a79d is described below
commit 6cc2232a79d19bea4b0dc3df388623541aef2d4b
Author: Wei Lee <[email protected]>
AuthorDate: Tue Jul 29 06:29:39 2025 +0300
fix(hitl): handle hitl details when task instance is retried (#53824)
---
.../api_fastapi/execution_api/routes/hitl.py | 53 ++++++----
.../execution_api/versions/head/test_hitl.py | 116 ++++++++++++---------
.../airflow/providers/standard/operators/hitl.py | 4 +-
task-sdk/src/airflow/sdk/execution_time/hitl.py | 2 +-
.../tests/task_sdk/execution_time/test_hitl.py | 6 +-
5 files changed, 110 insertions(+), 71 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
index e1f3736662b..36d50b19fa7 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
@@ -40,32 +40,49 @@ log = structlog.get_logger(__name__)
"/{task_instance_id}",
status_code=status.HTTP_201_CREATED,
)
-def add_hitl_detail(
+def upsert_hitl_detail(
task_instance_id: UUID,
payload: HITLDetailRequest,
session: SessionDep,
) -> HITLDetailRequest:
- """Get Human-in-the-loop detail for a specific Task Instance."""
+ """
+ Create a Human-in-the-loop detail for a specific Task Instance.
+
+ There're 3 cases handled here.
+
+ 1. If a HITLOperator task instance does not have a HITLDetail,
+ a new HITLDetail is created without a response section.
+ 2. If a HITLOperator task instance has a HITLDetail but lacks a response,
+ the existing HITLDetail is returned.
+ This situation occurs when a task instance is cleared before a response
is received.
+ 3. If a HITLOperator task instance has both a HITLDetail and a response
section,
+ the existing response is removed, and the HITLDetail is returned.
+ This happens when a task instance is cleared after a response has been
received.
+ This design ensures that each task instance has only one HITLDetail.
+ """
ti_id_str = str(task_instance_id)
hitl_detail_model =
session.scalar(select(HITLDetail).where(HITLDetail.ti_id == ti_id_str))
- if hitl_detail_model:
- raise HTTPException(
- status.HTTP_409_CONFLICT,
- f"Human-in-the-loop detail for Task Instance with id {ti_id_str}
already exists.",
+ if not hitl_detail_model:
+ hitl_detail_model = HITLDetail(
+ ti_id=ti_id_str,
+ options=payload.options,
+ subject=payload.subject,
+ body=payload.body,
+ defaults=payload.defaults,
+ multiple=payload.multiple,
+ params=payload.params,
)
+ session.add(hitl_detail_model)
+ elif hitl_detail_model.response_received:
+ # Cleanup the response part of HITLDetail as we only store one
response for one task instance.
+ # It normally happens after retry, we keep only the latest response.
+ hitl_detail_model.user_id = None
+ hitl_detail_model.response_at = None
+ hitl_detail_model.chosen_options = None
+ hitl_detail_model.params_input = {}
+ session.add(hitl_detail_model)
- hitl_detail = HITLDetail(
- ti_id=ti_id_str,
- options=payload.options,
- subject=payload.subject,
- body=payload.body,
- defaults=payload.defaults,
- multiple=payload.multiple,
- params=payload.params,
- )
- session.add(hitl_detail)
- session.commit()
- return HITLDetailRequest.model_validate(hitl_detail)
+ return HITLDetailRequest.model_validate(hitl_detail_model)
@router.patch("/{task_instance_id}")
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py
index 33247304770..ec61542c485 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py
@@ -16,43 +16,60 @@
# under the License.
from __future__ import annotations
-from datetime import datetime
-
import pytest
-import time_machine
-from uuid6 import uuid7
+from httpx import Client
from tests_common.test_utils.db import AIRFLOW_V_3_1_PLUS
if not AIRFLOW_V_3_1_PLUS:
pytest.skip("Human in the loop public API compatible with Airflow >=
3.0.1", allow_module_level=True)
+from datetime import datetime
from typing import TYPE_CHECKING, Any
+import time_machine
+
+from airflow._shared.timezones.timezone import convert_to_utc
from airflow.models.hitl import HITLDetail
if TYPE_CHECKING:
+ from fastapi.testclient import TestClient
+ from sqlalchemy.orm import Session
+
from airflow.models.taskinstance import TaskInstance
+ from tests_common.pytest_plugin import CreateTaskInstance
+
pytestmark = pytest.mark.db_test
-TI_ID = uuid7()
+
+default_hitl_detail_request_kwargs: dict[str, Any] = {
+ # ti_id decided at a later stage
+ "subject": "This is subject",
+ "body": "this is body",
+ "options": ["Approve", "Reject"],
+ "defaults": ["Approve"],
+ "multiple": False,
+ "params": {"input_1": 1},
+}
+expected_empty_hitl_detail_response_part: dict[str, Any] = {
+ "response_at": None,
+ "chosen_options": None,
+ "user_id": None,
+ "params_input": {},
+ "response_received": False,
+}
@pytest.fixture
-def sample_ti(create_task_instance) -> TaskInstance:
+def sample_ti(create_task_instance: CreateTaskInstance) -> TaskInstance:
return create_task_instance()
@pytest.fixture
-def sample_hitl_detail(session, sample_ti) -> HITLDetail:
+def sample_hitl_detail(session: Session, sample_ti: TaskInstance) ->
HITLDetail:
hitl_detail_model = HITLDetail(
ti_id=sample_ti.id,
- options=["Approve", "Reject"],
- subject="This is subject",
- body="this is body",
- defaults=["Approve"],
- multiple=False,
- params={"input_1": 1},
+ **default_hitl_detail_request_kwargs,
)
session.add(hitl_detail_model)
session.commit()
@@ -61,54 +78,65 @@ def sample_hitl_detail(session, sample_ti) -> HITLDetail:
@pytest.fixture
-def expected_sample_hitl_detail_dict(sample_ti) -> dict[str, Any]:
+def expected_sample_hitl_detail_dict(sample_ti: TaskInstance) -> dict[str,
Any]:
return {
- "body": "this is body",
- "defaults": ["Approve"],
- "multiple": False,
- "options": ["Approve", "Reject"],
- "params": {"input_1": 1},
- "params_input": {},
- "response_at": None,
- "chosen_options": None,
- "response_received": False,
- "subject": "This is subject",
"ti_id": sample_ti.id,
- "user_id": None,
+ **default_hitl_detail_request_kwargs,
+ **expected_empty_hitl_detail_response_part,
}
-def test_add_hitl_detail(client, create_task_instance, session) -> None:
[email protected](
+ "existing_hitl_detail_args",
+ [
+ None,
+ default_hitl_detail_request_kwargs,
+ {
+ **default_hitl_detail_request_kwargs,
+ **{
+ "params_input": {"input_1": 2},
+ "response_at": convert_to_utc(datetime(2025, 7, 3, 0, 0, 0)),
+ "chosen_options": ["Reject"],
+ "user_id": "Fallback to defaults",
+ },
+ },
+ ],
+ ids=[
+ "no existing hitl detail",
+ "existing hitl detail without response",
+ "existing hitl detail with response",
+ ],
+)
+def test_upsert_hitl_detail(
+ client: TestClient,
+ create_task_instance: CreateTaskInstance,
+ session: Session,
+ existing_hitl_detail_args: dict[str, Any],
+) -> None:
ti = create_task_instance()
session.commit()
+ if existing_hitl_detail_args:
+ session.add(HITLDetail(ti_id=ti.id, **existing_hitl_detail_args))
+ session.commit()
+
response = client.post(
f"/execution/hitl-details/{ti.id}",
json={
"ti_id": ti.id,
- "options": ["Approve", "Reject"],
- "subject": "This is subject",
- "body": "this is body",
- "defaults": ["Approve"],
- "multiple": False,
- "params": {"input_1": 1},
+ **default_hitl_detail_request_kwargs,
},
)
assert response.status_code == 201
assert response.json() == {
"ti_id": ti.id,
- "options": ["Approve", "Reject"],
- "subject": "This is subject",
- "body": "this is body",
- "defaults": ["Approve"],
- "multiple": False,
- "params": {"input_1": 1},
+ **default_hitl_detail_request_kwargs,
}
@time_machine.travel(datetime(2025, 7, 3, 0, 0, 0), tick=False)
@pytest.mark.usefixtures("sample_hitl_detail")
-def test_update_hitl_detail(client, sample_ti) -> None:
+def test_update_hitl_detail(client: Client, sample_ti: TaskInstance) -> None:
response = client.patch(
f"/execution/hitl-details/{sample_ti.id}",
json={
@@ -128,13 +156,7 @@ def test_update_hitl_detail(client, sample_ti) -> None:
@pytest.mark.usefixtures("sample_hitl_detail")
-def test_get_hitl_detail(client, sample_ti) -> None:
+def test_get_hitl_detail(client: Client, sample_ti: TaskInstance) -> None:
response = client.get(f"/execution/hitl-details/{sample_ti.id}")
assert response.status_code == 200
- assert response.json() == {
- "params_input": {},
- "response_at": None,
- "chosen_options": None,
- "response_received": False,
- "user_id": None,
- }
+ assert response.json() == expected_empty_hitl_detail_response_part
diff --git
a/providers/standard/src/airflow/providers/standard/operators/hitl.py
b/providers/standard/src/airflow/providers/standard/operators/hitl.py
index 891a78908e4..f89b840c101 100644
--- a/providers/standard/src/airflow/providers/standard/operators/hitl.py
+++ b/providers/standard/src/airflow/providers/standard/operators/hitl.py
@@ -34,7 +34,7 @@ from airflow.providers.standard.exceptions import
HITLTimeoutError, HITLTriggerE
from airflow.providers.standard.triggers.hitl import HITLTrigger,
HITLTriggerEventSuccessPayload
from airflow.providers.standard.utils.skipmixin import SkipMixin
from airflow.sdk.definitions.param import ParamsDict
-from airflow.sdk.execution_time.hitl import add_hitl_detail
+from airflow.sdk.execution_time.hitl import upsert_hitl_detail
if TYPE_CHECKING:
from airflow.sdk.definitions.context import Context
@@ -98,7 +98,7 @@ class HITLOperator(BaseOperator):
"""Add a Human-in-the-loop Response and then defer to HITLTrigger and
wait for user input."""
ti_id = context["task_instance"].id
# Write Human-in-the-loop input request to DB
- add_hitl_detail(
+ upsert_hitl_detail(
ti_id=ti_id,
options=self.options,
subject=self.subject,
diff --git a/task-sdk/src/airflow/sdk/execution_time/hitl.py
b/task-sdk/src/airflow/sdk/execution_time/hitl.py
index 6f6409b3dca..348f14b838a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/hitl.py
+++ b/task-sdk/src/airflow/sdk/execution_time/hitl.py
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
from airflow.api_fastapi.execution_api.datamodels.hitl import
HITLDetailResponse
-def add_hitl_detail(
+def upsert_hitl_detail(
ti_id: UUID,
options: list[str],
subject: str,
diff --git a/task-sdk/tests/task_sdk/execution_time/test_hitl.py
b/task-sdk/tests/task_sdk/execution_time/test_hitl.py
index a936812bbb7..5f7cd373484 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_hitl.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_hitl.py
@@ -23,16 +23,16 @@ from airflow.sdk import timezone
from airflow.sdk.api.datamodels._generated import HITLDetailResponse
from airflow.sdk.execution_time.comms import CreateHITLDetailPayload
from airflow.sdk.execution_time.hitl import (
- add_hitl_detail,
get_hitl_detail_content_detail,
update_htil_detail_response,
+ upsert_hitl_detail,
)
TI_ID = uuid7()
-def test_add_hitl_detail(mock_supervisor_comms) -> None:
- add_hitl_detail(
+def test_upsert_hitl_detail(mock_supervisor_comms) -> None:
+ upsert_hitl_detail(
ti_id=TI_ID,
options=["Approve", "Reject"],
subject="Subject",
