Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/common/exceptions.py: 94%

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. 

17 

18from __future__ import annotations 

19 

20import logging 

21import traceback 

22from abc import ABC, abstractmethod 

23from enum import Enum 

24from typing import Generic, TypeVar 

25 

26from fastapi import FastAPI, HTTPException, Request, status 

27from sqlalchemy.exc import DataError, IntegrityError, SQLAlchemyError 

28 

29from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT 

30from airflow.configuration import conf 

31from airflow.exceptions import DeserializationError 

32from airflow.utils.strings import get_random_string 

33 

34T = TypeVar("T", bound=Exception) 

35DBError = TypeVar("DBError", bound=SQLAlchemyError) 

36 

37log = logging.getLogger(__name__) 

38 

39 

40class BaseErrorHandler(Generic[T], ABC): 

41 """Base class for error handlers.""" 

42 

43 def __init__(self, exception_cls: T) -> None: 

44 self.exception_cls = exception_cls 

45 

46 @abstractmethod 

47 def exception_handler(self, request: Request, exc: T): 

48 """exception_handler method.""" 

49 raise NotImplementedError 

50 

51 

52class _DatabaseDialect(Enum): 

53 SQLITE = "sqlite" 

54 MYSQL = "mysql" 

55 POSTGRES = "postgres" 

56 

57 

58class _DatabaseErrorHandler(BaseErrorHandler[DBError]): 

59 """ 

60 Base for handlers that turn a SQLAlchemy error into an actionable HTTP response. 

61 

62 The failing statement is logged under a random lookup id and echoed back to the 

63 caller only when ``[api] expose_stacktrace`` is set; otherwise the response just 

64 points at that id in the api server logs. Subclasses set ``status_code`` and 

65 ``reason`` and may override ``_should_handle`` to skip exceptions they do not own. 

66 """ 

67 

68 status_code: int 

69 reason: str 

70 

71 def _should_handle(self, exc: DBError) -> bool: 

72 return True 

73 

74 def _raise_database_error_response(self, exc: DBError) -> None: 

75 statement = getattr(exc, "statement", "hidden") 

76 orig_error = getattr(exc, "orig", "hidden") 

77 exception_id = get_random_string() 

78 stacktrace = "".join(traceback.format_tb(exc.__traceback__)) 

79 log_message = f"Error with id {exception_id}, statement: {statement}\n{stacktrace}" 

80 log.error(log_message) 

81 

82 if conf.get("api", "expose_stacktrace") == "True": 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true

83 message = log_message 

84 statement_out = str(statement) 

85 orig_error_out = str(orig_error) 

86 else: 

87 message = ( 

88 "Serious error when handling your request. Check logs for more details - " 

89 f"you will find it in api server when you look for ID {exception_id}" 

90 ) 

91 statement_out = "hidden" 

92 orig_error_out = "hidden" 

93 

94 raise HTTPException( 

95 status_code=self.status_code, 

96 detail={ 

97 "reason": self.reason, 

98 "statement": statement_out, 

99 "orig_error": orig_error_out, 

100 "message": message, 

101 }, 

102 ) 

103 

104 def exception_handler(self, request: Request, exc: DBError): 

105 if not self._should_handle(exc): 

106 return 

107 self._raise_database_error_response(exc) 

108 

109 

110class _UniqueConstraintErrorHandler(_DatabaseErrorHandler[IntegrityError]): 

111 """Translate a unique-constraint ``IntegrityError`` into a 409, matched per database dialect.""" 

112 

113 status_code = status.HTTP_409_CONFLICT 

114 reason = "Unique constraint violation" 

115 

116 unique_constraint_error_prefix_dict: dict[_DatabaseDialect, str] = { 

117 _DatabaseDialect.SQLITE: "UNIQUE constraint failed", 

118 _DatabaseDialect.MYSQL: "Duplicate entry", 

119 _DatabaseDialect.POSTGRES: "violates unique constraint", 

120 } 

121 

122 def __init__(self): 

123 super().__init__(IntegrityError) 

124 self.dialect: _DatabaseDialect | None = None 

125 

126 def _should_handle(self, exc: IntegrityError) -> bool: 

127 return self._is_dialect_matched(exc) 

128 

129 def _is_dialect_matched(self, exc: IntegrityError) -> bool: 

130 """Check if the exception matches the unique constraint error message for any dialect.""" 

131 exc_orig_str = str(exc.orig) 

132 for dialect, error_msg in self.unique_constraint_error_prefix_dict.items(): 

133 if error_msg in exc_orig_str: 

134 self.dialect = dialect 

135 return True 

136 return False 

137 

138 

139class DataErrorHandler(_DatabaseErrorHandler[DataError]): 

140 """ 

141 Translate a ``sqlalchemy.exc.DataError`` into a 422. 

142 

143 The database rejected a value that passed Pydantic validation (too long, out of 

144 range, or the wrong type for its column), so it is a client error, not a 500. 

145 """ 

146 

147 status_code = HTTP_422_UNPROCESSABLE_CONTENT 

148 reason = "Value rejected by database" 

149 

150 def __init__(self): 

151 super().__init__(DataError) 

152 

153 

154class DagErrorHandler(BaseErrorHandler[DeserializationError]): 

155 """Handler for Dag related errors.""" 

156 

157 def __init__(self): 

158 super().__init__(DeserializationError) 

159 

160 def exception_handler(self, request: Request, exc: DeserializationError): 

161 """Handle Dag deserialization exceptions.""" 

162 raise HTTPException( 

163 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

164 detail=f"An error occurred while trying to deserialize Dag: {exc}", 

165 ) 

166 

167 

168class SQLAlchemyErrorHandler(_DatabaseErrorHandler[SQLAlchemyError]): 

169 """Generic handler for SQLAlchemyError -> 500 responses.""" 

170 

171 status_code = status.HTTP_500_INTERNAL_SERVER_ERROR 

172 reason = "Database error" 

173 

174 def __init__(self): 

175 super().__init__(SQLAlchemyError) 

176 

177 

178ERROR_HANDLERS: list[BaseErrorHandler] = [ 

179 _UniqueConstraintErrorHandler(), 

180 DataErrorHandler(), 

181 SQLAlchemyErrorHandler(), 

182 DagErrorHandler(), 

183] 

184 

185 

186def init_error_handlers(app: FastAPI) -> None: 

187 for handler in ERROR_HANDLERS: 

188 app.add_exception_handler(handler.exception_cls, handler.exception_handler)