acryl-datahub-cloud 0.3.10.4__py3-none-any.whl → 0.3.11__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

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

acryl_datahub_cloud/_sdk_extras/__init__.py

@@ -0,0 +1,4 @@
+from acryl_datahub_cloud._sdk_extras.assertion import Assertion
+from acryl_datahub_cloud._sdk_extras.assertions_client import AssertionsClient
+
+__all__ = ["Assertion", "AssertionsClient"]

acryl_datahub_cloud/_sdk_extras/assertion.py

@@ -0,0 +1,15 @@
+"""
+This file contains the Assertion class, which is used to represent an assertion in DataHub.
+
+The Assertion class is currently not implemented, this is a placeholder for future implementation.
+"""
+
+from typing import Union
+
+from datahub.metadata.urns import AssertionUrn
+
+
+class Assertion:
+    def __init__(self, urn: Union[str, AssertionUrn]):
+        print(f"The Assertion class is currently not implemented. Urn provided: {urn}")
+        self.urn = AssertionUrn(urn)

acryl_datahub_cloud/_sdk_extras/assertions_client.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Union
+
+from acryl_datahub_cloud._sdk_extras.assertion import Assertion
+from datahub.metadata.urns import AssertionUrn
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+
+class AssertionsClient:
+    def __init__(self, client: DataHubClient):
+        self.client = client
+
+    def get_assertions(
+        self, urn: Union[str, list[str], AssertionUrn, list[AssertionUrn]]
+    ) -> list[Assertion]:
+        print(
+            "get_assertions is not implemented, this is a placeholder. Returning empty list."
+        )
+        print(f"urn provided: {urn}")
+        return []

acryl_datahub_cloud/acryl_cs_issues/acryl_customer.py

@@ -219,6 +219,7 @@ class AcrylCustomer:
         self._emitted_containers: Dict[str, bool] = {}
 
     def _get_owner_from_assignee(self, assignee: ExternalUser) -> str:
+        assert assignee.email, "Assignee must have an email"
         owner_urn_options = [
             urn
             for urn in self.graph.get_urns_by_filter(
@@ -233,7 +234,6 @@ class AcrylCustomer:
             )
         ]
         if not owner_urn_options:
-            assert assignee.email, "Assignee must have an email"
             owner_urn = "urn:li:corpuser:" + assignee.email
             self.graph.emit(
                 MetadataChangeProposalWrapper(

acryl_datahub_cloud/action_request/action_request_owner_source.py

@@ -2,6 +2,7 @@ import logging
 from typing import Dict, Iterable, List, Optional
 
 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,
@@ -14,10 +15,8 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
     ActionRequestInfoClass,
 )
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.utilities.urns.urn import guess_entity_type
 
-
 logger = logging.getLogger(__name__)
 
 

acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -163,7 +163,7 @@ class DataHubBasedS3Dataset:
                     self.schema = row.arrow_schema()
                 else:
                     # hail mary: infer schema from the first row and cast everything to string
-                    self.schema = pa.schema([(key, pa.string()) for key in row.keys()])
+                    self.schema = pa.schema([(key, pa.string()) for key in row])
                     self.stringify_row = True
 
             self._initialize_local_file()
@@ -172,7 +172,7 @@ class DataHubBasedS3Dataset:
                 self.schema,
                 compression=self.config.file_compression,
             )
-        if isinstance(row, BaseModel) or isinstance(row, BaseModelRow):
+        if isinstance(row, (BaseModel, BaseModelRow)):
             # for anything extending BaseModel, we want to use the dict representation
             write_row: Dict[str, Any] = row.dict()
         elif isinstance(row, dict):
@@ -274,11 +274,7 @@ class DataHubBasedS3Dataset:
         self, duckdb_columns: List[Tuple[str, str]]
     ) -> SchemaMetadataClass:
         def get_type_from_dtype(dtype: str) -> SchemaFieldDataTypeClass:
-            if "int" in dtype:
-                return SchemaFieldDataTypeClass(type=NumberTypeClass())
-            elif "float" in dtype:
-                return SchemaFieldDataTypeClass(type=NumberTypeClass())
-            elif "number" in dtype:
+            if "int" in dtype or "float" in dtype or "number" in dtype:
                 return SchemaFieldDataTypeClass(type=NumberTypeClass())
             elif "bool" in dtype:
                 return SchemaFieldDataTypeClass(type=BooleanTypeClass())

acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -5,6 +5,9 @@ from enum import Enum
 from typing import Any, Callable, Dict, Iterable, List, Optional
 
 import pandas as pd
+from pydantic import BaseModel
+
+from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.graph.filters import RawSearchFilterRule
@@ -16,9 +19,6 @@ from datahub.metadata.schema_classes import (
     FormStateClass,
     FormTypeClass,
 )
-from pydantic import BaseModel
-
-from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
 
 logger = logging.getLogger(__name__)
 
@@ -257,6 +257,7 @@ class DataHubFormReportingData(FormData):
                 for prompt_id, response_time in zip(
                     search_row.completedFormsCompletedPromptIds,
                     search_row.completedFormsCompletedPromptResponseTimes,
+                    strict=False,
                 )
                 if prompt_id in form_prompts
             }
