acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/iceberg/iceberg.py

@@ -12,7 +12,7 @@ from pyiceberg.exceptions import (
     NoSuchNamespaceError,
     NoSuchPropertyException,
     NoSuchTableError,
-    ServerError,
+    RESTError,
 )
 from pyiceberg.schema import Schema, SchemaVisitorPerPrimitiveType, visit
 from pyiceberg.table import Table
@@ -118,7 +118,7 @@ logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
 
 
 @platform_name("Iceberg")
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @config_class(IcebergSourceConfig)
 @capability(
     SourceCapability.PLATFORM_INSTANCE,
@@ -134,7 +134,9 @@ logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
     SourceCapability.OWNERSHIP,
     "Automatically ingests ownership information from table properties based on `user_ownership_property` and `group_ownership_property`",
 )
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 class IcebergSource(StatefulIngestionSourceBase):
     """
     ## Integration Details
@@ -152,6 +154,10 @@ class IcebergSource(StatefulIngestionSourceBase):
         self.report: IcebergSourceReport = IcebergSourceReport()
         self.config: IcebergSourceConfig = config
         self.ctx: PipelineContext = ctx
+        self.stamping_processor = AutoSystemMetadata(
+            self.ctx
+        )  # single instance used only when processing namespaces
+        self.namespaces: List[Tuple[Identifier, str]] = []
 
     @classmethod
     def create(cls, config_dict: Dict, ctx: PipelineContext) -> "IcebergSource":
@@ -194,9 +200,9 @@ class IcebergSource(StatefulIngestionSourceBase):
             auto_lowercase_dataset_urns,
             auto_materialize_referenced_tags_terms,
             partial(
-                auto_fix_duplicate_schema_field_paths, platform=self._infer_platform()
+                auto_fix_duplicate_schema_field_paths, platform=self.infer_platform()
             ),
-            partial(auto_fix_empty_field_paths, platform=self._infer_platform()),
+            partial(auto_fix_empty_field_paths, platform=self.infer_platform()),
             partial(auto_workunit_reporter, self.get_report()),
             auto_patch_last_modified,
             EnsureAspectSizeProcessor(self.get_report()).ensure_aspect_size,
@@ -244,6 +250,13 @@ class IcebergSource(StatefulIngestionSourceBase):
                     context=str(namespace),
                     exc=e,
                 )
+            except RESTError as e:
+                self.report.warning(
+                    title="Iceberg REST Server Error",
+                    message="Iceberg REST Server returned error status when trying to list tables for a namespace, skipping it.",
+                    context=str(namespace),
+                    exc=e,
+                )
             except Exception as e:
                 self.report.report_failure(
                     title="Error when processing a namespace",
@@ -320,10 +333,10 @@ class IcebergSource(StatefulIngestionSourceBase):
                     context=dataset_name,
                     exc=e,
                 )
-            except ServerError as e:
+            except RESTError as e:
                 self.report.warning(
                     title="Iceberg REST Server Error",
-                    message="Iceberg returned 500 HTTP status when trying to process a table, skipping it.",
+                    message="Iceberg REST Server returned error status when trying to process a table, skipping it.",
                     context=dataset_name,
                     exc=e,
                 )
@@ -363,7 +376,7 @@ class IcebergSource(StatefulIngestionSourceBase):
                 )
 
         try:
-            catalog = self.config.get_catalog()
+            self.catalog = self.config.get_catalog()
         except Exception as e:
             self.report.report_failure(
                 title="Failed to initialize catalog object",
@@ -373,33 +386,7 @@ class IcebergSource(StatefulIngestionSourceBase):
             return
 
         try:
-            stamping_processor = AutoSystemMetadata(self.ctx)
-            namespace_ids = self._get_namespaces(catalog)
-            namespaces: List[Tuple[Identifier, str]] = []
-            for namespace in namespace_ids:
-                namespace_repr = ".".join(namespace)
-                LOGGER.debug(f"Processing namespace {namespace_repr}")
-                namespace_urn = make_container_urn(
-                    NamespaceKey(
-                        namespace=namespace_repr,
-                        platform=self.platform,
-                        instance=self.config.platform_instance,
-                        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, namespace_properties
-                ):
-                    yield stamping_processor.stamp_wu(
-                        MetadataChangeProposalWrapper(
-                            entityUrn=namespace_urn, aspect=aspect
-                        ).as_workunit()
-                    )
-            LOGGER.debug("Namespaces ingestion completed")
+            yield from self._process_namespaces()
         except Exception as e:
             self.report.report_failure(
                 title="Failed to list namespaces",
@@ -413,13 +400,70 @@ class IcebergSource(StatefulIngestionSourceBase):
             args_list=[
                 (dataset_path, namespace_urn)
                 for dataset_path, namespace_urn in self._get_datasets(
-                    catalog, namespaces
+                    self.catalog, self.namespaces
                 )
             ],
             max_workers=self.config.processing_threads,
         ):
             yield wu
 
+    def _try_processing_namespace(
+        self, namespace: Identifier
+    ) -> Iterable[MetadataWorkUnit]:
+        namespace_repr = ".".join(namespace)
+        try:
+            LOGGER.debug(f"Processing namespace {namespace_repr}")
+            namespace_urn = make_container_urn(
+                NamespaceKey(
+                    namespace=namespace_repr,
+                    platform=self.platform,
+                    instance=self.config.platform_instance,
+                    env=self.config.env,
+                )
+            )
+
+            namespace_properties: Properties = self.catalog.load_namespace_properties(
+                namespace
+            )
+            for aspect in self._create_iceberg_namespace_aspects(
+                namespace, namespace_properties
+            ):
+                yield self.stamping_processor.stamp_wu(
+                    MetadataChangeProposalWrapper(
+                        entityUrn=namespace_urn, aspect=aspect
+                    ).as_workunit()
+                )
+            self.namespaces.append((namespace, namespace_urn))
+        except NoSuchNamespaceError as e:
+            self.report.report_warning(
+                title="Failed to retrieve namespace properties",
+                message="Couldn't find the namespace, was it deleted during the ingestion?",
+                context=namespace_repr,
+                exc=e,
+            )
+            return
+        except RESTError as e:
+            self.report.warning(
+                title="Iceberg REST Server Error",
+                message="Iceberg REST Server returned error status when trying to retrieve namespace properties, skipping it.",
+                context=str(namespace),
+                exc=e,
+            )
+        except Exception as e:
+            self.report.report_failure(
+                title="Failed to process namespace",
+                message="Unhandled exception happened during processing of the namespace",
+                context=namespace_repr,
+                exc=e,
+            )
+
+    def _process_namespaces(self) -> Iterable[MetadataWorkUnit]:
+        namespace_ids = self._get_namespaces(self.catalog)
+        for namespace in namespace_ids:
+            yield from self._try_processing_namespace(namespace)
+
+        LOGGER.debug("Namespaces ingestion completed")
+
     def _create_iceberg_table_aspects(
         self, dataset_name: str, table: Table, namespace_urn: str
     ) -> Iterable[_Aspect]:
@@ -522,11 +566,11 @@ class IcebergSource(StatefulIngestionSourceBase):
         custom_properties["format-version"] = str(table.metadata.format_version)
         custom_properties["partition-spec"] = str(self._get_partition_aspect(table))
         last_modified: Optional[int] = table.metadata.last_updated_ms
-        if table.current_snapshot():
-            custom_properties["snapshot-id"] = str(table.current_snapshot().snapshot_id)
-            custom_properties["manifest-list"] = table.current_snapshot().manifest_list
+        if current_snapshot := table.current_snapshot():
+            custom_properties["snapshot-id"] = str(current_snapshot.snapshot_id)
+            custom_properties["manifest-list"] = current_snapshot.manifest_list
             if not last_modified:
-                last_modified = int(table.current_snapshot().timestamp_ms)
+                last_modified = int(current_snapshot.timestamp_ms)
         if "created-at" in custom_properties:
             try:
                 dt = dateutil_parser.isoparse(custom_properties["created-at"])
@@ -792,9 +836,6 @@ class ToAvroSchemaIcebergVisitor(SchemaVisitorPerPrimitiveType[Dict[str, Any]]):
             "native_data_type": str(timestamp_type),
         }
 
-    # visit_timestamptz() is required when using pyiceberg >= 0.5.0, which is essentially a duplicate
-    # of visit_timestampz().  The function has been renamed from visit_timestampz().
-    # Once Datahub can upgrade its pyiceberg dependency to >=0.5.0, the visit_timestampz() function can be safely removed.
     def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> Dict[str, Any]:
         # Avro supports 2 types of timestamp:
         #  - Timestamp: independent of a particular timezone or calendar (TZ information is lost)
@@ -811,22 +852,6 @@ class ToAvroSchemaIcebergVisitor(SchemaVisitorPerPrimitiveType[Dict[str, Any]]):
             "native_data_type": str(timestamptz_type),
         }
 
-    def visit_timestampz(self, timestamptz_type: TimestamptzType) -> Dict[str, Any]:
-        # Avro supports 2 types of timestamp:
-        #  - Timestamp: independent of a particular timezone or calendar (TZ information is lost)
-        #  - Local Timestamp: represents a timestamp in a local timezone, regardless of what specific time zone is considered local
-        # utcAdjustment: bool = True
-        return {
-            "type": "long",
-            "logicalType": "timestamp-micros",
-            # Commented out since Avro's Python implementation (1.11.0) does not support local-timestamp-micros, even though it exists in the spec.
-            # See bug report: https://issues.apache.org/jira/browse/AVRO-3476 and PR https://github.com/apache/avro/pull/1634
-            # "logicalType": "timestamp-micros"
-            # if timestamp_type.adjust_to_utc
-            # else "local-timestamp-micros",
-            "native_data_type": str(timestamptz_type),
-        }
-
     def visit_string(self, string_type: StringType) -> Dict[str, Any]:
         return {
             "type": "string",
@@ -845,3 +870,42 @@ class ToAvroSchemaIcebergVisitor(SchemaVisitorPerPrimitiveType[Dict[str, Any]]):
             "type": "bytes",
             "native_data_type": str(binary_type),
         }
+
+    def visit_timestamp_ns(self, timestamp_ns_type: Any) -> Dict[str, Any]:
+        # Handle nanosecond precision timestamps
+        # Avro supports 2 types of timestamp:
+        #  - Timestamp: independent of a particular timezone or calendar (TZ information is lost)
+        #  - Local Timestamp: represents a timestamp in a local timezone, regardless of what specific time zone is considered local
+        return {
+            "type": "long",
+            "logicalType": "timestamp-micros",
+            # Commented out since Avro's Python implementation (1.11.0) does not support local-timestamp-micros, even though it exists in the spec.
+            # See bug report: https://issues.apache.org/jira/browse/AVRO-3476 and PR https://github.com/apache/avro/pull/1634
+            # "logicalType": "timestamp-micros"
+            # if timestamp_ns_type.adjust_to_utc
+            # else "local-timestamp-micros",
+            "native_data_type": str(timestamp_ns_type),
+        }
+
+    def visit_timestamptz_ns(self, timestamptz_ns_type: Any) -> Dict[str, Any]:
+        # Handle nanosecond precision timestamps with timezone
+        # Avro supports 2 types of timestamp:
+        #  - Timestamp: independent of a particular timezone or calendar (TZ information is lost)
+        #  - Local Timestamp: represents a timestamp in a local timezone, regardless of what specific time zone is considered local
+        return {
+            "type": "long",
+            "logicalType": "timestamp-micros",
+            # Commented out since Avro's Python implementation (1.11.0) does not support local-timestamp-micros, even though it exists in the spec.
+            # See bug report: https://issues.apache.org/jira/browse/AVRO-3476 and PR https://github.com/apache/avro/pull/1634
+            # "logicalType": "timestamp-micros"
+            # if timestamptz_ns_type.adjust_to_utc
+            # else "local-timestamp-micros",
+            "native_data_type": str(timestamptz_ns_type),
+        }
+
+    def visit_unknown(self, unknown_type: Any) -> Dict[str, Any]:
+        # Handle unknown types
+        return {
+            "type": "string",
+            "native_data_type": str(unknown_type),
+        }

datahub/ingestion/source/iceberg/iceberg_profiler.py

@@ -12,6 +12,7 @@ from pyiceberg.types import (
     IcebergType,
     IntegerType,
     LongType,
+    PrimitiveType,
     TimestampType,
     TimestamptzType,
     TimeType,
@@ -22,6 +23,7 @@ from pyiceberg.utils.datetime import (
     to_human_timestamp,
     to_human_timestamptz,
 )
+from typing_extensions import TypeGuard
 
 from datahub.emitter.mce_builder import get_sys_time
 from datahub.ingestion.source.iceberg.iceberg_common import (
@@ -65,7 +67,7 @@ class IcebergProfiler:
         aggregated_values: Dict[int, Any],
         manifest_values: Dict[int, bytes],
     ) -> None:
-        for field_id, value_encoded in manifest_values.items():  # type: int, Any
+        for field_id, value_encoded in manifest_values.items():
             try:
                 field = schema.find_field(field_id)
             except ValueError:
@@ -240,7 +242,7 @@ class IcebergProfiler:
             return None
 
     @staticmethod
-    def _is_numeric_type(type: IcebergType) -> bool:
+    def _is_numeric_type(type: IcebergType) -> TypeGuard[PrimitiveType]:
         return isinstance(
             type,
             (

datahub/ingestion/source/identity/azure_ad.py

@@ -167,7 +167,7 @@ class AzureADSourceReport(StaleEntityRemovalSourceReport):
 @config_class(AzureADConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 class AzureADSource(StatefulIngestionSourceBase):
     """

