jason810496 commented on code in PR #53907: URL: https://github.com/apache/airflow/pull/53907#discussion_r2276714584
########## airflow-core/src/airflow/api_fastapi/core_api/services/public/hitl_shared_links.py: ########## @@ -0,0 +1,382 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Services for HITL shared links functionality.""" + +from __future__ import annotations + +import base64 +import json +import os +from collections.abc import Callable +from datetime import datetime, timedelta +from functools import wraps +from typing import Any, Literal, TypedDict +from urllib.parse import urlparse, urlunparse + +import structlog +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from fastapi import HTTPException, status +from sqlalchemy.orm import Session + +from airflow._shared.timezones.timezone import utcnow +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.core_api.datamodels.hitl import ( + HITLDetailResponse, + UpdateHITLDetailPayload, +) +from airflow.api_fastapi.core_api.security import GetUserDep +from airflow.api_fastapi.core_api.services.public.hitl import ( + _get_task_instance, + _update_hitl_detail, +) +from airflow.configuration import conf +from airflow.exceptions import HITLSharedLinkTimeout + +log = structlog.get_logger(__name__) + +HITL_LINK_TYPE = Literal["redirect", "respond"] + + +class HITLSharedLinkConfig(TypedDict): + """Configuration for generating a Human-in-the-look shared link.""" + + link_type: HITL_LINK_TYPE + chosen_options: list[str] + params_input: dict[str, Any] + expires_at: str + + +class HITLSharedLinkData(TypedDict): + """The data used to generate a Human-in-the-look shared token.""" + + # task instance identifier + ti_id: str + dag_id: str + dag_run_id: str + task_id: str + map_index: int | None + + shared_link_config: HITLSharedLinkConfig + + +def hitl_shared_link_enabled(func: Callable) -> Callable: Review Comment: Nice one! Having decorator is even better then using FastAPI Depends for check whether does the user enable HITL. ########## airflow-core/src/airflow/api_fastapi/core_api/services/public/hitl.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import structlog +from fastapi import HTTPException, status +from sqlalchemy import select +from sqlalchemy.orm import joinedload + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.hitl import ( + HITLDetail, + HITLDetailResponse, + UpdateHITLDetailPayload, +) +from airflow.api_fastapi.core_api.security import GetUserDep +from airflow.models.hitl import HITLDetail as HITLDetailModel +from airflow.models.taskinstance import TaskInstance as TI + +hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitlDetails") + +log = structlog.get_logger(__name__) + + +def _get_task_instance( + dag_id: str, + dag_run_id: str, + task_id: str, + session: SessionDep, + map_index: int | None = None, +) -> TI: + query = select(TI).where( + TI.dag_id == dag_id, + TI.run_id == dag_run_id, + TI.task_id == task_id, + ) + + if map_index is not None: + query = query.where(TI.map_index == map_index) + + task_instance = session.scalar(query) + if task_instance is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", + ) + if map_index is None and task_instance.map_index != -1: + raise HTTPException( + status.HTTP_404_NOT_FOUND, "Task instance is mapped, add the map_index value to the URL" + ) + + return task_instance + + +def _update_hitl_detail( + dag_id: str, + dag_run_id: str, + task_id: str, + update_hitl_detail_payload: UpdateHITLDetailPayload, + user: GetUserDep, + session: SessionDep, + map_index: int | None = None, +) -> HITLDetailResponse: + task_instance = _get_task_instance( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + session=session, + map_index=map_index, + ) + ti_id_str = str(task_instance.id) + hitl_detail_model = session.scalar(select(HITLDetailModel).where(HITLDetailModel.ti_id == ti_id_str)) + if not hitl_detail_model: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"Human-in-the-loop detail does not exist for Task Instance with id {ti_id_str}", + ) + + if hitl_detail_model.response_received: + raise HTTPException( + status.HTTP_409_CONFLICT, + f"Human-in-the-loop detail has already been updated for Task Instance with id {ti_id_str} " Review Comment: ```suggestion f"Human-in-the-loop detail has already been updated for Task Instance with id `{ti_id_str}` " ``` ########## airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl_shared_links.py: ########## @@ -0,0 +1,293 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from json import JSONDecodeError + +import structlog +from cryptography.exceptions import InvalidTag +from fastapi import Depends, HTTPException, Request, status +from fastapi.responses import RedirectResponse +from pydantic import HttpUrl + +from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.hitl import HITLDetailResponse +from airflow.api_fastapi.core_api.datamodels.hitl_shared_links import ( + GenerateSharedLinkRequest, + SharedLinkResponse, +) +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.security import GetUserDep, requires_access_dag +from airflow.api_fastapi.core_api.services.public.hitl_shared_links import ( + _generate_hitl_shared_link, + _retrieve_hitl_redirect_url, + _update_hitl_detail_through_token, + hitl_shared_link_enabled, +) +from airflow.exceptions import HITLSharedLinkTimeout + +hitl_shared_links_router = AirflowRouter(tags=["HumanInTheLoopSharedLinks"], prefix="/hitlSharedLinks") + +log = structlog.get_logger(__name__) + + +@hitl_shared_link_enabled +@hitl_shared_links_router.post( + "/generate/{dag_id}/{dag_run_id}/{task_id}", Review Comment: I don't have a strong opinion, but I'd like to raise a discussion about RESTful route naming. Is it necessary to include verbs like `generate`, `redirect`, or `respond` in the route names? In my view, using `hitlSharedLinks` as a prefix is already meaningful enough for the intended operations: - POST, GET `/hitlSharedLinks/{DAG_ID}/{DAG_RUN_ID}/{TASK_ID}` - POST, GET `/hitlSharedLinks/{token}` ########## airflow-core/src/airflow/api_fastapi/core_api/services/public/hitl.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import structlog +from fastapi import HTTPException, status +from sqlalchemy import select +from sqlalchemy.orm import joinedload + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.hitl import ( + HITLDetail, + HITLDetailResponse, + UpdateHITLDetailPayload, +) +from airflow.api_fastapi.core_api.security import GetUserDep +from airflow.models.hitl import HITLDetail as HITLDetailModel +from airflow.models.taskinstance import TaskInstance as TI + +hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitlDetails") + +log = structlog.get_logger(__name__) + + +def _get_task_instance( + dag_id: str, + dag_run_id: str, + task_id: str, + session: SessionDep, + map_index: int | None = None, +) -> TI: + query = select(TI).where( + TI.dag_id == dag_id, + TI.run_id == dag_run_id, + TI.task_id == task_id, + ) + + if map_index is not None: + query = query.where(TI.map_index == map_index) + + task_instance = session.scalar(query) + if task_instance is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", + ) + if map_index is None and task_instance.map_index != -1: + raise HTTPException( + status.HTTP_404_NOT_FOUND, "Task instance is mapped, add the map_index value to the URL" + ) + + return task_instance + + +def _update_hitl_detail( + dag_id: str, + dag_run_id: str, + task_id: str, + update_hitl_detail_payload: UpdateHITLDetailPayload, + user: GetUserDep, + session: SessionDep, + map_index: int | None = None, +) -> HITLDetailResponse: + task_instance = _get_task_instance( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + session=session, + map_index=map_index, + ) + ti_id_str = str(task_instance.id) + hitl_detail_model = session.scalar(select(HITLDetailModel).where(HITLDetailModel.ti_id == ti_id_str)) + if not hitl_detail_model: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"Human-in-the-loop detail does not exist for Task Instance with id {ti_id_str}", Review Comment: ```suggestion f"Human-in-the-loop detail does not exist for Task Instance with id `{ti_id_str}`", ``` ########## airflow-core/src/airflow/api_fastapi/core_api/services/public/hitl.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import structlog +from fastapi import HTTPException, status +from sqlalchemy import select +from sqlalchemy.orm import joinedload + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.hitl import ( + HITLDetail, + HITLDetailResponse, + UpdateHITLDetailPayload, +) +from airflow.api_fastapi.core_api.security import GetUserDep +from airflow.models.hitl import HITLDetail as HITLDetailModel +from airflow.models.taskinstance import TaskInstance as TI + +hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitlDetails") + +log = structlog.get_logger(__name__) + + +def _get_task_instance( + dag_id: str, + dag_run_id: str, + task_id: str, + session: SessionDep, + map_index: int | None = None, +) -> TI: + query = select(TI).where( + TI.dag_id == dag_id, + TI.run_id == dag_run_id, + TI.task_id == task_id, + ) + + if map_index is not None: + query = query.where(TI.map_index == map_index) + + task_instance = session.scalar(query) + if task_instance is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", + ) + if map_index is None and task_instance.map_index != -1: + raise HTTPException( + status.HTTP_404_NOT_FOUND, "Task instance is mapped, add the map_index value to the URL" + ) + + return task_instance + + +def _update_hitl_detail( + dag_id: str, + dag_run_id: str, + task_id: str, + update_hitl_detail_payload: UpdateHITLDetailPayload, + user: GetUserDep, + session: SessionDep, + map_index: int | None = None, +) -> HITLDetailResponse: + task_instance = _get_task_instance( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + session=session, + map_index=map_index, + ) + ti_id_str = str(task_instance.id) + hitl_detail_model = session.scalar(select(HITLDetailModel).where(HITLDetailModel.ti_id == ti_id_str)) Review Comment: Should we `joinedload` the TI here as well? The statement in `_get_hitl_detail ` include `.options(joinedload(HITLDetailModel.task_instance))`. ########## airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl_shared_links.py: ########## @@ -0,0 +1,409 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import datetime +from typing import TYPE_CHECKING, Any +from unittest.mock import ANY + +import pytest +from fastapi import status +from sqlalchemy import select + +from airflow._shared.timezones.timezone import convert_to_utc +from airflow.api_fastapi.core_api.services.public.hitl_shared_links import ( + HITL_LINK_TYPE, + HITLSharedLinkConfig, + HITLSharedLinkData, + decode_shared_data, +) +from airflow.models.hitl import HITLDetail +from airflow.models.taskinstance import TaskInstance as TI + +from tests_common.pytest_plugin import time_machine +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from fastapi.testclient import TestClient + from sqlalchemy.orm import Session + + from tests_common.pytest_plugin import CreateTaskInstance, TaskInstance + +pytestmark = pytest.mark.db_test + +TEST_SERVER_PREFIX = "http://testserver/api/v2"; + +DAG_ID = "test_hitl_shared_link_dag" +DAG_RUN_ID = "run_id" +TASK_ID = "sample_hitl_task" + +CURRENT_TIME = convert_to_utc(datetime.datetime(2025, 8, 14, 0, 0, 0)) +ONE_DAY_BEFORE_CURRENT = CURRENT_TIME - datetime.timedelta(days=1) +ONE_DAY_AFTER_CURRENT = CURRENT_TIME + datetime.timedelta(days=1) + + [email protected] +def enabled(): + with conf_vars({("api", "hitl_enable_shared_links"): True}): + yield + + [email protected] +def sample_ti(create_task_instance: CreateTaskInstance, session: Session) -> TaskInstance: + ti = create_task_instance(dag_id=DAG_ID, run_id=DAG_RUN_ID, task_id=TASK_ID, session=session) + session.commit() + return ti + + [email protected]("map_index", [None, -1]) +class TestGenerateSharedLink: + """Tests for /hitlSharedLinks/generate endpoints""" + + def gen_endpoint(self, map_index: int | None) -> str: + base = f"/hitlSharedLinks/generate/{DAG_ID}/{DAG_RUN_ID}/{TASK_ID}" + return base if map_index is None else f"{base}/{map_index}" + + @pytest.mark.parametrize( + "link_type, chosen_options_in_payload, params_input_payload", + (("redirect", None, None), ("respond", ["option 1"], None)), + ) + @time_machine.travel(CURRENT_TIME, tick=False) + def test_should_respond_201( + self, + test_client: TestClient, + sample_ti: TaskInstance, + link_type: HITL_LINK_TYPE, + chosen_options_in_payload: list[str] | None, + params_input_payload: dict[str, Any] | None, + map_index: int | None, + ) -> None: + endpoint = self.gen_endpoint(map_index) + secret_key = test_client.app.state.secret_key # type: ignore[attr-defined] + + expected_data = HITLSharedLinkData( + ti_id=sample_ti.id, + dag_id=DAG_ID, + dag_run_id=DAG_RUN_ID, + task_id=TASK_ID, + map_index=map_index, + shared_link_config=HITLSharedLinkConfig( + link_type=link_type, + expires_at=ONE_DAY_AFTER_CURRENT.isoformat(), + chosen_options=chosen_options_in_payload or [], + params_input=params_input_payload or {}, + ), + ) + + resp = test_client.post( + endpoint, + json={ + "link_type": link_type, + "expires_at": None, + "chosen_options": chosen_options_in_payload, + "params_input": params_input_payload, + }, + ) + assert resp.status_code == status.HTTP_201_CREATED + url = resp.json()["url"] + token = url.split("/")[-1] + assert url.startswith(f"{TEST_SERVER_PREFIX}/hitlSharedLinks/{link_type}/") + assert decode_shared_data(secret_key=secret_key, token=token) == expected_data Review Comment: Do we need additional unit test for `decode_shared_data` ? Since we are using `decode_shared_data` to validate the expected result in test for API. ( No strong opinion, as the `decode_shared_data` seems quite trivial ) ########## airflow-core/src/airflow/api_fastapi/core_api/services/public/hitl.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import structlog +from fastapi import HTTPException, status +from sqlalchemy import select +from sqlalchemy.orm import joinedload + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.hitl import ( + HITLDetail, + HITLDetailResponse, + UpdateHITLDetailPayload, +) +from airflow.api_fastapi.core_api.security import GetUserDep +from airflow.models.hitl import HITLDetail as HITLDetailModel +from airflow.models.taskinstance import TaskInstance as TI + +hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitlDetails") + +log = structlog.get_logger(__name__) + + +def _get_task_instance( + dag_id: str, + dag_run_id: str, + task_id: str, + session: SessionDep, + map_index: int | None = None, +) -> TI: + query = select(TI).where( + TI.dag_id == dag_id, + TI.run_id == dag_run_id, + TI.task_id == task_id, + ) + + if map_index is not None: + query = query.where(TI.map_index == map_index) + + task_instance = session.scalar(query) + if task_instance is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", + ) + if map_index is None and task_instance.map_index != -1: + raise HTTPException( + status.HTTP_404_NOT_FOUND, "Task instance is mapped, add the map_index value to the URL" + ) + + return task_instance + + +def _update_hitl_detail( + dag_id: str, + dag_run_id: str, + task_id: str, + update_hitl_detail_payload: UpdateHITLDetailPayload, + user: GetUserDep, + session: SessionDep, + map_index: int | None = None, +) -> HITLDetailResponse: + task_instance = _get_task_instance( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + session=session, + map_index=map_index, + ) + ti_id_str = str(task_instance.id) + hitl_detail_model = session.scalar(select(HITLDetailModel).where(HITLDetailModel.ti_id == ti_id_str)) Review Comment: It seems we can even get TaskInstance with HITLDetailModel in single join query. -- 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]
