acryl-datahub-cloud 0.3.7.9.1__tar.gz → 0.3.8__tar.gz

{acryl_datahub_cloud-0.3.7.9.1 → acryl_datahub_cloud-0.3.8}/PKG-INFO

@@ -1,82 +1,90 @@
 Metadata-Version: 2.2
 Name: acryl-datahub-cloud
-Version: 0.3.7.9.1
+Version: 0.3.8
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
 Provides-Extra: datahub-lineage-features
+Requires-Dist: duckdb; extra == "datahub-lineage-features"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
-Requires-Dist: pyarrow; extra == "datahub-lineage-features"
-Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
 Requires-Dist: pandas; extra == "datahub-lineage-features"
-Requires-Dist: duckdb; extra == "datahub-lineage-features"
+Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
+Requires-Dist: pyarrow; extra == "datahub-lineage-features"
 Provides-Extra: datahub-reporting-forms
-Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
+Requires-Dist: duckdb; extra == "datahub-reporting-forms"
 Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Requires-Dist: pandas; extra == "datahub-reporting-forms"
-Requires-Dist: duckdb; extra == "datahub-reporting-forms"
+Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
+Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
 Provides-Extra: datahub-reporting-extract-graph
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
+Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
+Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
-Provides-Extra: datahub-usage-feature-reporting
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-feature-reporting"
-Requires-Dist: duckdb; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-feature-reporting"
-Requires-Dist: numpy<2; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pyarrow; extra == "datahub-usage-feature-reporting"
-Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-feature-reporting"
-Requires-Dist: boto3; extra == "datahub-usage-feature-reporting"
-Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-feature-reporting"
-Requires-Dist: polars<=1.16.0; extra == "datahub-usage-feature-reporting"
+Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
+Provides-Extra: datahub-usage-reporting
+Requires-Dist: polars==1.19.0; extra == "datahub-usage-reporting"
+Requires-Dist: duckdb; extra == "datahub-usage-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-reporting"
+Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
+Requires-Dist: boto3; extra == "datahub-usage-reporting"
+Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
+Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
+Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
+Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
+Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
+Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
+Provides-Extra: datahub-metadata-sharing
+Requires-Dist: tenacity; extra == "datahub-metadata-sharing"
 Provides-Extra: acryl-cs-issues
-Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
 Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Requires-Dist: openai; extra == "acryl-cs-issues"
 Requires-Dist: zenpy; extra == "acryl-cs-issues"
+Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
 Provides-Extra: all
-Requires-Dist: opensearch-py==2.4.2; extra == "all"
-Requires-Dist: jinja2; extra == "all"
-Requires-Dist: openai; extra == "all"
 Requires-Dist: pandas; extra == "all"
-Requires-Dist: duckdb; extra == "all"
-Requires-Dist: pyarrow<=18.0.0; extra == "all"
+Requires-Dist: boto3; extra == "all"
 Requires-Dist: zenpy; extra == "all"
+Requires-Dist: termcolor==2.5.0; extra == "all"
 Requires-Dist: numpy<2; extra == "all"
+Requires-Dist: pyarrow<=18.0.0; extra == "all"
+Requires-Dist: pydantic<2; extra == "all"
 Requires-Dist: slack-sdk; extra == "all"
+Requires-Dist: polars==1.19.0; extra == "all"
+Requires-Dist: duckdb; extra == "all"
+Requires-Dist: openai; extra == "all"
+Requires-Dist: elasticsearch==7.13.4; extra == "all"
 Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: tenacity; extra == "all"
 Requires-Dist: scipy<=1.14.1; extra == "all"
-Requires-Dist: pydantic<2; extra == "all"
-Requires-Dist: boto3; extra == "all"
-Requires-Dist: elasticsearch==7.13.4; extra == "all"
-Requires-Dist: polars<=1.16.0; extra == "all"
+Requires-Dist: opensearch-py==2.4.2; extra == "all"
+Requires-Dist: jinja2; extra == "all"
 Provides-Extra: dev
-Requires-Dist: opensearch-py==2.4.2; extra == "dev"
-Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: polars==1.19.0; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
 Requires-Dist: openai; extra == "dev"
 Requires-Dist: pandas; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
-Requires-Dist: pyarrow<=18.0.0; extra == "dev"
-Requires-Dist: acryl-datahub[dev]; extra == "dev"
+Requires-Dist: elasticsearch==7.13.4; extra == "dev"
+Requires-Dist: boto3; extra == "dev"
+Requires-Dist: pyarrow; extra == "dev"
+Requires-Dist: tenacity; extra == "dev"
+Requires-Dist: termcolor==2.5.0; extra == "dev"
 Requires-Dist: zenpy; extra == "dev"
 Requires-Dist: numpy<2; extra == "dev"
