Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/execution_api/routes/hitl.py: 36%

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 uuid import UUID 

20 

21import structlog 

22from cadwyn import VersionedAPIRouter 

23from fastapi import HTTPException, Security, status 

24from sqlalchemy import select 

25 

26from airflow._shared.timezones import timezone 

27from airflow.api_fastapi.common.db.common import SessionDep 

28from airflow.api_fastapi.execution_api.datamodels.hitl import ( 

29 HITLDetailRequest, 

30 HITLDetailResponse, 

31 UpdateHITLDetailPayload, 

32) 

33from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute, require_auth 

34from airflow.models.hitl import HITLDetail 

35 

36router = VersionedAPIRouter( 

37 route_class=ExecutionAPIRoute, 

38 dependencies=[ 

39 # Validates that the JWT sub matches the task_instance_id path parameter. 

40 Security(require_auth, scopes=["ti:self"]), 

41 ], 

42) 

43 

44log = structlog.get_logger(__name__) 

45 

46 

47@router.post( 

48 "/{task_instance_id}", 

49 status_code=status.HTTP_201_CREATED, 

50) 

51def upsert_hitl_detail( 

52 task_instance_id: UUID, 

53 payload: HITLDetailRequest, 

54 session: SessionDep, 

55) -> HITLDetailRequest: 

56 """ 

57 Create a Human-in-the-loop detail for a specific Task Instance. 

58 

59 There're 3 cases handled here. 

60 

61 1. If a HITLOperator task instance does not have a HITLDetail, 

62 a new HITLDetail is created without a response section. 

63 2. If a HITLOperator task instance has a HITLDetail but lacks a response, 

64 the existing HITLDetail is returned. 

65 This situation occurs when a task instance is cleared before a response is received. 

66 3. If a HITLOperator task instance has both a HITLDetail and a response section, 

67 the existing response is removed, and the HITLDetail is returned. 

68 This happens when a task instance is cleared after a response has been received. 

69 This design ensures that each task instance has only one HITLDetail. 

70 """ 

71 hitl_detail_model = session.scalar(select(HITLDetail).where(HITLDetail.ti_id == task_instance_id)) 

72 if not hitl_detail_model: 

73 hitl_detail_model = HITLDetail( 

74 ti_id=task_instance_id, 

75 options=payload.options, 

76 subject=payload.subject, 

77 body=payload.body, 

78 defaults=payload.defaults, 

79 multiple=payload.multiple, 

80 params=payload.params, 

81 assignees=[user.model_dump() for user in payload.assigned_users], 

82 ) 

83 session.add(hitl_detail_model) 

84 elif hitl_detail_model.response_received: 

85 # Cleanup the response part of HITLDetail as we only store one response for one task instance. 

86 # It normally happens after retry, we keep only the latest response. 

87 hitl_detail_model.responded_by = None 

88 hitl_detail_model.responded_at = None 

89 hitl_detail_model.chosen_options = None 

90 hitl_detail_model.params_input = {} 

91 session.add(hitl_detail_model) 

92 

93 return HITLDetailRequest.model_validate(hitl_detail_model) 

94 

95 

96def _check_hitl_detail_exists(hitl_detail_model: HITLDetail | None) -> HITLDetail: 

97 if not hitl_detail_model: 

98 raise HTTPException( 

99 status.HTTP_404_NOT_FOUND, 

100 detail={ 

101 "reason": "not_found", 

102 "message": ( 

103 "HITLDetail not found. " 

104 "This happens most likely due to clearing task instance before receiving response." 

105 ), 

106 }, 

107 ) 

108 

109 return hitl_detail_model 

110 

111 

112@router.patch("/{task_instance_id}") 

113def update_hitl_detail( 

114 task_instance_id: UUID, 

115 payload: UpdateHITLDetailPayload, 

116 session: SessionDep, 

117) -> HITLDetailResponse: 

118 """Update the response part of a Human-in-the-loop detail for a specific Task Instance.""" 

119 hitl_detail_model_result = session.execute( 

120 select(HITLDetail).where(HITLDetail.ti_id == task_instance_id) 

121 ).scalar() 

122 hitl_detail_model = _check_hitl_detail_exists(hitl_detail_model_result) 

123 if hitl_detail_model.response_received: 

124 raise HTTPException( 

125 status.HTTP_409_CONFLICT, 

126 f"Human-in-the-loop detail for Task Instance with id {task_instance_id} already exists.", 

127 ) 

128 

129 hitl_detail_model.responded_by = None 

130 hitl_detail_model.responded_at = timezone.utcnow() 

131 hitl_detail_model.chosen_options = payload.chosen_options 

132 hitl_detail_model.params_input = payload.params_input 

133 session.add(hitl_detail_model) 

134 return HITLDetailResponse.from_hitl_detail_orm(hitl_detail_model) 

135 

136 

137@router.get( 

138 "/{task_instance_id}", 

139 status_code=status.HTTP_200_OK, 

140) 

141def get_hitl_detail( 

142 task_instance_id: UUID, 

143 session: SessionDep, 

144) -> HITLDetailResponse: 

145 """Get Human-in-the-loop detail for a specific Task Instance.""" 

146 hitl_detail_model_result = session.execute( 

147 select(HITLDetail).where(HITLDetail.ti_id == task_instance_id), 

148 ).scalar() 

149 hitl_detail_model = _check_hitl_detail_exists(hitl_detail_model_result) 

150 return HITLDetailResponse.from_hitl_detail_orm(hitl_detail_model)