acryl-datahub 1.1.0.5rc3__py3-none-any.whl → 1.1.0.5rc5__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -663,6 +663,7 @@ class SnowflakeQueriesSource(Source):
     def close(self) -> None:
         self.connection.close()
         self.queries_extractor.close()
+        super().close()
 
 
 # Make sure we don't try to generate too much info for a single query.

datahub/ingestion/source/sql/athena.py

@@ -29,7 +29,10 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import make_s3_urn
-from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetContainerSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
@@ -321,9 +324,18 @@ class Partitionitem:
 @capability(
     SourceCapability.DATA_PROFILING,
     "Optionally enabled via configuration. Profiling uses sql queries on whole table which can be expensive operation.",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
 )
-@capability(SourceCapability.LINEAGE_COARSE, "Supported for S3 tables")
-@capability(SourceCapability.LINEAGE_FINE, "Supported for S3 tables")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 class AthenaSource(SQLAlchemySource):
     """

datahub/ingestion/source/sql/mssql/source.py

@@ -27,6 +27,7 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.mssql.job_models import (
     JobStep,
     MSSQLDataFlow,
@@ -177,10 +178,18 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_column_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
 )
 class SQLServerSource(SQLAlchemySource):
     """

datahub/ingestion/source/sql/sql_common.py

@@ -54,6 +54,7 @@ from datahub.ingestion.source.common.data_reader import DataReader
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
@@ -305,10 +306,12 @@ class ProfileMetadata:
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default to get lineage for views via `include_view_lineage`",
+    subtype_modifier=[SourceCapabilityModifier.VIEW],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Enabled by default to get lineage for views via `include_view_column_lineage`",
+    subtype_modifier=[SourceCapabilityModifier.VIEW],
 )
 @capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 @capability(

datahub/ingestion/source/sql/sql_generic_profiler.py

@@ -57,10 +57,11 @@ class GenericProfiler:
         platform: Optional[str] = None,
         profiler_args: Optional[Dict] = None,
     ) -> Iterable[MetadataWorkUnit]:
