acryl-datahub-cloud 0.3.6.9rc1__tar.gz → 0.3.7rc1__tar.gz

{acryl_datahub_cloud-0.3.6.9rc1 → acryl_datahub_cloud-0.3.7rc1}/PKG-INFO

@@ -1,77 +1,77 @@
 Metadata-Version: 2.1
 Name: acryl-datahub-cloud
-Version: 0.3.6.9rc1
+Version: 0.3.7rc1
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
 Provides-Extra: datahub-lineage-features
-Requires-Dist: pyarrow; extra == "datahub-lineage-features"
 Requires-Dist: pandas; extra == "datahub-lineage-features"
-Requires-Dist: duckdb; extra == "datahub-lineage-features"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
 Requires-Dist: pydantic<2; 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"
 Provides-Extra: datahub-reporting-forms
-Requires-Dist: pyarrow; 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"
+Requires-Dist: duckdb; extra == "datahub-reporting-forms"
 Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Provides-Extra: datahub-reporting-extract-graph
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
-Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pydantic<2; 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: boto3; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
-Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
+Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
 Provides-Extra: datahub-usage-feature-reporting
-Requires-Dist: pyarrow; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-feature-reporting"
-Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-feature-reporting"
-Requires-Dist: duckdb; extra == "datahub-usage-feature-reporting"
+Requires-Dist: scipy; extra == "datahub-usage-feature-reporting"
 Requires-Dist: numpy<2; extra == "datahub-usage-feature-reporting"
+Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-feature-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-feature-reporting"
+Requires-Dist: pydantic<2; extra == "datahub-usage-feature-reporting"
 Requires-Dist: polars; extra == "datahub-usage-feature-reporting"
-Requires-Dist: scipy; extra == "datahub-usage-feature-reporting"
+Requires-Dist: pyarrow; extra == "datahub-usage-feature-reporting"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-feature-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-feature-reporting"
+Requires-Dist: duckdb; extra == "datahub-usage-feature-reporting"
 Requires-Dist: boto3; extra == "datahub-usage-feature-reporting"
 Provides-Extra: acryl-cs-issues
-Requires-Dist: jinja2; extra == "acryl-cs-issues"
-Requires-Dist: zenpy; extra == "acryl-cs-issues"
-Requires-Dist: openai; 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"
+Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Provides-Extra: all
-Requires-Dist: pyarrow; extra == "all"
-Requires-Dist: pandas; extra == "all"
+Requires-Dist: scipy; extra == "all"
 Requires-Dist: jinja2; extra == "all"
-Requires-Dist: zenpy; extra == "all"
-Requires-Dist: elasticsearch==7.13.4; extra == "all"
 Requires-Dist: duckdb; extra == "all"
 Requires-Dist: numpy<2; extra == "all"
-Requires-Dist: slack-sdk; extra == "all"
-Requires-Dist: polars; extra == "all"
 Requires-Dist: openai; extra == "all"
-Requires-Dist: scipy; extra == "all"
-Requires-Dist: opensearch-py==2.4.2; extra == "all"
+Requires-Dist: elasticsearch==7.13.4; extra == "all"
+Requires-Dist: zenpy; extra == "all"
+Requires-Dist: pandas; extra == "all"
 Requires-Dist: pydantic<2; extra == "all"
+Requires-Dist: polars; extra == "all"
+Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: slack-sdk; extra == "all"
+Requires-Dist: opensearch-py==2.4.2; extra == "all"
 Requires-Dist: boto3; extra == "all"
 Provides-Extra: dev
-Requires-Dist: pyarrow; extra == "dev"
-Requires-Dist: pandas; extra == "dev"
+Requires-Dist: scipy; extra == "dev"
 Requires-Dist: jinja2; extra == "dev"
-Requires-Dist: zenpy; extra == "dev"
-Requires-Dist: elasticsearch==7.13.4; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
-Requires-Dist: acryl-datahub[dev]; extra == "dev"
 Requires-Dist: numpy<2; extra == "dev"
-Requires-Dist: slack-sdk; extra == "dev"
-Requires-Dist: polars; extra == "dev"
 Requires-Dist: openai; extra == "dev"
-Requires-Dist: scipy; extra == "dev"
-Requires-Dist: opensearch-py==2.4.2; extra == "dev"
+Requires-Dist: elasticsearch==7.13.4; extra == "dev"
+Requires-Dist: zenpy; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
 Requires-Dist: pydantic<2; extra == "dev"
+Requires-Dist: opensearch-py==2.4.2; extra == "dev"
+Requires-Dist: pyarrow; extra == "dev"
+Requires-Dist: acryl-datahub[dev]; extra == "dev"
+Requires-Dist: polars; extra == "dev"
+Requires-Dist: slack-sdk; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
 Requires-Dist: boto3; extra == "dev"

{acryl_datahub_cloud-0.3.6.9rc1 → acryl_datahub_cloud-0.3.7rc1}/setup.py

