acryl-datahub-cloud 0.3.7.9rc1__py3-none-any.whl → 0.3.8rc1__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.7.9rc1",
+  "version": "0.3.8rc1",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

acryl_datahub_cloud/acryl_cs_issues/source.py

@@ -100,7 +100,6 @@ class AcrylCSIssuesSource(Source):
     def _provision_platform(
         self, platform: str, logo_url: str, graph: DataHubGraph
     ) -> None:
-
         platform_urn = make_data_platform_urn(platform)
         if not graph.exists(platform_urn):
             platform_info = DataPlatformInfoClass(

acryl_datahub_cloud/api/__init__.py

@@ -0,0 +1 @@
+from acryl_datahub_cloud.api.client import AcrylGraph

acryl_datahub_cloud/api/client.py

@@ -0,0 +1,6 @@
+from acryl_datahub_cloud.api.entity_versioning import EntityVersioningAPI
+from datahub.ingestion.graph.client import DataHubGraph
+
+
+class AcrylGraph(EntityVersioningAPI, DataHubGraph):
+    pass

acryl_datahub_cloud/api/entity_versioning.py

@@ -0,0 +1,167 @@
+import uuid
+from typing import Optional
+
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.schema_classes import (
+    VersionPropertiesClass,
+    VersionSetPropertiesClass,
+)
+from datahub.metadata.urns import VersionSetUrn
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+class EntityVersioningAPI(DataHubGraph):
+    LINK_VERSION_MUTATION = """
+        mutation($input: LinkVersionInput!) {
+            linkAssetVersion(input: $input)
+        }
+    """
+
+    UNLINK_VERSION_MUTATION = """
+        mutation($input: UnlinkVersionInput!) {
+            unlinkAssetVersion(input: $input)
+        }
+    """
+
+    def link_asset_to_version_set(
+        self,
+        asset_urn: str,
+        version_set_urn: Optional[str],
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> str:
+        """Sets an entity as the latest version of a version set.
+
+        Can also be used to create a new version set, with `asset_urn` as the first version.
+
+        Args:
+            asset_urn: URN of the entity.
+            version_set_urn: URN of the version set, or None to generate a new version set urn
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `asset_urn` was linked.
+        """
+
+        entity_type = guess_entity_type(asset_urn)
+        if version_set_urn is None:
+            version_set_urn = VersionSetUrn(str(uuid.uuid4()), entity_type).urn()
+        elif guess_entity_type(version_set_urn) != "versionSet":
+            raise ValueError(f"Expected version set URN, got {version_set_urn}")
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if entity_version:
+            raise ValueError(
+                f"Asset {asset_urn} is already a version of {entity_version.versionSet}"
+            )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "linkedEntity": asset_urn,
+                "version": label,
+                "comment": comment,
+            }
+        }
+        self.execute_graphql(self.LINK_VERSION_MUTATION, variables)
+        return version_set_urn
+
+    def link_asset_to_versioned_asset(
+        self,
+        new_asset_urn: str,
+        old_asset_urn: str,
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> str:
+        """Sets an entity as the latest version of an existing versioned entity.
+
+        Args:
+            new_asset_urn: URN of the new latest entity.
+            old_asset_urn: URN of an existing versioned entity to link onto.
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `new_asset_urn` was linked.
+        """
+
+        new_entity_type = guess_entity_type(new_asset_urn)
+        old_entity_type = guess_entity_type(old_asset_urn)
+        if new_entity_type != old_entity_type:
+            raise ValueError(
+                f"Expected URNs of the same type, got {new_entity_type} and {old_entity_type}"
+            )
+
+        new_entity_version = self.get_aspect(new_asset_urn, VersionPropertiesClass)
+        if new_entity_version:
+            raise ValueError(
+                f"Asset {new_asset_urn} is already a version of {new_entity_version.versionSet}"
+            )
+        old_entity_version = self.get_aspect(old_asset_urn, VersionPropertiesClass)
+        if not old_entity_version:
+            raise ValueError(f"Asset {old_asset_urn} is not versioned")
+
+        version_set_urn = old_entity_version.versionSet
+        self.link_asset_to_version_set(
+            new_asset_urn, version_set_urn, label, comment=comment
+        )
+        return version_set_urn
+
+    def unlink_asset_from_version_set(self, asset_urn: str) -> Optional[str]:
+        """Unlinks an entity from its version set.
+
+        Args:
+            asset_urn: URN of the entity to unlink from its version set.
+
+        Returns:
+            If successful, the URN of the version set from which `asset_urn` was unlinked.
+        """
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if not entity_version:
+            raise ValueError(f"Asset {asset_urn} is not versioned")
+
+        variables = {
+            "input": {
+                "versionSet": entity_version.versionSet,
+                "unlinkedEntity": asset_urn,
+            }
+        }
+        if self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables):
+            return entity_version.versionSet
+        else:
+            return None
+
+    def unlink_latest_asset_from_version_set(
+        self, version_set_urn: str
+    ) -> Optional[str]:
+        """Unlinks the latest version of a version set.
+
+        Args:
+            version_set_urn: URN of the version set.
+
+        Returns:
+            If successful, the URN of the entity that was unlinked from `version_set_urn`.
+        """
+
+        version_set_properties = self.get_aspect(
+            version_set_urn, VersionSetPropertiesClass
+        )
+        if not version_set_properties:
+            raise ValueError(
+                f"Version set {version_set_urn} does not exist or has no versions"
+            )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "unlinkedEntity": version_set_properties.latest,
+            }
+        }
+        if self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables):
+            return version_set_properties.latest
+        else:
+            return None