+        # We don't run ge profiling queries if table profiling is enabled or if the row count is 0.
         ge_profile_requests: List[GEProfilerRequest] = [
             cast(GEProfilerRequest, request)
             for request in requests
-            if not request.profile_table_level_only
+            if not request.profile_table_level_only or request.table.rows_count == 0
         ]
         table_level_profile_requests: List[TableProfilerRequest] = [
             request for request in requests if request.profile_table_level_only

datahub/ingestion/source/sql/teradata.py

@@ -445,7 +445,10 @@ class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig):
 @capability(SourceCapability.DOMAINS, "Enabled by default")
 @capability(SourceCapability.CONTAINERS, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
-@capability(SourceCapability.DELETION_DETECTION, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.DELETION_DETECTION,
+    "Enabled by default when stateful ingestion is turned on",
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
 @capability(SourceCapability.LINEAGE_FINE, "Optionally enabled via configuration")

datahub/ingestion/source/sql/vertica.py

@@ -25,6 +25,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_reader import DataReader
+from datahub.ingestion.source.common.subtypes import (
+    DatasetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -113,6 +117,10 @@ class VerticaConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default, can be disabled via configuration `include_view_lineage` and `include_projection_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.PROJECTIONS,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
@@ -497,7 +505,7 @@ class VerticaSource(SQLAlchemySource):
             changeType=ChangeTypeClass.UPSERT,
             entityUrn=dataset_urn,
             aspectName="subTypes",
-            aspect=SubTypesClass(typeNames=["Projections"]),
+            aspect=SubTypesClass(typeNames=[DatasetSubTypes.PROJECTIONS]),
         ).as_workunit()
 
         if self.config.domain:

datahub/ingestion/source/tableau/tableau.py

@@ -80,6 +80,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     BIContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -867,10 +868,14 @@ def report_user_role(report: TableauSourceReport, server: Server) -> None:
 @capability(
     SourceCapability.USAGE_STATS,
     "Dashboard/Chart view counts, enabled using extract_usage_stats config",
+    subtype_modifier=[
+        SourceCapabilityModifier.DASHBOARD,
+        SourceCapabilityModifier.CHART,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Enabled by default when stateful ingestion is turned on.",
+    "Enabled by default via stateful ingestion.",
 )
 @capability(SourceCapability.OWNERSHIP, "Requires recipe configuration")
 @capability(SourceCapability.TAGS, "Requires recipe configuration")

datahub/ingestion/source/unity/source.py

@@ -1020,29 +1020,45 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     ) -> Iterable[MetadataWorkUnit]:
         if self.ctx.graph and self.platform_resource_repository:
             for tag in tags:
-                platform_resource_id = UnityCatalogTagPlatformResourceId.from_tag(
-                    platform_instance=self.platform_instance_name,
-                    platform_resource_repository=self.platform_resource_repository,
-                    tag=tag,
-                )
-                logger.debug(f"Created platform resource {platform_resource_id}")
+                try:
+                    platform_resource_id = UnityCatalogTagPlatformResourceId.from_tag(
+                        platform_instance=self.platform_instance_name,
+                        platform_resource_repository=self.platform_resource_repository,
+                        tag=tag,
+                    )
+                    logger.debug(f"Created platform resource {platform_resource_id}")
 
-                unity_catalog_tag = UnityCatalogTagPlatformResource.get_from_datahub(
-                    platform_resource_id, self.platform_resource_repository, False
-                )
-                if (
-                    tag.to_datahub_tag_urn().urn()
-                    not in unity_catalog_tag.datahub_linked_resources().urns
-                ):
-                    unity_catalog_tag.datahub_linked_resources().add(
-                        tag.to_datahub_tag_urn().urn()
+                    unity_catalog_tag = (
+                        UnityCatalogTagPlatformResource.get_from_datahub(
+                            platform_resource_id,
+                            self.platform_resource_repository,
+                            False,
+                        )
                     )
-                    platform_resource = unity_catalog_tag.as_platform_resource()
-                    for mcp in platform_resource.to_mcps():
-                        yield MetadataWorkUnit(
-                            id=f"platform_resource-{platform_resource.id}",
-                            mcp=mcp,
+                    if (
+                        tag.to_datahub_tag_urn().urn()
+                        not in unity_catalog_tag.datahub_linked_resources().urns
+                    ):
+                        unity_catalog_tag.datahub_linked_resources().add(
+                            tag.to_datahub_tag_urn().urn()
                         )
+                        platform_resource = unity_catalog_tag.as_platform_resource()
+                        for mcp in platform_resource.to_mcps():
+                            yield MetadataWorkUnit(
+                                id=f"platform_resource-{platform_resource.id}",
+                                mcp=mcp,
+                            )
+                except Exception as e:
+                    logger.exception(
+                        f"Error processing platform resource for tag {tag}"
+                    )
+                    self.report.report_warning(
+                        message="Error processing platform resource for tag",
+                        context=str(tag),
+                        title="Error processing platform resource for tag",
+                        exc=e,
+                    )
+                    continue
 
     def _create_schema_metadata_aspect(
         self, table: Table

datahub/ingestion/transformer/add_dataset_ownership.py

@@ -71,8 +71,24 @@ class AddDatasetOwnership(OwnershipTransformer):
 
         server_ownership = graph.get_ownership(entity_urn=urn)
         if server_ownership:
-            owners = {owner.owner: owner for owner in server_ownership.owners}
-            owners.update({owner.owner: owner for owner in mce_ownership.owners})
+            owners = {
+                (
+                    owner.owner,
+                    owner.type,
+                    owner.typeUrn,
+                ): owner
+                for owner in server_ownership.owners
+            }
+            owners.update(
+                {
+                    (
+                        owner.owner,
+                        owner.type,
+                        owner.typeUrn,
+                    ): owner
+                    for owner in mce_ownership.owners
+                }
+            )
             mce_ownership.owners = list(owners.values())
 
         return mce_ownership