feat(ingestion/iceberg): Add capability to extract namespace properties to the iceberg ingestor (#13238)

Author: skrydal

metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py

@@ -16,7 +16,7 @@
 )
 from pyiceberg.schema import Schema, SchemaVisitorPerPrimitiveType, visit
 from pyiceberg.table import Table
-from pyiceberg.typedef import Identifier
+from pyiceberg.typedef import Identifier, Properties
 from pyiceberg.types import (
     BinaryType,
     BooleanType,
@@ -387,8 +387,13 @@ def _process_dataset(
                         env=self.config.env,
                     )
                 )
+                namespace_properties: Properties = catalog.load_namespace_properties(
+                    namespace
+                )
                 namespaces.append((namespace, namespace_urn))
-                for aspect in self._create_iceberg_namespace_aspects(namespace):
+                for aspect in self._create_iceberg_namespace_aspects(
+                    namespace, namespace_properties
+                ):
                     yield stamping_processor.stamp_wu(
                         MetadataChangeProposalWrapper(
                             entityUrn=namespace_urn, aspect=aspect
@@ -608,12 +613,23 @@ def get_report(self) -> SourceReport:
         return self.report
 
     def _create_iceberg_namespace_aspects(
-        self, namespace: Identifier
+        self, namespace: Identifier, properties: Properties
     ) -> Iterable[_Aspect]:
         namespace_repr = ".".join(namespace)
+        custom_properties: Dict[str, str] = {}
+        for k, v in properties.items():
+            try:
+                custom_properties[str(k)] = str(v)
+            except Exception as e:
+                LOGGER.warning(
+                    f"Exception when trying to parse namespace properties for {namespace_repr}. Exception: {e}"
+                )
         yield Status(removed=False)
         yield ContainerProperties(
-            name=namespace_repr, qualifiedName=namespace_repr, env=self.config.env
+            name=namespace_repr,
+            qualifiedName=namespace_repr,
+            env=self.config.env,
+            customProperties=custom_properties,
         )
         yield SubTypes(typeNames=[DatasetContainerSubTypes.NAMESPACE])
         dpi = self._get_dataplatform_instance_aspect()

metadata-ingestion/tests/integration/iceberg/iceberg_deleted_table_mcps_golden.json

@@ -23,7 +23,9 @@
     "aspectName": "containerProperties",
     "aspect": {
         "json": {
-            "customProperties": {},
+            "customProperties": {
+                "location": "s3a://warehouse/wh/nyc"
+            },
             "name": "nyc",
             "qualifiedName": "nyc",
             "env": "PROD"

metadata-ingestion/tests/integration/iceberg/iceberg_ingest_mcps_golden.json

@@ -22,7 +22,9 @@
     "aspectName": "containerProperties",
     "aspect": {
         "json": {
-            "customProperties": {},
+            "customProperties": {
+                "location": "s3a://warehouse/wh/nyc"
+            },
             "name": "nyc",
             "qualifiedName": "nyc",
             "env": "PROD"

metadata-ingestion/tests/integration/iceberg/iceberg_profile_mcps_golden.json

@@ -22,7 +22,9 @@
     "aspectName": "containerProperties",
     "aspect": {
         "json": {
-            "customProperties": {},
+            "customProperties": {
+                "location": "s3a://warehouse/wh/nyc"
+            },
             "name": "nyc",
             "qualifiedName": "nyc",
             "env": "PROD"

metadata-ingestion/tests/unit/test_iceberg.py

@@ -1,4 +1,5 @@
 import uuid
+from collections import defaultdict
 from decimal import Decimal
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 from unittest import TestCase
@@ -540,13 +541,20 @@ def test_avro_decimal_bytes_nullable() -> None:
 
 
 class MockCatalog:
-    def __init__(self, tables: Dict[str, Dict[str, Callable[[], Table]]]):
+    def __init__(
+        self,
+        tables: Dict[str, Dict[str, Callable[[], Table]]],
+        namespace_properties: Optional[Dict[str, Dict[str, str]]] = None,
+    ):
         """
 
         :param tables: Dictionary containing namespaces as keys and dictionaries containing names of tables (keys) and
                        their metadata as values
         """
         self.tables = tables
+        self.namespace_properties = (
+            namespace_properties if namespace_properties else defaultdict(dict)
+        )
 
     def list_namespaces(self) -> Iterable[Tuple[str]]:
         return [*[(key,) for key in self.tables]]
@@ -557,6 +565,9 @@ def list_tables(self, namespace: str) -> Iterable[Tuple[str, str]]:
     def load_table(self, dataset_path: Tuple[str, str]) -> Table:
         return self.tables[dataset_path[0]][dataset_path[1]]()
 
+    def load_namespace_properties(self, namespace: Tuple[str, ...]) -> Dict[str, str]:
+        return self.namespace_properties[namespace[0]]
+
 
 class MockCatalogExceptionListingTables(MockCatalog):
     def list_tables(self, namespace: str) -> Iterable[Tuple[str, str]]:
@@ -1189,3 +1200,39 @@ def _raise_other_value_error_exception():
                 "('namespaceA', 'table5') <class 'Exception'>: ",
             ],
         )
+
+
+def test_ingesting_namespace_properties() -> None:
+    source = with_iceberg_source(processing_threads=2)
+    custom_properties = {"prop1": "foo", "prop2": "bar"}
+    mock_catalog = MockCatalog(
+        tables={
+            "namespaceA": {},  # mapped to urn:li:container:390e031441265aae5b7b7ae8d51b0c1f
+            "namespaceB": {},  # mapped to urn:li:container:74727446a56420d80ff3b1abf2a18087
+        },
+        namespace_properties={"namespaceA": {}, "namespaceB": custom_properties},
+    )
+    with patch(
+        "datahub.ingestion.source.iceberg.iceberg.IcebergSourceConfig.get_catalog"
+    ) as get_catalog:
+        get_catalog.return_value = mock_catalog
+        wu: List[MetadataWorkUnit] = [*source.get_workunits_internal()]
+        aspects: Dict[str, Dict[str, Any]] = defaultdict(dict)
+        for unit in wu:
+            assert isinstance(unit.metadata, MetadataChangeProposalWrapper)
+            mcpw: MetadataChangeProposalWrapper = unit.metadata
+            assert mcpw.entityUrn
+            assert mcpw.aspectName
+            aspects[mcpw.entityUrn][mcpw.aspectName] = mcpw.aspect
+        assert (
+            aspects["urn:li:container:390e031441265aae5b7b7ae8d51b0c1f"][
+                "containerProperties"
+            ].customProperties
+            == {}
+        )
+        assert (
+            aspects["urn:li:container:74727446a56420d80ff3b1abf2a18087"][
+                "containerProperties"
+            ].customProperties
+            == custom_properties
+        )