-Requires-Dist: slack-sdk; extra == "dev"
-Requires-Dist: pyarrow; extra == "dev"
+Requires-Dist: acryl-datahub[dev]; extra == "dev"
 Requires-Dist: scipy<=1.14.1; extra == "dev"
+Requires-Dist: pyarrow<=18.0.0; extra == "dev"
+Requires-Dist: opensearch-py==2.4.2; extra == "dev"
 Requires-Dist: pydantic<2; extra == "dev"
-Requires-Dist: boto3; extra == "dev"
-Requires-Dist: elasticsearch==7.13.4; extra == "dev"
-Requires-Dist: polars<=1.16.0; extra == "dev"
+Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: slack-sdk; extra == "dev"
 Dynamic: provides-extra
 Dynamic: requires-dist

{acryl_datahub_cloud-0.3.7.9.1 → acryl_datahub_cloud-0.3.8}/setup.py

@@ -13,19 +13,25 @@ open_search_common = {"opensearch-py==2.4.2"}
 
 plugins = {
     "datahub-lineage-features": stats_common | open_search_common,
-    "datahub-reporting-forms": stats_common | aws_common,
+    "datahub-reporting-forms": stats_common 
+    | aws_common
+    | {
+        "termcolor==2.5.0",
+    },
     "datahub-reporting-extract-graph": stats_common | aws_common | open_search_common,
     "datahub-reporting-extract-sql": stats_common | aws_common,
-    "datahub-usage-feature-reporting": stats_common
+    "datahub-usage-reporting": stats_common
     | aws_common
     | {
         "opensearch-py==2.4.2",
-        "polars<=1.16.0",
+        "polars==1.19.0",
         "elasticsearch==7.13.4",
         "numpy<2",
         "scipy<=1.14.1",
         "pyarrow<=18.0.0",
+        "termcolor==2.5.0",
     },
+    "datahub-metadata-sharing": {"tenacity"},
     "acryl-cs-issues": {"zenpy", "openai", "jinja2", "slack-sdk"},
 }
 
@@ -39,7 +45,8 @@ dev_requirements = {
             "datahub-reporting-extract-graph",
             "datahub-reporting-extract-sql",
             "datahub-lineage-features",
-            "datahub-usage-feature-reporting",
+            "datahub-usage-reporting",
+            "datahub-metadata-sharing",
             "acryl-cs-issues",
         ]
         for dependency in plugins[plugin]
@@ -63,9 +70,23 @@ setup(
                 "datahub-reporting-extract-sql = acryl_datahub_cloud.datahub_reporting.extract_sql:DataHubReportingExtractSQLSource",
                 "datahub-lineage-features = acryl_datahub_cloud.lineage_features.source:DataHubLineageFeaturesSource",
                 "datahub-usage-reporting = acryl_datahub_cloud.datahub_usage_reporting.usage_feature_reporter:DataHubUsageFeatureReportingSource",
+                "datahub-metadata-sharing = acryl_datahub_cloud.datahub_metadata_sharing.metadata_sharing_source:DataHubMetadataSharingSource",
                 "acryl-cs-issues = acryl_datahub_cloud.acryl_cs_issues.source:AcrylCSIssuesSource",
             ],
         },
+        "include_package_data": True,
+        "package_data": {
+            "acryl_datahub_cloud": [
+                "*.json",
+                "metadata/*.avsc",
+                "metadata/schemas/*.avsc",
+
+            ],
+            "acryl_datahub_cloud.datahub_metadata_sharing": [
+                "scroll_shared_entities.gql",
+                "share_entity.gql",
+            ],
+        },
     },
     extras_require={
         **{plugin: list(dependencies) for (plugin, dependencies) in plugins.items()},
@@ -74,4 +95,4 @@ setup(
         ),
         "dev": list(dev_requirements),
     }
-)
+)

{acryl_datahub_cloud-0.3.7.9.1 → acryl_datahub_cloud-0.3.8}/src/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-0.3.7.9.1 → acryl_datahub_cloud-0.3.8}/src/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-0.3.8/src/acryl_datahub_cloud/api/__init__.py

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

acryl_datahub_cloud-0.3.8/src/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-0.3.8/src/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-0.3.8/src/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-0.3.8/src/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()