@@ -18,7 +18,7 @@ plugins = {
     "datahub-reporting-extract-sql": stats_common | aws_common,
     "datahub-usage-feature-reporting": stats_common
     | aws_common
-    | {"opensearch-py==2.4.2", "polars", "elasticsearch==7.13.4", "numpy<2", "scipy"},
+    | {"opensearch-py==2.4.2", "polars", "elasticsearch==7.13.4", "numpy<2", "scipy", "pyarrow"},
     "acryl-cs-issues": {"zenpy", "openai", "jinja2", "slack-sdk"},
 }
 

{acryl_datahub_cloud-0.3.6.9rc1 → acryl_datahub_cloud-0.3.7rc1}/src/acryl_datahub_cloud/_codegen_config.json

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

{acryl_datahub_cloud-0.3.6.9rc1 → acryl_datahub_cloud-0.3.7rc1}/src/acryl_datahub_cloud/datahub_reporting/extract_graph.py

@@ -160,7 +160,7 @@ class DataHubReportingExtractGraphSource(Source):
                 ],
             }
 
-            index = f"{index_prefix}_graph_service_v1"
+            index = f"{index_prefix}graph_service_v1"
             response = server.create_pit(index, keep_alive="10m")
 
             # TODO: Save PIT, we can resume processing based on <pit, search_after> tuple

acryl_datahub_cloud-0.3.7rc1/src/acryl_datahub_cloud/datahub_usage_reporting/usage_feature_patch_builder.py