@@ -289,9 +290,10 @@ class DataHubFormReportingData(FormData):
         on_asset_scanned: Optional[Callable[[str], Any]] = None,
         on_form_scanned: Optional[Callable[[str], Any]] = None,
     ) -> Iterable[FormReportingRow]:
-        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__.keys()]
+        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__]
+        # TODO: Replace with the new search/filter SDK.
         result = self.graph.get_results_by_filter(
-            extra_or_filters=self.get_form_existence_or_filters(),
+            extra_or_filters=[{"and": self.get_form_existence_or_filters()}],
             extra_source_fields=extra_fields,
             skip_cache=True,
         )
@@ -388,6 +390,7 @@ class DataHubFormReportingData(FormData):
                     for (p, p_response_time) in zip(
                         search_row.incompleteFormsCompletedPromptIds,
                         search_row.incompleteFormsCompletedPromptResponseTimes,
+                        strict=False,
                     )
                     if p in form_prompts
                 ]:
@@ -485,6 +488,7 @@ class DataHubFormReportingData(FormData):
                     for (p, p_response_time) in zip(
                         search_row.completedFormsCompletedPromptIds,
                         search_row.completedFormsCompletedPromptResponseTimes,
+                        strict=False,
                     )
                     if p in form_prompts
                 ]:

acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -395,18 +395,10 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                     "last_modified_at": (
                         doc["_source"]["lastModifiedAt"]
                         if "lastModifiedAt" in doc["_source"]
-                        else (
-                            doc["_source"]["lastModifiedAt"]
-                            if "lastModifiedAt" in doc["_source"]
-                            else None
-                        )
+                        else (doc["_source"].get("lastModifiedAt", None))
                     ),
                     "platform": doc["_source"]["platform"],
-                    "removed": (
-                        doc["_source"]["removed"]
-                        if "removed" in doc["_source"]
-                        else False
-                    ),
+                    "removed": (doc["_source"].get("removed", False)),
                 }
 
             time_taken = timer.elapsed_seconds()
@@ -509,11 +501,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                         "eventGranularity": doc["_source"].get("eventGranularity"),
                         "totalSqlQueries": doc["_source"].get("totalSqlQueries", 0),
                         "uniqueUserCount": doc["_source"].get("uniqueUserCount", 0),
-                        "userCounts": (
-                            doc["_source"]["event"]["userCounts"]
-                            if "userCounts" in doc["_source"]["event"]
-                            else None
-                        ),
+                        "userCounts": (doc["_source"]["event"].get("userCounts", None)),
                         "platform": platform,
                     }
                 except KeyError as e:
@@ -525,7 +513,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
             time_taken = timer.elapsed_seconds()
             logger.info(f"DatasetUsage processing took {time_taken:.3f} seconds")
 
-    def search_score(
+    def search_score(  # noqa: C901
         self, urn: str, last_update_time: int, usage_percentile: int
     ) -> SearchRankingMultipliers:
         usage_search_score_multiplier = 1.0
@@ -622,10 +610,10 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                         [endpoint],
                         http_auth=(user, password),
                         use_ssl=(
-                            True
-                            if self.config.search_index
-                            and self.config.search_index.use_ssl
-                            else False
+                            bool(
+                                self.config.search_index
+                                and self.config.search_index.use_ssl
+                            )
                         ),
                     )
 
@@ -639,10 +627,10 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                         [endpoint],
                         http_auth=(user, password),
                         use_ssl=(
-                            True
-                            if self.config.search_index
-                            and self.config.search_index.use_ssl
-                            else False
+                            bool(
+                                self.config.search_index
+                                and self.config.search_index.use_ssl
+                            )
                         ),
                     )
 
@@ -737,7 +725,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                 polars.Duration(): pa.duration("ns"),
             }
 
-            if polars_dtype in [type(key) for key in type_mapping.keys()]:
+            if polars_dtype in [type(key) for key in type_mapping]:
                 return type_mapping[polars_dtype]
             elif polars_dtype == polars.Categorical:
                 return pa.dictionary(index_type=pa.int32(), value_type=pa.string())
@@ -1006,12 +994,9 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
     def generate_mcp_from_lazyframe(
         self, lazy_frame: polars.LazyFrame
     ) -> Iterable[MetadataWorkUnit]:
-        num = 0
         for row in lazy_frame.collect(
             streaming=self.config.experimental_full_streaming
         ).to_struct():
-            num += 1
-
             if "siblings" in row and row["siblings"]:
                 logger.info(f"Siblings found for urn: {row['urn']} -> row['siblings']")
 
@@ -1101,10 +1086,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
     def generate_query_usage_mcp_from_lazyframe(
         self, lazy_frame: polars.LazyFrame
     ) -> Iterable[MetadataWorkUnit]:
-        num = 0
         for row in lazy_frame.collect().iter_rows(named=True):
-            num += 1
-
             query_usage_features = QueryUsageFeaturesClass(
                 queryCountLast30Days=int(row.get("totalSqlQueries", 0) or 0),
                 queryCountTotal=None,  # This is not implemented
@@ -1287,7 +1269,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                 .is_not_null()
                 # We only want to downrank datasets that have a search score multiplier greater than 1. 1 is the minimum score of a dataset
                 .and_(polars.col("combinedSearchRankingMultiplier").ne(1))
-            )  # noqa: E712
+            )
             .filter(polars.col("removed") == False)  # noqa: E712
             .drop(["removed"])
             .drop(["last_modified_at"])