Add support for serialization of iceberg tables

airflow/serialization/serializers/iceberg.py

@@ -0,0 +1,76 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.utils.module_loading import qualname
+
+serializers = ["pyiceberg.table.Table"]
+deserializers = serializers
+stringifiers = serializers
+
+if TYPE_CHECKING:
+    from airflow.serialization.serde import U
+
+__version__ = 1
+
+
+def serialize(o: object) -> tuple[U, str, int, bool]:
+    from pyiceberg.table import Table
+
+    if not isinstance(o, Table):
+        return "", "", 0, False
+
+    from airflow.models.crypto import get_fernet
+
+    # we encrypt the catalog information here until we have
+    # global catalog management in airflow and the properties
+    # can have sensitive information
+    fernet = get_fernet()
+    properties = {}
+    for k, v in o.catalog.properties.items():
+        properties[k] = fernet.encrypt(v.encode("utf-8")).decode("utf-8")
+
+    data = {
+        "identifier": o.identifier,
+        "catalog_properties": properties,
+    }
+
+    return data, qualname(o), __version__, True
+
+
+def deserialize(classname: str, version: int, data: dict):
+    from pyiceberg.catalog import load_catalog
+    from pyiceberg.table import Table
+
+    from airflow.models.crypto import get_fernet
+
+    if version > __version__:
+        raise TypeError("serialized version is newer than class version")
+
+    if classname == qualname(Table):
+        fernet = get_fernet()
+        properties = {}
+        for k, v in data["catalog_properties"].items():
+            properties[k] = fernet.decrypt(v.encode("utf-8")).decode("utf-8")
+
+        catalog = load_catalog(data["identifier"][0], **properties)
+        return catalog.load_table((data["identifier"][1], data["identifier"][2]))
+
+    raise TypeError(f"do not know how to deserialize {classname}")

setup.py

@@ -453,6 +453,10 @@ def write_version(filename: str = str(AIRFLOW_SOURCES_ROOT / "airflow" / "git_ve
     "mongomock",
 ]
 
+_devel_only_iceberg = [
+    "pyiceberg>=0.5.0",
+]
+
 _devel_only_sentry = [
     "blinker",
 ]
@@ -492,6 +496,7 @@ def write_version(filename: str = str(AIRFLOW_SOURCES_ROOT / "airflow" / "git_ve
     *_devel_only_devscripts,
     *_devel_only_duckdb,
     *_devel_only_mongo,
+    *_devel_only_iceberg,
     *_devel_only_sentry,
     *_devel_only_static_checks,
     *_devel_only_tests,

tests/serialization/serializers/test_serializers.py

@@ -18,12 +18,16 @@
 
 import datetime
 import decimal
+from unittest.mock import patch
 
 import numpy as np
 import pendulum.tz
 import pytest
 from dateutil.tz import tzutc
 from pendulum import DateTime
+from pyiceberg.catalog import Catalog
+from pyiceberg.io import FileIO
+from pyiceberg.table import Table
 
 from airflow import PY39
 from airflow.models.param import Param, ParamsDict
@@ -175,3 +179,22 @@ def test_pandas(self):
         e = serialize(i)
         d = deserialize(e)
         assert i.equals(d)
+
+    @patch.object(Catalog, "__abstractmethods__", set())
+    @patch.object(FileIO, "__abstractmethods__", set())
+    @patch("pyiceberg.catalog.Catalog.load_table")
+    @patch("pyiceberg.catalog.load_catalog")
+    def test_iceberg(self, mock_load_catalog, mock_load_table):
+        uri = "http://rest.no.where"
+        catalog = Catalog("catalog", uri=uri)
+        identifier = ("catalog", "schema", "table")
+        mock_load_catalog.return_value = catalog
+
+        i = Table(identifier, "bar", catalog=catalog, metadata_location="", io=FileIO())
+        mock_load_table.return_value = i
+
+        e = serialize(i)
+        d = deserialize(e)
+        assert i == d
+        mock_load_catalog.assert_called_with("catalog", uri=uri)
+        mock_load_table.assert_called_with((identifier[1], identifier[2]))