datahub/ingestion/source/identity/okta.py

@@ -41,7 +41,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
 )
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     CorpGroupInfoClass,
     CorpUserInfoClass,
     GroupMembershipClass,
@@ -202,7 +201,7 @@ class OktaSourceReport(StaleEntityRemovalSourceReport):
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.DESCRIPTIONS, "Optionally enabled via configuration")
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 class OktaSource(StatefulIngestionSourceBase):
     """
@@ -332,18 +331,12 @@ class OktaSource(StatefulIngestionSourceBase):
                 yield MetadataWorkUnit(id=wu_id, mce=mce)
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpGroup",
                     entityUrn=datahub_corp_group_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="origin",
                     aspect=OriginClass(OriginTypeClass.EXTERNAL, "OKTA"),
                 ).as_workunit()
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpGroup",
                     entityUrn=datahub_corp_group_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="status",
                     aspect=StatusClass(removed=False),
                 ).as_workunit()
 
@@ -418,18 +411,12 @@ class OktaSource(StatefulIngestionSourceBase):
                 yield MetadataWorkUnit(id=wu_id, mce=mce)
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpuser",
                     entityUrn=datahub_corp_user_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="origin",
                     aspect=OriginClass(OriginTypeClass.EXTERNAL, "OKTA"),
                 ).as_workunit()
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpuser",
                     entityUrn=datahub_corp_user_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="status",
                     aspect=StatusClass(removed=False),
                 ).as_workunit()
 

datahub/ingestion/source/kafka/kafka.py

@@ -189,6 +189,22 @@ class KafkaConnectionTest:
     SourceCapability.SCHEMA_METADATA,
     "Schemas associated with each topic are extracted from the schema registry. Avro and Protobuf (certified), JSON (incubating). Schema references are supported.",
 )
+@capability(
+    SourceCapability.DATA_PROFILING,
+    "Not supported",
+    supported=False,
+)
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Not supported. If you use Kafka Connect, the kafka-connect source can generate lineage.",
+    supported=False,
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Not supported",
+    supported=False,
+)
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class KafkaSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following:

datahub/ingestion/source/kafka_connect/common.py

@@ -4,7 +4,7 @@ from typing import Dict, Iterable, List, Optional
 
 from pydantic.fields import Field
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel, LaxStr
 from datahub.configuration.source_common import (
     DatasetLineageProviderConfigBase,
     PlatformInstanceConfigMixin,
@@ -29,7 +29,7 @@ CONNECTOR_CLASS = "connector.class"
 class ProvidedConfig(ConfigModel):
     provider: str
     path_key: str
-    value: str
+    value: LaxStr
 
 
 class GenericConnectorConfig(ConfigModel):