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
20from typing import TYPE_CHECKING
21
22from airflow.utils.module_loading import qualname
23
24serializers = ["deltalake.table.DeltaTable"]
25deserializers = serializers
26stringifiers = serializers
27
28if TYPE_CHECKING:
29 from airflow.serialization.serde import U
30
31__version__ = 1
32
33
34def serialize(o: object) -> tuple[U, str, int, bool]:
35 from deltalake.table import DeltaTable
36
37 if not isinstance(o, DeltaTable):
38 return "", "", 0, False
39
40 from airflow.models.crypto import get_fernet
41
42 # we encrypt the information here until we have as part of the
43 # storage options can have sensitive information
44 fernet = get_fernet()
45 properties: dict = {}
46 for k, v in o._storage_options.items() if o._storage_options else {}:
47 properties[k] = fernet.encrypt(v.encode("utf-8")).decode("utf-8")
48
49 data = {
50 "table_uri": o.table_uri,
51 "version": o.version(),
52 "storage_options": properties,
53 }
54
55 return data, qualname(o), __version__, True
56
57
58def deserialize(classname: str, version: int, data: dict):
59 from deltalake.table import DeltaTable
60
61 from airflow.models.crypto import get_fernet
62
63 if version > __version__:
64 raise TypeError("serialized version is newer than class version")
65
66 if classname == qualname(DeltaTable):
67 fernet = get_fernet()
68 properties = {}
69 for k, v in data["storage_options"].items():
70 properties[k] = fernet.decrypt(v.encode("utf-8")).decode("utf-8")
71
72 if len(properties) == 0:
73 storage_options = None
74 else:
75 storage_options = properties
76
77 return DeltaTable(data["table_uri"], version=data["version"], storage_options=storage_options)
78
79 raise TypeError(f"do not know how to deserialize {classname}")