acryl-datahub-cloud 0.3.7.7rc6__tar.gz → 0.3.7.7rc8__tar.gz

{acryl_datahub_cloud-0.3.7.7rc6 → acryl_datahub_cloud-0.3.7.7rc8}/PKG-INFO

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

{acryl_datahub_cloud-0.3.7.7rc6 → acryl_datahub_cloud-0.3.7.7rc8}/src/acryl_datahub_cloud/_codegen_config.json

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

acryl_datahub_cloud-0.3.7.7rc8/src/acryl_datahub_cloud/datahub_usage_reporting/query_builder.py

@@ -0,0 +1,163 @@
+from typing import Dict
+
+
+class QueryBuilder:
+    @staticmethod
+    def get_soft_deleted_entities_query() -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}],
+        }
+
+    @staticmethod
+    def get_query_entities_query() -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "filter": {
+                        "bool": {
+                            "must_not": [
+                                {"term": {"source": "MANUAL"}},
+                            ]
+                        }
+                    }
+                }
+            },
+        }
+
+    @staticmethod
+    def get_upstreams_query() -> Dict:
+        return {
+            "sort": [{"destination.urn": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "must": [
+                        {"terms": {"destination.entityType": ["dataset"]}},
+                        {"terms": {"source.entityType": ["dataset"]}},
+                    ]
+                }
+            },
+        }
+
+    @staticmethod
+    def get_dashboard_usage_query(days: int) -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "filter": {
+                        "bool": {
+                            "must": [
+                                {
+                                    "range": {
+                                        "@timestamp": {
+                                            "gte": f"now-{days}d",
+                                            "lt": "now/d",
+                                        }
+                                    }
+                                },
+                                {"term": {"isExploded": False}},
+                            ]
+                        }
+                    }
+                }
+            },
+        }
+
+    @staticmethod
+    def get_dataset_usage_query(days: int) -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "filter": {
+                        "bool": {
+                            "must": [
+                                {
+                                    "range": {
+                                        "@timestamp": {
+                                            "gte": f"now-{days}d/d",
+                                            "lt": "now/d",
+                                        }
+                                    }
+                                },
+                                {"term": {"isExploded": False}},
+                                {"range": {"totalSqlQueries": {"gt": 0}}},
+                            ]
+                        }
+                    }
+                }
+            },
+        }
+
+    @staticmethod
+    def get_dataset_write_usage_raw_query(days: int) -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}, {"@timestamp": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "range": {
+                                "@timestamp": {"gte": f"now-{days}d/d", "lte": "now/d"}
+                            }
+                        },
+                        {"terms": {"operationType": ["INSERT", "UPDATE", "CREATE"]}},
+                    ]
+                }
+            },
+            "_source": {
+                "includes": ["urn", "@timestamp"],
+            },
+        }
+
+    @staticmethod
+    def get_dataset_write_usage_composite_query(days: int) -> Dict:
+        return {
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "range": {
+                                "@timestamp": {"gte": f"now-{days}d/d", "lte": "now/d"}
+                            }
+                        },
+                        {"terms": {"operationType": ["INSERT", "UPDATE", "CREATE"]}},
+                    ]
+                }
+            },
+            "aggs": {
+                "urn_count": {
+                    "composite": {
+                        "sources": [
+                            {"dataset_operationaspect_v1": {"terms": {"field": "urn"}}}
+                        ]
+                    }
+                }
+            },
+        }
+
+    @staticmethod
+    def get_query_usage_query(days: int) -> Dict:
+        return {
+            "sort": [{"urn": {"order": "asc"}}],
+            "query": {
+                "bool": {
+                    "filter": {
+                        "bool": {
+                            "must": [
+                                {
+                                    "range": {
+                                        "@timestamp": {
+                                            "gte": f"now-{days}d/d",
+                                            "lt": "now/d",
+                                        }
+                                    }
+                                },
+                                {"term": {"isExploded": False}},
+                            ]
+                        }
+                    }
+                }
+            },
+        }

{acryl_datahub_cloud-0.3.7.7rc6 → acryl_datahub_cloud-0.3.7.7rc8}/src/acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -20,6 +20,7 @@ from opensearchpy import OpenSearch
 from pydantic import Field
 from scipy.stats import expon
 
+from acryl_datahub_cloud.datahub_usage_reporting.query_builder import QueryBuilder
 from acryl_datahub_cloud.datahub_usage_reporting.usage_feature_patch_builder import (
     UsageFeaturePatchBuilder,
 )