@@ -0,0 +1,466 @@
+from typing import Callable, Dict, List, Optional, Union
+
+from acryl_datahub_cloud.metadata.schema_classes import UsageFeaturesClass
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import KafkaAuditHeaderClass, SystemMetadataClass
+
+
+class UsageFeaturePatchBuilder(MetadataPatchProposal):
+    def __init__(
+        self,
+        urn: str,
+        system_metadata: Optional[SystemMetadataClass] = None,
+        audit_header: Optional[KafkaAuditHeaderClass] = None,
+    ) -> None:
+        """Initializes a MonitorPatchBuilder instance.
+
+        Args:
+            urn: The URN of the monitor.
+            system_metadata: The system metadata of the data job (optional).
+            audit_header: The Kafka audit header of the data job (optional).
+        """
+        super().__init__(
+            urn, system_metadata=system_metadata, audit_header=audit_header
+        )
+
+    def setQueryCountLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the query count for the last 30 days.
+
+        Args:
+            count: The query count for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/queryCountLast30Days",
+            value=count,
+        )
+        return self
+
+    def setViewCountLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the view count for the last 30 days.
+
+        Args:
+            count: The view count for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/viewCountLast30Days",
+            value=count,
+        )
+        return self
+
+    def setViewCountTotal(self, count: Optional[int]) -> "UsageFeaturePatchBuilder":
+        """Sets the total view count.
+
+        Args:
+            count: The total view count.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/viewCountTotal",
+            value=count,
+        )
+        return self
+
+    def setViewCountPercentileLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the view count percentile for the last 30 days.
+
+        Args:
+            count: The view count percentile for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/viewCountPercentileLast30Days",
+            value=count,
+        )
+        return self
+
+    def setUsageCountLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the usage count for the last 30 days.
+
+        Args:
+            count: The usage count for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/usageCountLast30Days",
+            value=count,
+        )
+        return self
+
+    def setUniqueUserCountLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the unique user count for the last 30 days.
+
+        Args:
+            count: The unique user count for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/uniqueUserCountLast30Days",
+            value=count,
+        )
+        return self
+
+    def setWriteCountLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the write count for the last 30 days.
+
+        Args:
+            count: The write count for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/writeCountLast30Days",
+            value=count,
+        )
+        return self
+
+    def setQueryCountPercentileLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the query count percentile for the last 30 days.
+
+        Args:
+            count: The query count percentile for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/queryCountPercentileLast30Days",
+            value=count,
+        )
+        return self
+
+    def setQueryCountRankLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the query count rank for the last 30 days.
+
+        Args:
+            count: The query count rank for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/queryCountRankLast30Days",
+            value=count,
+        )
+        return self
+
+    def setUniqueUserPercentileLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the unique user percentile for the last 30 days.
+
+        Args:
+            count: The unique user percentile for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/uniqueUserPercentileLast30Days",
+            value=count,
+        )
+        return self
+
+    def setUniqueUserRankLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the unique user rank for the last 30 days.
+
+        Args:
+            count: The unique user rank for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/uniqueUserRankLast30Days",
+            value=count,
+        )
+        return self
+
+    def setWriteCountPercentileLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the write count percentile for the last 30 days.
+
+        Args:
+            count: The write count percentile for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/writeCountPercentileLast30Days",
+            value=count,
+        )
+        return self
+
+    def setWriteCountRankLast30Days(
+        self, count: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the write count rank for the last 30 days.
+
+        Args:
+            count: The write count rank for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/writeCountRankLast30Days",
+            value=count,
+        )
+        return self
+
+    def setTopUsersLast30Days(
+        self, users: Optional[List[str]]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the top users for the last 30 days.
+
+        Args:
+            users: The list of top users for the last 30 days.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/topUsersLast30Days",
+            value=users,
+        )
+        return self
+
+    def setSizeInBytesPercentile(
+        self, percentile: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the size in bytes percentile.
+
+        Args:
+            percentile: The size in bytes percentile.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/sizeInBytesPercentile",
+            value=percentile,
+        )
+        return self
+
+    def setSizeInBytesRank(self, rank: Optional[int]) -> "UsageFeaturePatchBuilder":
+        """Sets the size in bytes rank.
+
+        Args:
+            rank: The size in bytes rank.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/sizeInBytesRank",
+            value=rank,
+        )
+        return self
+
+    def setRowCountPercentile(
+        self, percentile: Optional[int]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the row count percentile.
+
+        Args:
+            percentile: The row count percentile.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/rowCountPercentile",
+            value=percentile,
+        )
+        return self
+
+    def setUsageSearchScoreMultiplier(
+        self, multiplier: Optional[float]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the usage search score multiplier.
+
+        Args:
+            multiplier: The usage search score multiplier.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/usageSearchScoreMultiplier",
+            value=multiplier,
+        )
+        return self
+
+    def setUsageFreshnessScoreMultiplier(
+        self, multiplier: Optional[float]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the usage freshness score multiplier.
+
+        Args:
+            multiplier: The usage freshness score multiplier.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/usageFreshnessScoreMultiplier",
+            value=multiplier,
+        )
+        return self
+
+    def setCustomDatahubScoreMultiplier(
+        self, multiplier: Optional[float]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the custom DataHub score multiplier.
+
+        Args:
+            multiplier: The custom DataHub score multiplier.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/customDatahubScoreMultiplier",
+            value=multiplier,
+        )
+        return self
+
+    def setCombinedSearchRankingMultiplier(
+        self, multiplier: Optional[float]
+    ) -> "UsageFeaturePatchBuilder":
+        """Sets the combined search ranking multiplier.
+
+        Args:
+            multiplier: The combined search ranking multiplier.
+
+        Returns:
+            UsageFeaturePatchBuilder: The instance of the builder.
+        """
+        self._add_patch(
+            UsageFeaturesClass.ASPECT_NAME,
+            "add",
+            path="/combinedSearchRankingMultiplier",
+            value=multiplier,
+        )
+        return self
+
+    def apply_usage_features(
+        self, usage_features: UsageFeaturesClass
+    ) -> "UsageFeaturePatchBuilder":
+        """
+        Applies the usage features to the patch builder.
+        :param usage_features: UsageFeaturesClass object containing the features to apply
+        :return: self (UsageFeaturePatchBuilder instance)
+        """
+        feature_setters: Dict[
+            str,
+            Union[
+                Callable[[Optional[int]], UsageFeaturePatchBuilder],
+                Callable[[Optional[float]], UsageFeaturePatchBuilder],
+                Callable[[Optional[List[str]]], UsageFeaturePatchBuilder],
+            ],
+        ] = {
+            "queryCountLast30Days": self.setQueryCountLast30Days,
+            "viewCountLast30Days": self.setViewCountLast30Days,
+            "viewCountTotal": self.setViewCountTotal,
+            "viewCountPercentileLast30Days": self.setViewCountPercentileLast30Days,
+            "usageCountLast30Days": self.setUsageCountLast30Days,
+            "uniqueUserCountLast30Days": self.setUniqueUserCountLast30Days,
+            "writeCountLast30Days": self.setWriteCountLast30Days,
+            "queryCountPercentileLast30Days": self.setQueryCountPercentileLast30Days,
+            "queryCountRankLast30Days": self.setQueryCountRankLast30Days,
+            "uniqueUserPercentileLast30Days": self.setUniqueUserPercentileLast30Days,
+            "uniqueUserRankLast30Days": self.setUniqueUserRankLast30Days,
+            "writeCountPercentileLast30Days": self.setWriteCountPercentileLast30Days,
+            "writeCountRankLast30Days": self.setWriteCountRankLast30Days,
+            "topUsersLast30Days": self.setTopUsersLast30Days,
+            "sizeInBytesPercentile": self.setSizeInBytesPercentile,
+            "sizeInBytesRank": self.setSizeInBytesRank,
+            "rowCountPercentile": self.setRowCountPercentile,
+            "usageSearchScoreMultiplier": self.setUsageSearchScoreMultiplier,
+            "usageFreshnessScoreMultiplier": self.setUsageFreshnessScoreMultiplier,
+            "customDatahubScoreMultiplier": self.setCustomDatahubScoreMultiplier,
+            "combinedSearchRankingMultiplier": self.setCombinedSearchRankingMultiplier,
+        }
+
+        for feature, setter in feature_setters.items():
+            value = getattr(usage_features, feature)
+            if value is not None:
+                assert value is not None
+                setter(value)
+
+        return self