1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
19
20import json
21from datetime import date, datetime
22from decimal import Decimal
23from typing import Any
24
25from flask.json.provider import JSONProvider
26
27from airflow.serialization.serde import CLASSNAME, DATA, SCHEMA_ID, deserialize, serialize
28from airflow.utils.timezone import convert_to_utc, is_naive
29
30
31class AirflowJsonProvider(JSONProvider):
32 """JSON Provider for Flask app to use WebEncoder."""
33
34 ensure_ascii: bool = True
35 sort_keys: bool = True
36
37 def dumps(self, obj, **kwargs):
38 kwargs.setdefault("ensure_ascii", self.ensure_ascii)
39 kwargs.setdefault("sort_keys", self.sort_keys)
40 return json.dumps(obj, **kwargs, cls=WebEncoder)
41
42 def loads(self, s: str | bytes, **kwargs):
43 return json.loads(s, **kwargs)
44
45
46class WebEncoder(json.JSONEncoder):
47 """This encodes values into a web understandable format. There is no deserializer.
48
49 This parses datetime, dates, Decimal and bytes. In order to parse the custom
50 classes and the other types, and since it's just to show the result in the UI,
51 we return repr(object) for everything else.
52 """
53
54 def default(self, o: Any) -> Any:
55 if isinstance(o, datetime):
56 if is_naive(o):
57 o = convert_to_utc(o)
58 return o.isoformat()
59
60 if isinstance(o, date):
61 return o.strftime("%Y-%m-%d")
62
63 if isinstance(o, Decimal):
64 data = serialize(o)
65 if isinstance(data, dict) and DATA in data:
66 return data[DATA]
67 if isinstance(o, bytes):
68 try:
69 return o.decode("unicode_escape")
70 except UnicodeDecodeError:
71 return repr(o)
72 try:
73 data = serialize(o)
74 if isinstance(data, dict) and CLASSNAME in data:
75 # this is here for backwards compatibility
76 if (
77 data[CLASSNAME].startswith("numpy")
78 or data[CLASSNAME] == "kubernetes.client.models.v1_pod.V1Pod"
79 ):
80 return data[DATA]
81 return data
82 except TypeError:
83 return repr(o)
84
85
86class XComEncoder(json.JSONEncoder):
87 """This encoder serializes any object that has attr, dataclass or a custom serializer."""
88
89 def default(self, o: object) -> Any:
90 try:
91 return serialize(o)
92 except TypeError:
93 return super().default(o)
94
95 def encode(self, o: Any) -> str:
96 # checked here and in serialize
97 if isinstance(o, dict) and (CLASSNAME in o or SCHEMA_ID in o):
98 raise AttributeError(f"reserved key {CLASSNAME} found in dict to serialize")
99
100 # tuples are not preserved by std python serializer
101 if isinstance(o, tuple):
102 o = self.default(o)
103
104 return super().encode(o)
105
106
107class XComDecoder(json.JSONDecoder):
108 """Deserialize dicts to objects if they contain the `__classname__` key, otherwise return the dict."""
109
110 def __init__(self, *args, **kwargs) -> None:
111 if not kwargs.get("object_hook"):
112 kwargs["object_hook"] = self.object_hook
113
114 super().__init__(*args, **kwargs)
115
116 def object_hook(self, dct: dict) -> object:
117 return deserialize(dct)
118
119 @staticmethod
120 def orm_object_hook(dct: dict) -> object:
121 """Create a readable representation of a serialized object."""
122 return deserialize(dct, False)
123
124
125# backwards compatibility
126AirflowJsonEncoder = WebEncoder