@@ -58,119 +59,6 @@ platform_regexp = re.compile(r"urn:li:dataset:\(urn:li:dataPlatform:(.+?),.*")
 dashboard_chart_platform_regexp = re.compile(r"urn:li:(?:dashboard|chart):\((.+?),.*")
 dbt_platform_regexp = re.compile(r"urn:li:dataset:\(urn:li:dataPlatform:dbt,.*\)")
 
-GET_SOFT_DELETED_ENTITIES = {
-    "sort": [{"urn": {"order": "asc"}}],
-}
-
-GET_QUERY_ENTITIES = {
-    "sort": [{"urn": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "filter": {
-                "bool": {
-                    "must_not": [
-                        {"term": {"source": "MANUAL"}},
-                    ]
-                }
-            }
-        }
-    },
-}
-
-GET_UPSTREAMS = {
-    "sort": [{"destination.urn": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "must": [
-                {"terms": {"destination.entityType": ["dataset"]}},
-                {"terms": {"source.entityType": ["dataset"]}},
-            ]
-        }
-    },
-}
-
-GET_DASHBOARD_USAGE_QUERY = {
-    "sort": [{"urn": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "filter": {
-                "bool": {
-                    "must": [
-                        {"range": {"@timestamp": {"gte": "now-30d", "lt": "now/d"}}},
-                        {"term": {"isExploded": False}},
-                    ]
-                }
-            }
-        }
-    },
-}
-
-GET_DATASET_USAGE_QUERY = {
-    "sort": [{"urn": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "filter": {
-                "bool": {
-                    "must": [
-                        {"range": {"@timestamp": {"gte": "now-30d/d", "lt": "now/d"}}},
-                        {"term": {"isExploded": False}},
-                        {"range": {"totalSqlQueries": {"gt": 0}}},
-                    ]
-                }
-            }
-        }
-    },
-}
-
-DATASET_WRITE_USAGE_RAW_QUERY = {
-    "sort": [{"urn": {"order": "asc"}}, {"@timestamp": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "must": [
-                {"range": {"@timestamp": {"gte": "now-30d/d", "lte": "now/d"}}},
-                {"terms": {"operationType": ["INSERT", "UPDATE", "CREATE"]}},
-            ]
-        }
-    },
-    "_source": {
-        "includes": ["urn", "@timestamp"],
-    },
-}
-
-DATASET_WRITE_USAGE_COMPOSITE_QUERY = {
-    "query": {
-        "bool": {
-            "must": [
-                {"range": {"@timestamp": {"gte": "now-30d/d", "lte": "now/d"}}},
-                {"terms": {"operationType": ["INSERT", "UPDATE", "CREATE"]}},
-            ]
-        }
-    },
-    "aggs": {
-        "urn_count": {
-            "composite": {
-                "sources": [{"dataset_operationaspect_v1": {"terms": {"field": "urn"}}}]
-            }
-        }
-    },
-}
-
-GET_QUERY_USAGE_QUERY = {
-    "sort": [{"urn": {"order": "asc"}}],
-    "query": {
-        "bool": {
-            "filter": {
-                "bool": {
-                    "must": [
-                        {"range": {"@timestamp": {"gte": "now-30d/d", "lt": "now/d"}}},
-                        {"term": {"isExploded": False}},
-                    ]
-                }
-            }
-        }
-    },
-}
-
 
 class S3ClientConfig(ConfigModel):
     bucket: str = os.getenv("DATA_BUCKET", "")
@@ -206,7 +94,13 @@ class RankingPolicy(ConfigModel):
     regexp_based_factors: List[RegexpFactor] = []
 
 
-class DataHubUsageFeatureReportingSourceConfig(StatefulIngestionConfigBase):
+class DataHubUsageFeatureReportingSourceConfig(
+    ConfigModel, StatefulIngestionConfigBase
+):
+    lookback_days: int = Field(
+        30, description="Number of days to look back for usage data."
+    )
+
     server: Optional[DatahubClientConfig] = Field(
         None, description="Optional configuration for the DataHub server connection."
     )
@@ -882,7 +776,9 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         if self.config.streaming_mode:
             wdf = self.load_es_data_to_lf(
                 index="dataset_operationaspect_v1",
-                query=DATASET_WRITE_USAGE_RAW_QUERY,
+                query=QueryBuilder.get_dataset_write_usage_raw_query(
+                    self.config.lookback_days
+                ),
                 read_function=self.write_stat_raw_batch,
                 schema={"urn": polars.Categorical, "platform": polars.Categorical},
             )
