acryl-datahub-cloud 0.3.12rc15__tar.gz → 0.3.12.1rc1__tar.gz

{acryl_datahub_cloud-0.3.12rc15 → acryl_datahub_cloud-0.3.12.1rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: acryl-datahub-cloud
-Version: 0.3.12rc15
+Version: 0.3.12.1rc1
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
 Requires-Dist: croniter
@@ -8,90 +8,91 @@ Requires-Dist: pytz
 Requires-Dist: types-croniter
 Provides-Extra: datahub-lineage-features
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
+Requires-Dist: pyarrow; extra == "datahub-lineage-features"
+Requires-Dist: duckdb; extra == "datahub-lineage-features"
+Requires-Dist: tenacity; extra == "datahub-lineage-features"
 Requires-Dist: pandas; extra == "datahub-lineage-features"
 Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
-Requires-Dist: duckdb; extra == "datahub-lineage-features"
-Requires-Dist: pyarrow; extra == "datahub-lineage-features"
 Provides-Extra: datahub-reporting-forms
-Requires-Dist: pandas; extra == "datahub-reporting-forms"
 Requires-Dist: boto3; extra == "datahub-reporting-forms"
+Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
+Requires-Dist: duckdb; extra == "datahub-reporting-forms"
 Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
+Requires-Dist: pandas; extra == "datahub-reporting-forms"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
-Requires-Dist: duckdb; extra == "datahub-reporting-forms"
-Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
 Provides-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: pyarrow; extra == "datahub-reporting-extract-graph"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
-Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
-Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
 Provides-Extra: datahub-usage-reporting
-Requires-Dist: boto3; extra == "datahub-usage-reporting"
 Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
-Requires-Dist: duckdb; extra == "datahub-usage-reporting"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
+Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
+Requires-Dist: boto3; extra == "datahub-usage-reporting"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
-Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-reporting"
 Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
-Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
+Requires-Dist: duckdb; extra == "datahub-usage-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-reporting"
+Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
 Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
 Provides-Extra: datahub-metadata-sharing
 Requires-Dist: tenacity; extra == "datahub-metadata-sharing"
 Provides-Extra: datahub-action-request-owner
 Requires-Dist: tenacity; extra == "datahub-action-request-owner"
 Provides-Extra: acryl-cs-issues
-Requires-Dist: zenpy; extra == "acryl-cs-issues"
+Requires-Dist: openai; extra == "acryl-cs-issues"
 Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
-Requires-Dist: openai; extra == "acryl-cs-issues"
+Requires-Dist: zenpy; extra == "acryl-cs-issues"
 Provides-Extra: datahub-forms-notifications
 Requires-Dist: tenacity; extra == "datahub-forms-notifications"
 Provides-Extra: all
-Requires-Dist: boto3; extra == "all"
-Requires-Dist: tenacity; extra == "all"
 Requires-Dist: scipy<=1.14.1; extra == "all"
-Requires-Dist: duckdb; extra == "all"
-Requires-Dist: jinja2; extra == "all"
+Requires-Dist: pydantic<2; extra == "all"
+Requires-Dist: slack-sdk; extra == "all"
+Requires-Dist: pyarrow<=18.0.0; extra == "all"
 Requires-Dist: pyarrow; extra == "all"
-Requires-Dist: pandas; extra == "all"
 Requires-Dist: elasticsearch==7.13.4; extra == "all"
-Requires-Dist: slack-sdk; extra == "all"
+Requires-Dist: numpy<2; extra == "all"
 Requires-Dist: openai; extra == "all"
 Requires-Dist: zenpy; extra == "all"
+Requires-Dist: tenacity; extra == "all"
+Requires-Dist: termcolor==2.5.0; extra == "all"
+Requires-Dist: jinja2; extra == "all"
+Requires-Dist: boto3; extra == "all"
 Requires-Dist: opensearch-py==2.4.2; extra == "all"
-Requires-Dist: pyarrow<=18.0.0; extra == "all"
+Requires-Dist: duckdb; extra == "all"
+Requires-Dist: pandas; extra == "all"
 Requires-Dist: polars==1.30.0; extra == "all"
-Requires-Dist: termcolor==2.5.0; extra == "all"
-Requires-Dist: pydantic<2; extra == "all"
-Requires-Dist: numpy<2; extra == "all"
 Provides-Extra: dev
 Requires-Dist: openai; extra == "dev"
+Requires-Dist: zenpy; extra == "dev"
+Requires-Dist: scipy<=1.14.1; extra == "dev"
+Requires-Dist: termcolor==2.5.0; extra == "dev"
+Requires-Dist: tenacity; extra == "dev"
+Requires-Dist: pydantic<2; extra == "dev"
+Requires-Dist: jinja2; extra == "dev"
 Requires-Dist: acryl-datahub[dev]; extra == "dev"
 Requires-Dist: boto3; extra == "dev"
-Requires-Dist: tenacity; extra == "dev"
-Requires-Dist: scipy<=1.14.1; extra == "dev"
-Requires-Dist: zenpy; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
 Requires-Dist: opensearch-py==2.4.2; extra == "dev"
 Requires-Dist: pyarrow<=18.0.0; extra == "dev"
 Requires-Dist: pyarrow; extra == "dev"
-Requires-Dist: polars==1.30.0; extra == "dev"
-Requires-Dist: jinja2; extra == "dev"
-Requires-Dist: pandas; extra == "dev"
-Requires-Dist: termcolor==2.5.0; extra == "dev"
-Requires-Dist: pydantic<2; extra == "dev"
 Requires-Dist: elasticsearch==7.13.4; extra == "dev"
-Requires-Dist: numpy<2; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
 Requires-Dist: slack-sdk; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: polars==1.30.0; extra == "dev"
+Requires-Dist: numpy<2; extra == "dev"
 Dynamic: provides-extra
 Dynamic: requires-dist

{acryl_datahub_cloud-0.3.12rc15 → acryl_datahub_cloud-0.3.12.1rc1}/setup.py

@@ -20,7 +20,7 @@ aws_common = {"boto3"}
 open_search_common = {"opensearch-py==2.4.2"}
 
 plugins = {
-    "datahub-lineage-features": stats_common | open_search_common,
+    "datahub-lineage-features": stats_common | open_search_common | {"tenacity"},
     "datahub-reporting-forms": stats_common
     | aws_common
     | {

{acryl_datahub_cloud-0.3.12rc15 → acryl_datahub_cloud-0.3.12.1rc1}/src/acryl_datahub_cloud/_codegen_config.json

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

acryl_datahub_cloud-0.3.12.1rc1/src/acryl_datahub_cloud/lineage_features/source.py

@@ -0,0 +1,356 @@
+import logging
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Iterable, List, Set
+
+from opensearchpy import OpenSearch
+from opensearchpy.exceptions import (
+    ConnectionError as OpenSearchConnectionError,
+    ConnectionTimeout,
+    RequestError,
+    TransportError,
+)
+from pydantic import validator
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from acryl_datahub_cloud.elasticsearch.config import ElasticSearchClientConfig
+from acryl_datahub_cloud.elasticsearch.graph_service import ElasticGraphRow
+from datahub.configuration import ConfigModel
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+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.source_report.ingestion_stage import IngestionStageReport
+from datahub.metadata.schema_classes import AuditStampClass, LineageFeaturesClass
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_ACTOR = "urn:li:corpuser:__datahub_system"
+
+
+class LineageFeaturesSourceConfig(ConfigModel):
+    enabled: bool = True
+    search_index: ElasticSearchClientConfig = ElasticSearchClientConfig()
+    query_timeout: int = 30
+    extract_batch_size: int = 3000
+    max_retries: int = 3
+    retry_delay_seconds: int = 5
+    retry_backoff_multiplier: float = 2.0
+
+    @validator("max_retries")
+    def validate_max_retries(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError("max_retries must be at least 1")
+        return v
+
+    @validator("retry_delay_seconds")
+    def validate_retry_delay_seconds(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError("retry_delay_seconds must be at least 1")
+        return v
+
+    @validator("retry_backoff_multiplier")
+    def validate_retry_backoff_multiplier(cls, v: float) -> float:
+        if v < 1.0:
+            raise ValueError("retry_backoff_multiplier must be at least 1.0")
+        return v
+
+
+@dataclass
+class LineageExtractGraphSourceReport(SourceReport, IngestionStageReport):
+    valid_urns_count: int = 0
+    upstream_count: int = 0
+    downstream_count: int = 0
+    edges_scanned: int = 0
+
+
+@platform_name(id="datahub", platform_name="DataHub")
+@config_class(LineageFeaturesSourceConfig)
+@support_status(SupportStatus.INCUBATING)
+class DataHubLineageFeaturesSource(Source):
+    """
+    DataHub Lineage Features Source that extracts lineage information from Elasticsearch/OpenSearch.
+    """
+
+    platform = "datahub"
+
+    def __init__(
+        self, config: LineageFeaturesSourceConfig, ctx: PipelineContext
+    ) -> None:
+        super().__init__(ctx)
+        self.config: LineageFeaturesSourceConfig = config
+        self.report = LineageExtractGraphSourceReport()
+        self.opened_files: List[str] = []
+
+        self.valid_urns: Set[str] = set()
+        self.upstream_counts: Dict[str, int] = defaultdict(int)
+        self.downstream_counts: Dict[str, int] = defaultdict(int)
+        self.last_print_time = time.time()
+
+    def _get_retry_decorator(
+        self,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Create a retry decorator based on config parameters"""
+
+        def should_retry_exception(exception: Exception) -> bool:
+            """Custom retry predicate for OpenSearch exceptions"""
+            if isinstance(
+                exception,
+                (
+                    OpenSearchConnectionError,
+                    ConnectionTimeout,
+                    RequestError,
+                    TransportError,
+                ),
+            ):
+                return True
+            # Also retry on general connection and timeout errors
+            if isinstance(exception, (ConnectionError, TimeoutError)):
+                return True
+            return False
+
+        return retry(
+            retry=retry_if_exception_type(
+                (
+                    OpenSearchConnectionError,
+                    ConnectionTimeout,
+                    RequestError,
+                    TransportError,
+                    ConnectionError,
+                    TimeoutError,
+                )
+            ),
+            stop=stop_after_attempt(self.config.max_retries),
+            wait=wait_exponential(
+                multiplier=self.config.retry_backoff_multiplier,
+                min=self.config.retry_delay_seconds,
+                max=30,
+            ),
+            before_sleep=before_sleep_log(logger, logging.WARNING),
+            reraise=True,
+        )
+
+    def _create_pit_with_retry(self, server: OpenSearch, index: str) -> str:
+        """Create a Point-in-Time (PIT) with retry logic"""
+        retry_decorator = self._get_retry_decorator()
+
+        @retry_decorator
+        def _create_pit() -> str:
+            logger.debug(f"Creating PIT for index: {index}")
+            response = server.create_pit(index, keep_alive="10m")
+            pit = response.get("pit_id")
+            if not pit:
+                raise Exception("Failed to create PIT - no pit_id returned")
+            logger.debug(f"Successfully created PIT: {pit}")
+            return pit
+
+        return _create_pit()
+
+    def _search_with_retry(
+        self, server: OpenSearch, query: dict, batch_size: int
+    ) -> dict:
+        """Execute search with retry logic"""
+        retry_decorator = self._get_retry_decorator()
+
+        @retry_decorator
+        def _search() -> dict:
+            logger.debug(f"Executing search with batch size: {batch_size}")
+            return server.search(
+                body=query,
+                size=batch_size,
+                params={"timeout": self.config.query_timeout},
+            )
+
+        return _search()
+
+    def _delete_pit_with_retry(self, server: OpenSearch, pit: str) -> None:
+        """Delete Point-in-Time (PIT) with retry logic"""
+        retry_decorator = self._get_retry_decorator()
+
+        @retry_decorator
+        def _delete_pit() -> None:
+            logger.debug(f"Deleting PIT: {pit}")
+            server.delete_pit(body={"pit_id": pit})
+            logger.debug(f"Successfully deleted PIT: {pit}")
+
+        _delete_pit()
+
+    def _create_opensearch_client_with_retry(self) -> OpenSearch:
+        """Create OpenSearch client with retry logic"""
+        retry_decorator = self._get_retry_decorator()
+
+        @retry_decorator
+        def _create_client() -> OpenSearch:
+            logger.debug(
+                f"Creating OpenSearch client for endpoint: {self.config.search_index.endpoint}"
+            )
+            return OpenSearch(
+                [self.config.search_index.endpoint],
+                http_auth=(
+                    self.config.search_index.username,
+                    self.config.search_index.password,
+                ),
+                use_ssl=self.config.search_index.use_ssl,
+            )
+
+        return _create_client()
+
+    def _update_report(self) -> None:
+        """
+        Information to see whether we are close to hitting the memory limits
+        """
+        self.report.valid_urns_count = len(self.valid_urns)
+        self.report.upstream_count = len(self.upstream_counts.keys())
+        self.report.downstream_count = len(self.downstream_counts.keys())
+
+    def _print_report(self) -> None:
+        """
+        Printing is required like this because the report is only printed
+        when the workunits are yielded
+        In case of background processes we won't know the progress if this is not done
+        """
+        time_taken = round(time.time() - self.last_print_time, 1)
+        # Print report every 2 minutes
+        if time_taken > 120:
+            self._update_report()
+            self.last_print_time = time.time()
+            logger.info(f"\n{self.report.as_string()}")
+
+    def process_batch(self, results: Iterable[dict]) -> None:
+        for doc in results:
+            self._print_report()
+            row = ElasticGraphRow.from_elastic_doc(doc["_source"])
+            self.report.edges_scanned += 1
+            if (
+                row.source_urn in self.valid_urns
+                and row.destination_urn in self.valid_urns
+            ):
+                self.upstream_counts[row.source_urn] += 1
+                self.downstream_counts[row.destination_urn] += 1
+
+    def populate_valid_urns(self) -> None:
+        graph = self.ctx.require_graph("Load non soft-deleted urns")
+        for urn in graph.get_urns_by_filter(batch_size=self.config.extract_batch_size):
+            self._print_report()
+            self.valid_urns.add(urn)
+
+    def get_workunits(self) -> Iterable[MetadataWorkUnit]:
+        with self.report.new_stage("Load valid URNs"):
+            self.populate_valid_urns()
+
+        timestamp = datetime.now(tz=timezone.utc)
+        server = self._create_opensearch_client_with_retry()
+
+        query = {
+            "query": {
+                "bool": {
+                    "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"}},
+            ],
+        }
+
+        index = f"{self.config.search_index.index_prefix}graph_service_v1"
+        pit = self._create_pit_with_retry(server, index)
+
+        # TODO: Save PIT, we can resume processing based on <pit, search_after> tuple
+        query.update({"pit": {"id": pit, "keep_alive": "10m"}})
+
+        # TODO: Using slicing we can parallelize the ES calls below:
+        # https://opensearch.org/docs/latest/search-plugins/searching-data/point-in-time/#search-slicing
+        batch_size = self.config.extract_batch_size
+        with self.report.new_stage("Extract lineage features"):
+            try:
+                while True:
+                    results = self._search_with_retry(server, query, batch_size)
+                    self.process_batch(results["hits"]["hits"])
+                    if len(results["hits"]["hits"]) < batch_size:
+                        break
+                    query.update({"search_after": results["hits"]["hits"][-1]["sort"]})
+            except Exception as e:
+                logger.error(f"Error during lineage extraction: {e}")
+                self.report.report_failure(
+                    title="Lineage extraction failed",
+                    message="Failed to extract lineage features from Elasticsearch",
+                    context=f"Error: {str(e)}",
+                    exc=e,
+                )
+                # Ensure PIT is cleaned up even on error
+                try:
+                    self._delete_pit_with_retry(server, pit)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Failed to cleanup PIT after error: {cleanup_error}"
+                    )
+                raise
+        # So previous stage's calculations are done
+        self.report.new_stage("Extract lineage features End")
+        self._update_report()
+        self._delete_pit_with_retry(server, pit)
+
+        # In Python 3.9, can be replaced by `self.self.upstream_counts.keys() | self.downstream_counts.keys()`
+        for urn in set(self.upstream_counts.keys()).union(
+            self.downstream_counts.keys()
+        ):
+            logger.debug(
+                f"{urn}: {self.upstream_counts[urn]}, {self.downstream_counts[urn]}"
+            )
+            wu = MetadataChangeProposalWrapper(
+                entityUrn=urn,
+                aspect=LineageFeaturesClass(
+                    upstreamCount=self.upstream_counts[urn],
+                    downstreamCount=self.downstream_counts[urn],
+                    computedAt=AuditStampClass(
+                        time=int(timestamp.timestamp() * 1000),
+                        actor=SYSTEM_ACTOR,
+                    ),
+                ),
+            ).as_workunit()
+            self.report.report_workunit(wu)
+            yield wu
+
+    def get_report(self) -> SourceReport:
+        return self.report
+
+    def close(self) -> None:
+        for file in self.opened_files:
+            os.remove(file)
+        return super().close()