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

datahub/ingestion/source/aws/platform_resource_repository.py

@@ -0,0 +1,30 @@
+import logging
+from typing import Optional
+
+from datahub.api.entities.external.external_entities import (
+    PlatformResourceRepository,
+)
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.aws.tag_entities import (
+    LakeFormationTagPlatformResource,
+    LakeFormationTagPlatformResourceId,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GluePlatformResourceRepository(
+    PlatformResourceRepository[
+        LakeFormationTagPlatformResourceId, LakeFormationTagPlatformResource
+    ]
+):
+    """AWS Glue-specific platform resource repository with tag-related operations."""
+
+    def __init__(
+        self,
+        graph: DataHubGraph,
+        platform_instance: Optional[str] = None,
+        catalog: Optional[str] = None,
+    ):
+        super().__init__(graph, platform_instance)
+        self.catalog = catalog

datahub/ingestion/source/aws/s3_boto_utils.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Iterable, Optional, Union
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Iterable, Optional, Union
 
 from datahub.emitter.mce_builder import make_tag_urn
 from datahub.ingestion.api.common import PipelineContext
@@ -11,9 +12,14 @@ from datahub.ingestion.source.aws.s3_util import (
 )
 from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
 
+if TYPE_CHECKING:
+    from mypy_boto3_s3.service_resource import ObjectSummary
+
 logging.getLogger("py4j").setLevel(logging.ERROR)
 logger: logging.Logger = logging.getLogger(__name__)
 
+LIST_OBJECTS_PAGE_SIZE = 1000
+
 
 def get_s3_tags(
     bucket_name: str,
@@ -74,16 +80,82 @@ def get_s3_tags(
     return new_tags
 
 
+@dataclass
+class DirEntry:
+    """
+    Intended to be similar to os.DirEntry, which contains a name, full path, and possibly
+    other attributes of a directory entry. Currently only used to represent S3 folder-like
+    paths.
+    """
+
+    name: str
+    path: str
+
+
 def list_folders_path(
-    s3_uri: str, aws_config: Optional[AwsConnectionConfig]
-) -> Iterable[str]:
+    s3_uri: str,
+    *,
+    startswith: str = "",
+    aws_config: Optional[AwsConnectionConfig] = None,
+) -> Iterable[DirEntry]:
+    """
+    Given an S3 URI to a folder or bucket, return all sub-folders underneath that URI,
+    optionally filtering by startswith. Returned entries never contain a trailing slash.
+    """
+
+    if not is_s3_uri(s3_uri):
+        raise ValueError("Not a s3 URI: " + s3_uri)
+    if aws_config is None:
+        raise ValueError("aws_config not set. Cannot browse s3")
+
+    if not s3_uri.endswith("/"):
+        s3_uri += "/"
+
+    bucket_name = get_bucket_name(s3_uri)
+    if not bucket_name:
+        # No bucket name means we only have the s3[an]:// protocol, not a full bucket and
+        # prefix.
+        for folder in list_buckets(startswith, aws_config):
+            yield DirEntry(name=folder, path=f"{s3_uri}{folder}")
+        return
+
+    prefix = get_bucket_relative_path(s3_uri) + startswith
+    for folder in list_folders(bucket_name, prefix, aws_config):
+        folder = folder.removesuffix("/").split("/")[-1]
+        yield DirEntry(name=folder, path=f"{s3_uri}{folder}")
+
+
+def list_objects_recursive_path(
+    s3_uri: str,
+    *,
+    startswith: str = "",
+    aws_config: Optional[AwsConnectionConfig] = None,
+) -> Iterable["ObjectSummary"]:
+    """
+    Given an S3 URI to a folder or bucket, return all objects underneath that URI, optionally
+    filtering by startswith.
+    """
+
     if not is_s3_uri(s3_uri):
         raise ValueError("Not a s3 URI: " + s3_uri)
     if aws_config is None:
         raise ValueError("aws_config not set. Cannot browse s3")
+    if startswith and "/" in startswith:
+        raise ValueError(f"startswith contains forward slash: {repr(startswith)}")
+
+    if not s3_uri.endswith("/"):
+        s3_uri += "/"
+
     bucket_name = get_bucket_name(s3_uri)
-    prefix = get_bucket_relative_path(s3_uri)
-    yield from list_folders(bucket_name, prefix, aws_config)
+    if not bucket_name:
+        # No bucket name means we only have the s3[an]:// protocol, not a full bucket and
+        # prefix.
+        for bucket_name in list_buckets(startswith, aws_config):
+            yield from list_objects_recursive(bucket_name, "", aws_config)
+        return
+
+    prefix = get_bucket_relative_path(s3_uri) + startswith
+    yield from list_objects_recursive(bucket_name, prefix, aws_config)
 
 
 def list_folders(
@@ -99,3 +171,26 @@ def list_folders(
             if folder.endswith("/"):
                 folder = folder[:-1]
             yield f"{folder}"
+
+
+def list_buckets(
+    prefix: str, aws_config: Optional[AwsConnectionConfig]
+) -> Iterable[str]:
+    if aws_config is None:
+        raise ValueError("aws_config not set. Cannot browse s3")
+    s3_client = aws_config.get_s3_client()
+    paginator = s3_client.get_paginator("list_buckets")
+    for page in paginator.paginate(Prefix=prefix):
+        for o in page.get("Buckets", []):
+            yield str(o.get("Name"))
+
+
+def list_objects_recursive(
+    bucket_name: str, prefix: str, aws_config: Optional[AwsConnectionConfig]
+) -> Iterable["ObjectSummary"]:
+    if aws_config is None:
+        raise ValueError("aws_config not set. Cannot browse s3")
+    s3_resource = aws_config.get_s3_resource()
+    bucket = s3_resource.Bucket(bucket_name)
+    for obj in bucket.objects.filter(Prefix=prefix).page_size(LIST_OBJECTS_PAGE_SIZE):
+        yield obj

datahub/ingestion/source/aws/tag_entities.py

@@ -0,0 +1,270 @@
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.source.aws.platform_resource_repository import (
+        GluePlatformResourceRepository,
+    )
+
+from pydantic import BaseModel
+
+from datahub.api.entities.external.external_entities import (
+    ExternalEntity,
+    ExternalEntityId,
+    LinkedResourceSet,
+)
+from datahub.api.entities.external.lake_formation_external_entites import (
+    LakeFormationTag,
+)
+from datahub.api.entities.platformresource.platform_resource import (
+    PlatformResource,
+    PlatformResourceKey,
+)
+from datahub.metadata.urns import TagUrn
+from datahub.utilities.urns.urn import Urn
+
+logger = logging.getLogger(__name__)
+
+
+class LakeFormationTagSyncContext(BaseModel):
+    # it is intentionally empty
+    platform_instance: Optional[str] = None
+    catalog: Optional[str] = None
+
+    # Making it compatible with SyncContext interface
+    def get_platform_instance(self) -> Optional[str]:
+        return self.platform_instance
+
+
+class LakeFormationTagPlatformResourceId(ExternalEntityId):
+    """
+    A LakeFormationTag is a unique identifier for a Lakeformation tag.
+    """
+
+    tag_key: str
+    tag_value: Optional[str] = None
+    platform_instance: Optional[str] = None
+    catalog: Optional[str] = None
+    exists_in_lake_formation: bool = False
+    persisted: bool = False
+
+    # this is a hack to make sure the property is a string and not private pydantic field
+    @staticmethod
+    def _RESOURCE_TYPE() -> str:
+        return "LakeFormationTagPlatformResource"
+
+    def to_platform_resource_key(self) -> PlatformResourceKey:
+        return PlatformResourceKey(
+            platform="glue",
+            resource_type=str(LakeFormationTagPlatformResourceId._RESOURCE_TYPE()),
+            primary_key=f"{self.catalog}.{self.tag_key}:{self.tag_value}"
+            if self.catalog
+            else f"{self.tag_key}:{self.tag_value}",
+            platform_instance=self.platform_instance,
+        )
+
+    @classmethod
+    def get_or_create_from_tag(
+        cls,
+        tag: LakeFormationTag,
+        platform_resource_repository: "GluePlatformResourceRepository",
+        exists_in_lake_formation: bool = False,
+        catalog_id: Optional[str] = None,
+    ) -> "LakeFormationTagPlatformResourceId":
+        """
+        Creates a LakeFormationTagPlatformResourceId from a LakeFormationTag.
+        """
+
+        # Use catalog_id if provided, otherwise fall back to repository catalog
+        effective_catalog = catalog_id or platform_resource_repository.catalog
+
+        existing_platform_resource = cls.search_by_urn(
+            tag.to_datahub_tag_urn().urn(),
+            platform_resource_repository=platform_resource_repository,
+            tag_sync_context=LakeFormationTagSyncContext(
+                platform_instance=platform_resource_repository.platform_instance,
+                catalog=effective_catalog,
+            ),
+        )
+        if existing_platform_resource:
+            logger.info(
+                f"Found existing LakeFormationTagPlatformResourceId for tag {tag.key}: {existing_platform_resource}"
+            )
+            return existing_platform_resource
+
+        return LakeFormationTagPlatformResourceId(
+            tag_key=str(tag.key),
+            tag_value=str(tag.value) if tag.value is not None else None,
+            platform_instance=platform_resource_repository.platform_instance,
+            catalog=effective_catalog,
+            exists_in_lake_formation=exists_in_lake_formation,
+            persisted=False,
+        )
+
+    @classmethod
+    def search_by_urn(
+        cls,
+        urn: str,
+        platform_resource_repository: "GluePlatformResourceRepository",
+        tag_sync_context: LakeFormationTagSyncContext,
+    ) -> Optional["LakeFormationTagPlatformResourceId"]:
+        """
+        Search for existing Lake Formation tag entity by URN using repository caching.
+
+        This method now delegates to the repository's search_entity_by_urn method to ensure
+        consistent caching behavior across all platform implementations.
+        """
+        # Use repository's cached search method instead of duplicating search logic
+        existing_entity_id = platform_resource_repository.search_entity_by_urn(urn)
+
+        if existing_entity_id:
+            # Verify platform instance and catalog match
+            if (
+                existing_entity_id.platform_instance
+                == tag_sync_context.platform_instance
+                and existing_entity_id.catalog == tag_sync_context.catalog
+            ):
+                logger.info(
+                    f"Found existing LakeFormationTagPlatformResourceId for URN {urn}: {existing_entity_id}"
+                )
+                # Create a new ID with the correct state instead of mutating
+                return LakeFormationTagPlatformResourceId(
+                    tag_key=existing_entity_id.tag_key,
+                    tag_value=existing_entity_id.tag_value,
+                    platform_instance=existing_entity_id.platform_instance,
+                    catalog=existing_entity_id.catalog,
+                    exists_in_lake_formation=True,  # This tag exists in Lake Formation
+                    persisted=True,  # And it's persisted in DataHub
+                )
+
+        logger.info(
+            f"No mapped tag found for URN {urn} with platform instance {tag_sync_context.platform_instance}. Creating a new LakeFormationTagPlatformResourceId."
+        )
+        return None
+
+    @classmethod
+    def from_datahub_urn(
+        cls,
+        urn: str,
+        platform_resource_repository: "GluePlatformResourceRepository",
+        tag_sync_context: LakeFormationTagSyncContext,
+    ) -> "LakeFormationTagPlatformResourceId":
+        """
+        Creates a UnityCatalogTagPlatformResourceId from a DataHub URN.
+        """
+        # First we check if we already have a mapped platform resource for this
+        # urn that is of the type UnityCatalogTagPlatformResource
+        # If we do, we can use it to create the UnityCatalogTagPlatformResourceId
+        # Else, we need to generate a new UnityCatalogTagPlatformResourceId
+        existing_platform_resource_id = cls.search_by_urn(
+            urn, platform_resource_repository, tag_sync_context
+        )
+        if existing_platform_resource_id:
+            logger.info(
+                f"Found existing LakeFormationTagPlatformResourceId for URN {urn}: {existing_platform_resource_id}"
+            )
+            return existing_platform_resource_id
+
+        # Otherwise, we need to create a new UnityCatalogTagPlatformResourceId
+        new_tag_id = cls.generate_tag_id(tag_sync_context, urn)
+        if new_tag_id:
+            # we then check if this tag has already been ingested as a platform
+            # resource in the platform resource repository
+            resource_key = platform_resource_repository.get(
+                new_tag_id.to_platform_resource_key()
+            )
+            if resource_key:
+                logger.info(
+                    f"Tag {new_tag_id} already exists in platform resource repository with {resource_key}"
+                )
+                # Create a new ID with the correct state instead of mutating
+                return LakeFormationTagPlatformResourceId(
+                    tag_key=new_tag_id.tag_key,
+                    tag_value=new_tag_id.tag_value,
+                    platform_instance=new_tag_id.platform_instance,
+                    catalog=new_tag_id.catalog,
+                    exists_in_lake_formation=True,  # This tag exists in Lake Formation
+                    persisted=new_tag_id.persisted,
+                )
+            return new_tag_id
+        raise ValueError(f"Unable to create LakeFormationTagId from DataHub URN: {urn}")
+
+    @classmethod
+    def generate_tag_id(
+        cls, tag_sync_context: LakeFormationTagSyncContext, urn: str
+    ) -> "LakeFormationTagPlatformResourceId":
+        parsed_urn = Urn.from_string(urn)
+        entity_type = parsed_urn.entity_type
+        if entity_type == "tag":
+            new_tag_id = LakeFormationTagPlatformResourceId.from_datahub_tag(
+                TagUrn.from_string(urn), tag_sync_context
+            )
+        else:
+            raise ValueError(f"Unsupported entity type {entity_type} for URN {urn}")
+        return new_tag_id
+
+    @classmethod
+    def from_datahub_tag(
+        cls, tag_urn: TagUrn, tag_sync_context: LakeFormationTagSyncContext
+    ) -> "LakeFormationTagPlatformResourceId":
+        tag = LakeFormationTag.from_urn(tag_urn)
+
+        return LakeFormationTagPlatformResourceId(
+            tag_key=str(tag.key),
+            tag_value=str(tag.value),
+            platform_instance=tag_sync_context.platform_instance,
+            catalog=tag_sync_context.catalog,
+            exists_in_lake_formation=False,
+        )
+
+
+class LakeFormationTagPlatformResource(ExternalEntity):
+    datahub_urns: LinkedResourceSet
+    managed_by_datahub: bool
+    id: LakeFormationTagPlatformResourceId
+    allowed_values: Optional[List[str]] = None
+
+    def get_id(self) -> ExternalEntityId:
+        return self.id
+
+    def is_managed_by_datahub(self) -> bool:
+        return self.managed_by_datahub
+
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        return self.datahub_urns
+
+    def as_platform_resource(self) -> PlatformResource:
+        return PlatformResource.create(
+            key=self.id.to_platform_resource_key(),
+            secondary_keys=[u for u in self.datahub_urns.urns],
+            value=self,
+        )
+
+    @classmethod
+    def create_default(
+        cls,
+        entity_id: ExternalEntityId,
+        managed_by_datahub: bool,
+    ) -> "LakeFormationTagPlatformResource":
+        """Create a default Lake Formation tag entity when none found in DataHub."""
+        # Type narrowing: we know this will be a LakeFormationTagPlatformResourceId
+        assert isinstance(entity_id, LakeFormationTagPlatformResourceId), (
+            f"Expected LakeFormationTagPlatformResourceId, got {type(entity_id)}"
+        )
+
+        # Create a new entity ID with correct default state instead of mutating
+        default_entity_id = LakeFormationTagPlatformResourceId(
+            tag_key=entity_id.tag_key,
+            tag_value=entity_id.tag_value,
+            platform_instance=entity_id.platform_instance,
+            catalog=entity_id.catalog,
+            exists_in_lake_formation=False,  # New entities don't exist in Lake Formation yet
+            persisted=False,  # New entities are not persisted yet
+        )
+
+        return cls(
+            id=default_entity_id,
+            datahub_urns=LinkedResourceSet(urns=[]),
+            managed_by_datahub=managed_by_datahub,
+            allowed_values=None,
+        )

datahub/ingestion/source/azure/azure_common.py

@@ -61,13 +61,13 @@ class AzureConnectionConfig(ConfigModel):
     def get_blob_service_client(self):
         return BlobServiceClient(
             account_url=f"https://{self.account_name}.blob.core.windows.net",
-            credential=f"{self.get_credentials()}",
+            credential=self.get_credentials(),
         )
 
     def get_data_lake_service_client(self) -> DataLakeServiceClient:
         return DataLakeServiceClient(
             account_url=f"https://{self.account_name}.dfs.core.windows.net",
-            credential=f"{self.get_credentials()}",
+            credential=self.get_credentials(),
         )
 
     def get_credentials(
@@ -81,7 +81,7 @@ class AzureConnectionConfig(ConfigModel):
             )
         return self.sas_token if self.sas_token is not None else self.account_key
 
-    @root_validator()
+    @root_validator(skip_on_failure=True)
     def _check_credential_values(cls, values: Dict) -> Dict:
         if (
             values.get("account_key")

datahub/ingestion/source/bigquery_v2/bigquery.py

@@ -4,6 +4,7 @@ import logging
 import os
 from typing import Iterable, List, Optional
 
+from datahub.configuration.common import AllowDenyPattern
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -44,9 +45,11 @@ from datahub.ingestion.source.bigquery_v2.queries_extractor import (
     BigQueryQueriesExtractorConfig,
 )
 from datahub.ingestion.source.bigquery_v2.usage import BigQueryUsageExtractor
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler
 from datahub.ingestion.source.state.redundant_run_skip_handler import (
     RedundantLineageRunSkipHandler,
+    RedundantQueriesRunSkipHandler,
     RedundantUsageRunSkipHandler,
 )
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
@@ -77,7 +80,14 @@ def cleanup(config: BigQueryV2Config) -> None:
     supported=False,
 )
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.BIGQUERY_PROJECT,
+        SourceCapabilityModifier.BIGQUERY_DATASET,
+    ],
+)
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
 @capability(
     SourceCapability.DATA_PROFILING,
@@ -99,6 +109,7 @@ def cleanup(config: BigQueryV2Config) -> None:
     SourceCapability.PARTITION_SUPPORT,
     "Enabled by default, partition keys and clustering keys are supported.",
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
     def __init__(self, ctx: PipelineContext, config: BigQueryV2Config):
         super().__init__(config, ctx)
@@ -135,7 +146,10 @@ class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
         redundant_lineage_run_skip_handler: Optional[RedundantLineageRunSkipHandler] = (
             None
         )
-        if self.config.enable_stateful_lineage_ingestion:
+        if (
+            self.config.enable_stateful_lineage_ingestion
+            and not self.config.use_queries_v2
+        ):
             redundant_lineage_run_skip_handler = RedundantLineageRunSkipHandler(
                 source=self,
                 config=self.config,
@@ -241,7 +255,23 @@ class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
             ).workunit_processor,
         ]
 
+    def _warn_deprecated_configs(self):
+        if (
+            self.config.match_fully_qualified_names is not None
+            and not self.config.match_fully_qualified_names
+            and self.config.schema_pattern is not None
+            and self.config.schema_pattern != AllowDenyPattern.allow_all()
+        ):
+            self.report.report_warning(
+                message="Please update `schema_pattern` to match against fully qualified schema name `<database_name>.<schema_name>` and set config `match_fully_qualified_names : True`."
+                "Current default `match_fully_qualified_names: False` is only to maintain backward compatibility. "
+                "The config option `match_fully_qualified_names` will be removed in future and the default behavior will be like `match_fully_qualified_names: True`.",
+                context="Config option deprecation warning",
+                title="Config option deprecation warning",
+            )
+
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        self._warn_deprecated_configs()
         projects = get_projects(
             self.bq_schema_extractor.schema_api,
             self.report,
@@ -270,28 +300,41 @@ class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
             ):
                 return
 
-            with self.report.new_stage(
-                f"*: {QUERIES_EXTRACTION}"
-            ), BigQueryQueriesExtractor(
-                connection=self.config.get_bigquery_client(),
-                schema_api=self.bq_schema_extractor.schema_api,
-                config=BigQueryQueriesExtractorConfig(
-                    window=self.config,
-                    user_email_pattern=self.config.usage.user_email_pattern,
-                    include_lineage=self.config.include_table_lineage,
-                    include_usage_statistics=self.config.include_usage_statistics,
-                    include_operations=self.config.usage.include_operational_stats,
-                    include_queries=self.config.include_queries,
-                    include_query_usage_statistics=self.config.include_query_usage_statistics,
-                    top_n_queries=self.config.usage.top_n_queries,
-                    region_qualifiers=self.config.region_qualifiers,
-                ),
-                structured_report=self.report,
-                filters=self.filters,
-                identifiers=self.identifiers,
-                schema_resolver=self.sql_parser_schema_resolver,
-                discovered_tables=self.bq_schema_extractor.table_refs,
-            ) as queries_extractor:
+            redundant_queries_run_skip_handler: Optional[
+                RedundantQueriesRunSkipHandler
+            ] = None
+            if self.config.enable_stateful_time_window:
+                redundant_queries_run_skip_handler = RedundantQueriesRunSkipHandler(
+                    source=self,
+                    config=self.config,
+                    pipeline_name=self.ctx.pipeline_name,
+                    run_id=self.ctx.run_id,
+                )
+
+            with (
+                self.report.new_stage(f"*: {QUERIES_EXTRACTION}"),
+                BigQueryQueriesExtractor(
+                    connection=self.config.get_bigquery_client(),
+                    schema_api=self.bq_schema_extractor.schema_api,
+                    config=BigQueryQueriesExtractorConfig(
+                        window=self.config,
+                        user_email_pattern=self.config.usage.user_email_pattern,
+                        include_lineage=self.config.include_table_lineage,
+                        include_usage_statistics=self.config.include_usage_statistics,
+                        include_operations=self.config.usage.include_operational_stats,
+                        include_queries=self.config.include_queries,
+                        include_query_usage_statistics=self.config.include_query_usage_statistics,
+                        top_n_queries=self.config.usage.top_n_queries,
+                        region_qualifiers=self.config.region_qualifiers,
+                    ),
+                    structured_report=self.report,
+                    filters=self.filters,
+                    identifiers=self.identifiers,
+                    redundant_run_skip_handler=redundant_queries_run_skip_handler,
+                    schema_resolver=self.sql_parser_schema_resolver,
+                    discovered_tables=self.bq_schema_extractor.table_refs,
+                ) as queries_extractor,
+            ):
                 self.report.queries_extractor = queries_extractor.report
                 yield from queries_extractor.get_workunits_internal()
         else: