Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/datamodels/hitl.py: 96%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17from __future__ import annotations 

18 

19from collections.abc import Iterable, Mapping 

20from datetime import datetime 

21from typing import Any 

22 

23from pydantic import Field, field_validator 

24 

25from airflow.api_fastapi.core_api.base import BaseModel 

26from airflow.api_fastapi.core_api.datamodels.task_instance_history import TaskInstanceHistoryResponse 

27from airflow.api_fastapi.core_api.datamodels.task_instances import TaskInstanceResponse 

28 

29 

30class UpdateHITLDetailPayload(BaseModel): 

31 """Schema for updating the content of a Human-in-the-loop detail.""" 

32 

33 chosen_options: list[str] = Field(min_length=1) 

34 params_input: Mapping = Field(default_factory=dict) 

35 

36 

37class HITLDetailResponse(BaseModel): 

38 """Response of updating a Human-in-the-loop detail.""" 

39 

40 responded_by: HITLUser 

41 responded_at: datetime 

42 chosen_options: list[str] = Field(min_length=1) 

43 params_input: Mapping = Field(default_factory=dict) 

44 

45 

46class HITLUser(BaseModel): 

47 """Schema for a Human-in-the-loop users.""" 

48 

49 id: str 

50 name: str 

51 

52 

53class BaseHITLDetail(BaseModel): 

54 """The common part within HITLDetail and HITLDetailHistory.""" 

55 

56 # User Request Detail 

57 options: list[str] = Field(min_length=1) 

58 subject: str 

59 body: str | None = None 

60 defaults: list[str] | None = None 

61 multiple: bool = False 

62 params: Mapping = Field(default_factory=dict) 

63 assigned_users: list[HITLUser] = Field(default_factory=list) 

64 created_at: datetime 

65 

66 # Response Content Detail 

67 responded_by_user: HITLUser | None = None 

68 responded_at: datetime | None = None 

69 chosen_options: list[str] | None = None 

70 params_input: dict[str, Any] = Field(default_factory=dict) 

71 

72 response_received: bool = False 

73 

74 @field_validator("params", mode="before") 

75 @classmethod 

76 def get_params(cls, params: dict[str, Any]) -> dict[str, Any]: 

77 """Convert params attribute to dict representation.""" 

78 return { 

79 key: value 

80 if BaseHITLDetail._is_param(value) 

81 else { 

82 "value": value, 

83 "description": None, 

84 "schema": {}, 

85 } 

86 for key, value in params.items() 

87 } 

88 

89 @staticmethod 

90 def _is_param(value: Any) -> bool: 

91 return isinstance(value, dict) and all(key in value for key in ("description", "schema", "value")) 

92 

93 

94class HITLDetail(BaseHITLDetail): 

95 """Schema for Human-in-the-loop detail.""" 

96 

97 task_instance: TaskInstanceResponse 

98 

99 

100class HITLDetailCollection(BaseModel): 

101 """Schema for a collection of Human-in-the-loop details.""" 

102 

103 hitl_details: Iterable[HITLDetail] 

104 total_entries: int 

105 

106 

107class HITLDetailHistory(BaseHITLDetail): 

108 """Schema for Human-in-the-loop detail history.""" 

109 

110 task_instance: TaskInstanceHistoryResponse