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

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 

19import json 

20from collections.abc import Iterable 

21from datetime import datetime 

22from typing import Any 

23 

24from pydantic import AliasPath, Field, field_validator 

25 

26from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel 

27 

28 

29class XComResponse(BaseModel): 

30 """Serializer for a xcom item.""" 

31 

32 key: str 

33 timestamp: datetime 

34 logical_date: datetime | None 

35 map_index: int 

36 task_id: str 

37 dag_id: str 

38 run_id: str 

39 dag_display_name: str = Field(validation_alias=AliasPath("dag_run", "dag_model", "dag_display_name")) 

40 task_display_name: str = Field(validation_alias=AliasPath("task", "task_display_name")) 

41 run_after: datetime = Field(validation_alias=AliasPath("dag_run", "run_after")) 

42 

43 

44def _stringify_if_needed(value): 

45 """ 

46 Check whether value is JSON-encodable (recursively if needed); stringify it if not. 

47 

48 The list of JSON-ecodable types are taken from Python documentation: 

49 https://docs.python.org/3/library/json.html#json.JSONEncoder 

50 """ 

51 if value is None or isinstance(value, (str, int, float, bool)): 

52 return value 

53 if isinstance(value, dict): 

54 return {str(k): _stringify_if_needed(v) for k, v in value.items()} 

55 if isinstance(value, (list, tuple)): 55 ↛ 57line 55 didn't jump to line 57 because the condition on line 55 was always true

56 return [_stringify_if_needed(v) for v in value] 

57 return str(value) 

58 

59 

60class XComResponseNative(XComResponse): 

61 """XCom response serializer with native return type.""" 

62 

63 value: Any 

64 

65 @field_validator("value", mode="before") 

66 def value_to_json_serializable(cls, v): 

67 return _stringify_if_needed(v) 

68 

69 

70class XComResponseString(XComResponse): 

71 """XCom response serializer with string return type.""" 

72 

73 value: str | None 

74 

75 @field_validator("value", mode="before") 

76 def value_to_string(cls, v): 

77 return str(v) if v is not None else None 

78 

79 

80class XComCollectionResponse(BaseModel): 

81 """XCom Collection serializer for responses.""" 

82 

83 xcom_entries: Iterable[XComResponse] 

84 total_entries: int 

85 

86 

87def _check_forbidden_xcom_keys(value: Any) -> Any: 

88 """Recursively reject forbidden deserialization keys in user-provided XCom data.""" 

89 from airflow._shared.serialization import FORBIDDEN_XCOM_KEYS 

90 

91 def _walk(obj: Any, path: str = "value") -> None: 

92 if isinstance(obj, str): 

93 # A value submitted as a JSON string literal (e.g. ``json.dumps({...})``) 

94 # is stored verbatim and re-parsed into a dict/list on a 

95 # ``deserialize=true`` read, which would otherwise smuggle reserved keys 

96 # past the dict/list checks below. Re-parse and inspect the decoded 

97 # structure the same way the read path does. 

98 try: 

99 decoded = json.loads(obj) 

100 except (ValueError, TypeError): 

101 return 

102 if isinstance(decoded, (dict, list)): 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true

103 _walk(decoded, path) 

104 return 

105 if isinstance(obj, dict): 

106 found = FORBIDDEN_XCOM_KEYS & obj.keys() 

107 if found: 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true

108 raise ValueError( 

109 f"XCom {path} contains reserved serialization keys: {', '.join(sorted(found))}. " 

110 f"These keys are reserved for internal use." 

111 ) 

112 for k, v in obj.items(): 

113 _walk(v, f"{path}.{k}") 

114 elif isinstance(obj, (list, tuple)): 

115 for i, item in enumerate(obj): 

116 _walk(item, f"{path}[{i}]") 

117 

118 _walk(value) 

119 return value 

120 

121 

122class XComCreateBody(StrictBaseModel): 

123 """Payload serializer for creating an XCom entry.""" 

124 

125 key: str 

126 value: Any 

127 map_index: int = -1 

128 

129 @field_validator("value") 

130 @classmethod 

131 def _check_forbidden_keys(cls, value: Any) -> Any: 

132 return _check_forbidden_xcom_keys(value) 

133 

134 

135class XComUpdateBody(StrictBaseModel): 

136 """Payload serializer for updating an XCom entry.""" 

137 

138 value: Any 

139 map_index: int = -1 

140 

141 @field_validator("value") 

142 @classmethod 

143 def _check_forbidden_keys(cls, value: Any) -> Any: 

144 return _check_forbidden_xcom_keys(value)