acryl_datahub_cloud/datahub_metadata_sharing/metadata_sharing_source.py

@@ -0,0 +1,262 @@
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+from pydantic import BaseModel
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from acryl_datahub_cloud.datahub_metadata_sharing.query import (
+    GRAPHQL_SCROLL_SHARED_ENTITIES,
+    GRAPHQL_SHARE_ENTITY,
+)
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import Source, SourceReport
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.graph.client import DataHubGraph
+
+logger = logging.getLogger(__name__)
+
+
+class GraphQLError(Exception):
+    """Custom exception for GraphQL-specific errors"""
+
+    pass
+
+
+class DataHubMetadataSharingSourceConfig(BaseModel):
+    batch_size: int = 100
+    batch_delay_ms: int = 100
+    max_retries: int = 3
+    initial_retry_delay_ms: int = 1000
+
+
+@dataclass
+class DataHubMetadataSharingSourceReport(SourceReport):
+    entities_shared: int = 0
+    entities_failed: int = 0
+    implicit_entities_skipped: int = 0
+    batches_processed: int = 0
+
+
+@platform_name(id="datahub", platform_name="DataHub")
+@config_class(DataHubMetadataSharingSourceConfig)
+@support_status(SupportStatus.INCUBATING)
+class DataHubMetadataSharingSource(Source):
+    """MetadataSharing Source that reshares entities across DataHub instances"""
+
+    def __init__(
+        self, config: DataHubMetadataSharingSourceConfig, ctx: PipelineContext
+    ):
+        super().__init__(ctx)
+        self.config: DataHubMetadataSharingSourceConfig = config
+        self.report = DataHubMetadataSharingSourceReport()
+        self.graph: Optional[DataHubGraph] = None
+
+    @retry(
+        retry=retry_if_exception_type((GraphQLError, ConnectionError)),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True,
+    )
+    def execute_graphql_with_retry(
+        self, query: str, variables: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Execute GraphQL query with retry logic"""
+        if self.graph is None:
+            raise ValueError("Graph client not initialized")
+        response = self.graph.execute_graphql(query, variables=variables)
+        error = response.get("error")
+        if error:
+            raise GraphQLError(f"GraphQL error: {error}")
+        return response
+
+    def get_workunits(self) -> Iterable[MetadataWorkUnit]:
+        self.graph = self.ctx.require_graph("Loading default graph coordinates.")
+
+        self.reshare_entities()
+
+        # This source doesn't produce any work units
+        return []
+
+    def reshare_entities(self) -> None:
+        scroll_id: Optional[str] = None
+        current_batch_number: int = 1
+
+        try:
+            while True:
+                next_scroll_id, results = self.scroll_shared_entities(
+                    scroll_id, self.config.batch_size
+                )
+
+                for result in results:
+                    self._process_single_entity(result)
+
+                self.report.batches_processed = current_batch_number
+                self.report.info(
+                    message="Completed sharing batch of entities.",
+                    context=f"{current_batch_number} of size {self.config.batch_size}!",
+                )
+                current_batch_number += 1
+
+                if next_scroll_id is None:
+                    break
+
+                time.sleep(self.config.batch_delay_ms / 1000.0)
+
+        except Exception as e:
+            self.report.report_failure(
+                title="Failed to process batches",
+                message="Error occurred while processing one or more batches!",
+                context=f"message = {str(e)}",
+                exc=e,
+            )
+            return
+
+        self.report.info(
+            message="Completed sharing all entities.",
+            context=f"Successfully shared {self.report.entities_shared} entities, "
+            f"failed to share {self.report.entities_failed} entities.",
+        )
+
+    # Rest of the methods remain the same...
+
+    def _process_single_entity(self, result: Dict[str, Any]) -> None:
+        """Process a single entity result"""
+        entity_urn = result.get("entity", {}).get("urn", None)
+        share_results = (
+            result.get("entity", {}).get("share", {}).get("lastShareResults", [])
+        )
+
+        if entity_urn is None:
+            self.report.report_warning(
+                message="Failed to resolve entity urn for shared asset! Skipping...",
+                context=f"Response: {str(result)}",
+            )
+            return
+
+        for share_result in share_results:
+            try:
+                destination_data = share_result.get("destination", {})
+                destination_urn = destination_data.get("urn", "")
+                previous_status = share_result.get("status")
+                share_config = share_result.get("shareConfig", {})
+
+                # Important: If there is implicit entity, we should skip this urn.
+                # This means the entity was not EXPLICITLY shared, so we do not want to explicitly share here.
+                implicit_shared_entity = share_result.get("implicitShareEntity")
+                is_implicitly_shared = (
+                    implicit_shared_entity is not None
+                    and "urn" in implicit_shared_entity
+                )
+
+                if is_implicitly_shared:
+                    self.report.implicit_entities_skipped += 1
+                    continue
+
+                if previous_status != "SUCCESS":
+                    self.report.report_warning(
+                        message="Attempting to share a previously unsuccessful shared entity!",
+                        context=f"entity urn: {entity_urn}, destination urn: {destination_urn}",
+                    )
+
+                lineage_direction = self._determine_lineage_direction(share_config)
+
+                shared = self.share_entity(
+                    entity_urn=entity_urn,
+                    destination_urn=destination_urn,
+                    lineage_direction=lineage_direction,
+                )
+
+                if shared:
+                    self.report.entities_shared += 1
+                else:
+                    self.report.entities_failed += 1
+
+            except Exception as e:
+                self.report.report_warning(
+                    message="Failed to share single entity!",
+                    context=f"entity urn: {entity_urn}",
+                )
+                logger.exception(f"Error processing entity {entity_urn}", e)
+                self.report.entities_failed += 1
+
+    def _determine_lineage_direction(
+        self, share_config: Dict[str, Any]
+    ) -> Optional[str]:
+        """Determine lineage direction based on share config"""
+        include_upstreams = share_config.get("enableUpstreamLineage", False)
+        include_downstreams = share_config.get(
+            "enableDownstreamLineage", False
+        )  # Fixed typo
+
+        if include_upstreams and include_downstreams:
+            return "BOTH"
+        if include_upstreams:
+            return "UPSTREAM"
+        if include_downstreams:
+            return "DOWNSTREAM"
+        return None
+
+    def scroll_shared_entities(
+        self, scroll_id: Optional[str], count: int
+    ) -> Tuple[Optional[str], List[Dict[str, Any]]]:
+        """Scroll through shared entities with retry logic"""
+        response = self.execute_graphql_with_retry(
+            GRAPHQL_SCROLL_SHARED_ENTITIES,
+            variables={
+                "scrollId": scroll_id,
+                "count": count,
+            },
+        )
+
+        result = response.get("scrollAcrossEntities", {})
+        return result.get("nextScrollId"), result.get("searchResults", [])
+
+    def share_entity(
+        self, entity_urn: str, destination_urn: str, lineage_direction: Optional[str]
+    ) -> bool:
+        """Share entity with retry logic"""
+        try:
+            response = self.execute_graphql_with_retry(
+                GRAPHQL_SHARE_ENTITY,
+                variables={
+                    "entityUrn": entity_urn,
+                    "destinationUrn": destination_urn,
+                    "lineageDirection": lineage_direction,
+                },
+            )
+
+            result = response.get("shareEntity", {})
+            if not result.get("succeeded", False):
+                self.report.report_failure(
+                    title="Failed to Share Entity",
+                    message="Response returned that success failed for entity and destination!",
+                    context=f"entity urn: {entity_urn}, destination urn: {destination_urn}",
+                )
+                return False
+
+            return True
+
+        except Exception as e:
+            self.report.report_failure(
+                title="Failed to Share Entity",
+                message="Exception occurred while sharing entity",
+                context=f"entity urn: {entity_urn}, destination urn: {destination_urn}",
+                exc=e,
+            )
+            return False
+
+    def get_report(self) -> SourceReport:
+        return self.report

acryl_datahub_cloud/datahub_metadata_sharing/query.py

@@ -0,0 +1,7 @@
+import pathlib
+
+GRAPHQL_SCROLL_SHARED_ENTITIES = (
+    pathlib.Path(__file__).parent / "scroll_shared_entities.gql"
+).read_text()
+
+GRAPHQL_SHARE_ENTITY = (pathlib.Path(__file__).parent / "share_entity.gql").read_text()

acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -409,7 +409,6 @@ class DataHubBasedS3Dataset:
         physical_uri: str,
         local_file: str,
     ) -> Iterable[MetadataChangeProposalWrapper]:
-
         aspects: List = []
         mcps: List[MetadataChangeProposalWrapper] = self._update_presigned_url(
             dataset_urn, physical_uri
@@ -456,7 +455,6 @@ class DataHubBasedS3Dataset:
         physical_uri: str,
         dataset_properties: Optional[DatasetPropertiesClass] = None,
     ) -> List[MetadataChangeProposalWrapper]:
-
         if self.config.generate_presigned_url:
             external_url = self._generate_presigned_url(physical_uri)
         else:

acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -180,7 +180,6 @@ class DataHubFormReportingData(FormData):
     def form_assigned_date(
         self, search_row: DataHubDatasetSearchRow
     ) -> Dict[str, date]:
-
         form_assigned_dates: Dict[str, date] = {}
         forms = self.graph.get_aspect(search_row.urn, FormsClass)
         if not forms:

acryl_datahub_cloud/datahub_reporting/extract_graph.py

@@ -118,7 +118,6 @@ class DataHubReportingExtractGraphSource(Source):
         return skip_extract
 
     def get_workunits(self):
-
         self.graph = (
             self.ctx.require_graph("Loading default graph coordinates.")
             if self.config.server is None

acryl_datahub_cloud/datahub_reporting/extract_sql.py

@@ -118,7 +118,6 @@ class DataHubReportingExtractSQLSource(Source):
         return skip_extract
 
     def get_workunits(self):
-
         self.graph = (
             self.ctx.require_graph("Loading default graph coordinates.")
             if self.config.server is None

acryl_datahub_cloud/lineage_features/source.py

@@ -83,16 +83,31 @@ class DataHubLineageFeaturesSource(Source):
         query = {
             "query": {
                 "bool": {
-                    "filter": [
-                        {"term": {"source.entityType": "schemaField"}},
-                        {"term": {"destination.entityType": "schemaField"}},
+                    "should": [
+                        {"term": {"relationshipType": "Consumes"}},
                         {"term": {"relationshipType": "DownstreamOf"}},
+                        {"term": {"relationshipType": "TrainedBy"}},
+                        {"term": {"relationshipType": "UsedBy"}},
+                        {"term": {"relationshipType": "MemberOf"}},
+                        {"term": {"relationshipType": "DerivedFrom"}},
+                        {"term": {"relationshipType": "Produces"}},
+                        {"term": {"relationshipType": "DashboardContainsDashboard"}},
+                        {
+                            "bool": {
+                                "must": [
+                                    {"term": {"relationshipType": "Contains"}},
+                                    {"term": {"source.entityType": "dashboard"}},
+                                    {"term": {"destination.entityType": "chart"}},
+                                ]
+                            }
+                        },
                     ],
-                }
+                },
             },
             "sort": [
                 {"source.urn": {"order": "desc"}},
                 {"destination.urn": {"order": "desc"}},
+                {"relationshipType": {"order": "desc"}},
                 {"lifecycleOwner": {"order": "desc"}},
             ],
         }
@@ -124,7 +139,9 @@ class DataHubLineageFeaturesSource(Source):
         for urn in set(self.upstream_counts.keys()).union(
             self.downstream_counts.keys()
         ):
-            print(urn, self.upstream_counts[urn], self.downstream_counts[urn])
+            logger.debug(
+                f"{urn}: {self.upstream_counts[urn]}, {self.downstream_counts[urn]}"
+            )
             yield MetadataChangeProposalWrapper(
                 entityUrn=urn,
                 aspect=LineageFeaturesClass(