acryl-datahub-cloud 0.3.7.9.1__py3-none-any.whl → 0.3.8__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.7.9.1",
+  "version": "0.3.8",
   "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,267 @@
+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
+                else:
+                    scroll_id = next_scroll_id
+
+                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: Optional[Dict[str, Any]]
+    ) -> Optional[str]:
+        if share_config is None:
+            return None
+
+        """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_metadata_sharing/scroll_shared_entities.gql

@@ -0,0 +1,204 @@
+fragment shareDetails on Share {
+  lastShareResults {
+    destination {
+      urn
+    }
+    implicitShareEntity {
+      urn
+      type
+    }
+    created {
+      time
+      actor {
+          urn
+      }
+    }
+    lastAttempt {
+      time
+      actor {
+        urn
+      }
+    }
+    lastSuccess {
+      time
+      actor {
+        urn
+      }
+    }
+    status
+    message
+    shareConfig {
+      enableUpstreamLineage
+      enableDownstreamLineage
+    }
+    statusLastUpdated
+  }
+  lastUnshareResults {
+    destination {
+      urn
+    }
+    implicitShareEntity {
+      urn
+      type
+    }
+    created {
+      time
+      actor {
+        urn
+      }
+    }
+    lastAttempt {
+      time
+      actor {
+        urn
+      }
+    }
+    lastSuccess {
+      time
+      actor {
+        urn
+      }
+    }
+    status
+    message
+    shareConfig {
+      enableUpstreamLineage
+      enableDownstreamLineage
+    }
+    statusLastUpdated
+  }
+}
+
+
+query scrollSharedEntities($scrollId: String, $count: Int!) {
+  scrollAcrossEntities(input: {
+    types: [],
+    query: "*",
+    scrollId: $scrollId,
+    count: $count,
+    orFilters: [
+        {
+            and: [
+                {
+                    field: "shareDestinations",
+                    condition: EXISTS,
+                    values: []
+                }
+            ]
+        }
+    ]
+  }) {
+    nextScrollId
+    count
+    total
+    searchResults {
+        entity {
+            urn 
+            type
+            ...on Dataset {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Container {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Dashboard {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Chart {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on DataJob {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on DataFlow {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on MLModel {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on MLModelGroup {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on MLPrimaryKey {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on MLFeature {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on MLFeatureTable {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Domain {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on DataProduct {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Tag {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on CorpUser {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on CorpGroup {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on GlossaryTerm {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on GlossaryNode {
+                share {
+                    ...shareDetails
+                }
+            }
+            ...on Notebook {
+                share {
+                    ...shareDetails
+                }
+            }
+        }
+        matchedFields {
+            name
+            value
+        }
+        insights {
+            text
+            icon
+        }
+    }
+  }
+}

acryl_datahub_cloud/datahub_metadata_sharing/share_entity.gql

@@ -0,0 +1,9 @@
+mutation shareEntity($entityUrn: String!, $destinationUrn: String!, $lineageDirection: ShareLineageDirection) {
+    shareEntity(input: {
+        entityUrn: $entityUrn,
+        connectionUrn: $destinationUrn,
+        lineageDirection: $lineageDirection
+    }) {
+        succeeded
+    }
+}

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: