Lee-W commented on code in PR #53907: URL: https://github.com/apache/airflow/pull/53907#discussion_r2278522528
########## 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: + """Check if 'api.hitl_enable_shared_links' is set to True in the config.""" + + @wraps(func) + def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any: + if not conf.getboolean("api", "hitl_enable_shared_links", fallback=False): + raise HTTPException( + status.HTTP_403_FORBIDDEN, + ( + "Human-in-the-look link sharing is disabled. " + "To enable, please set 'api.hitl_enable_shared_links' to true" + ), + ) + return func(*args, **kwargs) + + return wrapper + + +def encode_shared_data( + *, + # encrypt key + secret_key: str, + # task instance identifier + ti_id: str, + dag_id: str, + dag_run_id: str, + task_id: str, + map_index: int | None, + # link generatinon config + shared_link_config: HITLSharedLinkConfig, +) -> str: + """ + Generate a secure token for Human-in-the-loop shared links. + + :param ti_id: task_instance ID + :param dag_id: Dag ID + :param dag_run_id: Dag run ID + :param task_id: Task ID + :param map_index: Map index for mapped tasks + :param shared_link_config: the configuration for generating a Human-in-the-loop shared token. + link_type: Type of link ('redirect' or 'respond') + chosen_options: Chosen options for respond links + params_input: Parameters input for respond links + expires_at: Custom expiration time in hours + :param session: Database session (required to get task instance UUID) + + :return: jwt-encoded token + """ + token_data: HITLSharedLinkData = HITLSharedLinkData( + ti_id=ti_id, + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task_id, + map_index=map_index, + shared_link_config=shared_link_config, + ) + + encoded_json = json.dumps(token_data).encode() + aesgcm = AESGCM(secret_key.encode()) + iv = os.urandom(12) + ciphertext = aesgcm.encrypt(iv, encoded_json, None) + return base64.urlsafe_b64encode(iv + ciphertext).decode("utf-8") + + +def decode_shared_data( + *, + secret_key: str, + token: str, +) -> HITLSharedLinkData: + """ + Decode a shared link token. + + :param token: Base64-encoded token + :return: Decoded token data + :raises ValueError: If token is invalid or expired + """ + raw = base64.urlsafe_b64decode(token.encode()) + if len(raw) < 13: + raise ValueError("Token too short") + + iv, ciphertext = raw[:12], raw[12:] + aesgcm = AESGCM(secret_key.encode()) + plaintext = aesgcm.decrypt(iv, ciphertext, None) + return json.loads(plaintext) + + +def _infer_base_url(base_url: str) -> str: + if not base_url: + base_url = conf.get("api", "base_url") + if not base_url: + raise ValueError("API base_url is not configured") + return base_url + + +def _validate_redirect_config( + *, chosen_options: list[str] | None, params_input: dict[str, Any] | None +) -> None: + if chosen_options or params_input: + raise ValueError( + '"chosen_options" and "params_input" should be empty "when "link_type" is set to "redirect"' + ) + + +def _validate_respond_config(*, chosen_options: list[str] | None) -> None: + if not chosen_options: Review Comment: Sounds like a good idea 👀 My original thought was to defer it to task execution so that we can reduce one db query. but fail early seems to be better -- 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]