@@ -891,7 +787,9 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
             wdf = polars.LazyFrame(
                 self.load_data_from_es(
                     "dataset_operationaspect_v1",
-                    DATASET_WRITE_USAGE_RAW_QUERY,
+                    QueryBuilder.get_dataset_write_usage_raw_query(
+                        self.config.lookback_days
+                    ),
                     self.write_stat_raw_batch,
                 ),
                 schema={"urn": polars.Categorical, "platform": polars.Categorical},
@@ -918,12 +816,12 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
     def load_write_usage_server_side_aggregation(
         self, soft_deleted_entities_df: polars.LazyFrame
     ) -> polars.LazyFrame:
-        query: Dict = DATASET_WRITE_USAGE_COMPOSITE_QUERY
-        query["aggs"]["urn_count"]["composite"]["size"] = self.config.extract_batch_size
         wdf = polars.LazyFrame(
             self.load_data_from_es(
                 "dataset_operationaspect_v1",
-                DATASET_WRITE_USAGE_COMPOSITE_QUERY,
+                QueryBuilder.get_dataset_write_usage_composite_query(
+                    self.config.lookback_days
+                ),
                 self.write_stat_batch,
                 aggregation_key="urn_count",
             ),
@@ -954,7 +852,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         upstreams_lf = polars.LazyFrame(
             self.load_data_from_es(
                 "graph_service_v1",
-                GET_UPSTREAMS,
+                QueryBuilder.get_upstreams_query(),
                 self.upstream_lineage_batch,
             ),
             schema={
@@ -1289,7 +1187,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         soft_deleted_df = polars.LazyFrame(
             self.load_data_from_es(
                 index=entity_index,
-                query=GET_SOFT_DELETED_ENTITIES,
+                query=QueryBuilder.get_soft_deleted_entities_query(),
                 process_function=self.soft_deleted_batch,
             ),
             schema={
@@ -1305,7 +1203,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         lf: polars.LazyFrame = polars.LazyFrame(
             self.load_data_from_es(
                 index=usage_index,
-                query=GET_DASHBOARD_USAGE_QUERY,
+                query=QueryBuilder.get_dashboard_usage_query(self.config.lookback_days),
                 process_function=self.process_dashboard_usage,
             ),
             schema={
@@ -1405,7 +1303,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         query_entities = polars.LazyFrame(
             self.load_data_from_es(
                 index=entity_index,
-                query=GET_QUERY_ENTITIES,
+                query=QueryBuilder.get_query_entities_query(),
                 process_function=self.queries_entities_batch,
             ),
             schema={
@@ -1420,7 +1318,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         lf: polars.LazyFrame = polars.LazyFrame(
             self.load_data_from_es(
                 index=usage_index,
-                query=GET_QUERY_USAGE_QUERY,
+                query=QueryBuilder.get_query_usage_query(self.config.lookback_days),
                 process_function=self.process_query_usage,
             ),
             schema={
@@ -1484,7 +1382,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         lf: polars.LazyFrame = polars.LazyFrame(
             self.load_data_from_es(
                 index=index,
-                query=GET_DATASET_USAGE_QUERY,
+                query=QueryBuilder.get_dataset_usage_query(self.config.lookback_days),
                 process_function=self.process_batch,
             ),
             schema={
@@ -1576,7 +1474,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         datasets_df = polars.LazyFrame(
             self.load_data_from_es(
                 index="datasetindex_v2",
-                query=GET_SOFT_DELETED_ENTITIES,
+                query=QueryBuilder.get_soft_deleted_entities_query(),
                 process_function=self.soft_deleted_batch,
             ),
             schema={

{acryl_datahub_cloud-0.3.7.7rc6 → acryl_datahub_cloud-0.3.7.7rc8}/src/acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/structured/__init__.py

@@ -13,6 +13,7 @@ from .....schema_classes import StructuredPropertiesClass
 from .....schema_classes import StructuredPropertyDefinitionClass
 from .....schema_classes import StructuredPropertyFilterStatusClass
 from .....schema_classes import StructuredPropertyKeyClass
+from .....schema_classes import StructuredPropertySettingsClass
 from .....schema_classes import StructuredPropertyValueAssignmentClass
 
 
@@ -22,6 +23,7 @@ StructuredProperties = StructuredPropertiesClass
 StructuredPropertyDefinition = StructuredPropertyDefinitionClass
 StructuredPropertyFilterStatus = StructuredPropertyFilterStatusClass
 StructuredPropertyKey = StructuredPropertyKeyClass
+StructuredPropertySettings = StructuredPropertySettingsClass
 StructuredPropertyValueAssignment = StructuredPropertyValueAssignmentClass
 
 # fmt: on