feat(ingest/iceberg): Improve iceberg connector (#12163)

Author: skrydal

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

@@ -10,6 +10,7 @@
     NoSuchNamespaceError,
     NoSuchPropertyException,
     NoSuchTableError,
+    ServerError,
 )
 from pyiceberg.schema import Schema, SchemaVisitorPerPrimitiveType, visit
 from pyiceberg.table import Table
@@ -145,6 +146,13 @@ def _get_datasets(self, catalog: Catalog) -> Iterable[Identifier]:
         self.report.report_no_listed_namespaces(len(namespaces))
         tables_count = 0
         for namespace in namespaces:
+            namespace_repr = ".".join(namespace)
+            if not self.config.namespace_pattern.allowed(namespace_repr):
+                LOGGER.info(
+                    f"Namespace {namespace_repr} is not allowed by config pattern, skipping"
+                )
+                self.report.report_dropped(f"{namespace_repr}.*")
+                continue
             try:
                 tables = catalog.list_tables(namespace)
                 tables_count += len(tables)
@@ -181,6 +189,9 @@ def _process_dataset(dataset_path: Identifier) -> Iterable[MetadataWorkUnit]:
             if not self.config.table_pattern.allowed(dataset_name):
                 # Dataset name is rejected by pattern, report as dropped.
                 self.report.report_dropped(dataset_name)
+                LOGGER.debug(
+                    f"Skipping table {dataset_name} due to not being allowed by the config pattern"
+                )
                 return
             try:
                 if not hasattr(thread_local, "local_catalog"):
@@ -219,6 +230,22 @@ def _process_dataset(dataset_path: Identifier) -> Iterable[MetadataWorkUnit]:
                 LOGGER.warning(
                     f"NoSuchTableError while processing table {dataset_path}, skipping it.",
                 )
+            except FileNotFoundError as e:
+                self.report.report_warning(
+                    "file-not-found",
+                    f"Encountered FileNotFoundError when trying to read manifest file for {dataset_name}. {e}",
+                )
+                LOGGER.warning(
+                    f"FileNotFoundError while processing table {dataset_path}, skipping it."
+                )
+            except ServerError as e:
+                self.report.report_warning(
+                    "iceberg-rest-server-error",
+                    f"Iceberg Rest Catalog returned 500 status due to an unhandled exception for {dataset_name}. Exception: {e}",
+                )
+                LOGGER.warning(
+                    f"Iceberg Rest Catalog server error (500 status) encountered when processing table {dataset_path}, skipping it."
+                )
             except Exception as e:
                 self.report.report_failure("general", f"Failed to create workunit: {e}")
                 LOGGER.exception(
@@ -269,7 +296,6 @@ def _create_iceberg_workunit(
                 ] = table.current_snapshot().manifest_list
             dataset_properties = DatasetPropertiesClass(
                 name=table.name()[-1],
-                tags=[],
                 description=table.metadata.properties.get("comment", None),
                 customProperties=custom_properties,
             )

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

@@ -68,6 +68,10 @@ class IcebergSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin)
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for tables to filter in ingestion.",
     )
+    namespace_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for namespaces to filter in ingestion.",
+    )
     user_ownership_property: Optional[str] = Field(
         default="owner",
         description="Iceberg table property to look for a `CorpUser` owner.  Can only hold a single user value.  If property has no value, no owner information will be emitted.",

metadata-ingestion/tests/unit/test_iceberg.py

@@ -10,6 +10,8 @@
     NoSuchIcebergTableError,
     NoSuchNamespaceError,
     NoSuchPropertyException,
+    NoSuchTableError,
+    ServerError,
 )
 from pyiceberg.io.pyarrow import PyArrowFileIO
 from pyiceberg.partitioning import PartitionSpec
@@ -39,6 +41,7 @@
     UUIDType,
 )
 
+from datahub.configuration.common import AllowDenyPattern
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.iceberg.iceberg import (
@@ -62,12 +65,12 @@
 )
 
 
-def with_iceberg_source(processing_threads: int = 1) -> IcebergSource:
+def with_iceberg_source(processing_threads: int = 1, **kwargs: Any) -> IcebergSource:
     catalog = {"test": {"type": "rest"}}
     return IcebergSource(
         ctx=PipelineContext(run_id="iceberg-source-test"),
         config=IcebergSourceConfig(
-            catalog=catalog, processing_threads=processing_threads
+            catalog=catalog, processing_threads=processing_threads, **kwargs
         ),
     )
 
@@ -542,27 +545,27 @@ def __init__(self, tables: Dict[str, Dict[str, Callable[[], Table]]]):
         """
         self.tables = tables
 
-    def list_namespaces(self) -> Iterable[str]:
-        return [*self.tables.keys()]
+    def list_namespaces(self) -> Iterable[Tuple[str]]:
+        return [*[(key,) for key in self.tables.keys()]]
 
     def list_tables(self, namespace: str) -> Iterable[Tuple[str, str]]:
-        return [(namespace, table) for table in self.tables[namespace].keys()]
+        return [(namespace[0], table) for table in self.tables[namespace[0]].keys()]
 
     def load_table(self, dataset_path: Tuple[str, str]) -> Table:
         return self.tables[dataset_path[0]][dataset_path[1]]()
 
 
 class MockCatalogExceptionListingTables(MockCatalog):
     def list_tables(self, namespace: str) -> Iterable[Tuple[str, str]]:
-        if namespace == "no_such_namespace":
+        if namespace == ("no_such_namespace",):
             raise NoSuchNamespaceError()
-        if namespace == "generic_exception":
+        if namespace == ("generic_exception",):
             raise Exception()
         return super().list_tables(namespace)
 
 
 class MockCatalogExceptionListingNamespaces(MockCatalog):
-    def list_namespaces(self) -> Iterable[str]:
+    def list_namespaces(self) -> Iterable[Tuple[str]]:
         raise Exception()
 
 
@@ -814,15 +817,157 @@ def test_proper_run_with_multiple_namespaces() -> None:
         )
 
 
+def test_filtering() -> None:
+    source = with_iceberg_source(
+        processing_threads=1,
+        table_pattern=AllowDenyPattern(deny=[".*abcd.*"]),
+        namespace_pattern=AllowDenyPattern(allow=["namespace1"]),
+    )
+    mock_catalog = MockCatalog(
+        {
+            "namespace1": {
+                "table_xyz": lambda: Table(
+                    identifier=("namespace1", "table_xyz"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace1/table_xyz",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace1/table_xyz",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+                "JKLtable": lambda: Table(
+                    identifier=("namespace1", "JKLtable"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace1/JKLtable",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace1/JKLtable",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+                "table_abcd": lambda: Table(
+                    identifier=("namespace1", "table_abcd"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace1/table_abcd",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace1/table_abcd",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+                "aaabcd": lambda: Table(
+                    identifier=("namespace1", "aaabcd"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace1/aaabcd",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace1/aaabcd",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+            },
+            "namespace2": {
+                "foo": lambda: Table(
+                    identifier=("namespace2", "foo"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace2/foo",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace2/foo",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+                "bar": lambda: Table(
+                    identifier=("namespace2", "bar"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace2/bar",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace2/bar",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+            },
+            "namespace3": {
+                "sales": lambda: Table(
+                    identifier=("namespace3", "sales"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace3/sales",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace3/sales",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+                "products": lambda: Table(
+                    identifier=("namespace2", "bar"),
+                    metadata=TableMetadataV2(
+                        partition_specs=[PartitionSpec(spec_id=0)],
+                        location="s3://abcdefg/namespace3/products",
+                        last_column_id=0,
+                        schemas=[Schema(schema_id=0)],
+                    ),
+                    metadata_location="s3://abcdefg/namespace3/products",
+                    io=PyArrowFileIO(),
+                    catalog=None,
+                ),
+            },
+        }
+    )
+    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()]
+        assert len(wu) == 2
+        urns = []
+        for unit in wu:
+            assert isinstance(unit.metadata, MetadataChangeEvent)
+            assert isinstance(unit.metadata.proposedSnapshot, DatasetSnapshotClass)
+            urns.append(unit.metadata.proposedSnapshot.urn)
+        TestCase().assertCountEqual(
+            urns,
+            [
+                "urn:li:dataset:(urn:li:dataPlatform:iceberg,namespace1.table_xyz,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:iceberg,namespace1.JKLtable,PROD)",
+            ],
+        )
+        assert source.report.tables_scanned == 2
+
+
 def test_handle_expected_exceptions() -> None:
     source = with_iceberg_source(processing_threads=3)
 
     def _raise_no_such_property_exception():
         raise NoSuchPropertyException()
 
-    def _raise_no_such_table_exception():
+    def _raise_no_such_iceberg_table_exception():
         raise NoSuchIcebergTableError()
 
+    def _raise_file_not_found_error():
+        raise FileNotFoundError()
+
+    def _raise_no_such_table_exception():
+        raise NoSuchTableError()
+
+    def _raise_server_error():
+        raise ServerError()
+
     mock_catalog = MockCatalog(
         {
             "namespaceA": {
@@ -876,6 +1021,9 @@ def _raise_no_such_table_exception():
                 ),
                 "table5": _raise_no_such_property_exception,
                 "table6": _raise_no_such_table_exception,
+                "table7": _raise_file_not_found_error,
+                "table8": _raise_no_such_iceberg_table_exception,
+                "table9": _raise_server_error,
             }
         }
     )
@@ -899,7 +1047,7 @@ def _raise_no_such_table_exception():
                 "urn:li:dataset:(urn:li:dataPlatform:iceberg,namespaceA.table4,PROD)",
             ],
         )
-        assert source.report.warnings.total_elements == 2
+        assert source.report.warnings.total_elements == 5
         assert source.report.failures.total_elements == 0
         assert source.report.tables_scanned == 4