acryl-datahub-cloud 0.3.8.3__py3-none-any.whl → 0.3.9__py3-none-any.whl

acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -370,7 +370,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
             for doc in results:
                 if "platform" not in doc["_source"] or not doc["_source"]["platform"]:
                     logger.warning(
-                        f"Platform not found in query { doc['_source']['urn']}. Skipping..."
+                        f"Platform not found in query {doc['_source']['urn']}. Skipping..."
                     )
                     continue
 
@@ -688,10 +688,10 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
     @staticmethod
     def polars_to_arrow_schema(
-        polars_schema: Dict[str, Union[DataTypeClass, polars.DataType]]
+        polars_schema: Dict[str, Union[DataTypeClass, polars.DataType]],
     ) -> pa.Schema:
         def convert_dtype(
-            polars_dtype: Union[DataTypeClass, polars.DataType]
+            polars_dtype: Union[DataTypeClass, polars.DataType],
         ) -> pa.DataType:
             type_mapping: Dict[Union[DataTypeClass, polars.DataType], pa.DataType] = {
                 polars.Boolean(): pa.bool_(),
@@ -949,14 +949,14 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         if self.config.dataset_usage_enabled:
             with self.report.dataset_usage_processing_time as timer:
-                self.report.report_ingestion_stage_start("generate dataset usage")
+                self.report.new_stage("generate dataset usage")
                 yield from self.generate_dataset_usage_mcps()
                 time_taken = timer.elapsed_seconds()
                 logger.info(f"Dataset Usage generation took {time_taken:.3f} seconds")
 
         if self.config.dashboard_usage_enabled:
             with self.report.dashboard_usage_processing_time as timer:
-                self.report.report_ingestion_stage_start("generate dashboard usage")
+                self.report.new_stage("generate dashboard usage")
                 yield from self.generate_dashboard_usage_mcps()
 
                 time_taken = timer.elapsed_seconds()
@@ -964,7 +964,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
         if self.config.chart_usage_enabled:
             with self.report.chart_usage_processing_time as timer:
-                self.report.report_ingestion_stage_start("generate chart usage")
+                self.report.new_stage("generate chart usage")
 
                 yield from self.generate_chart_usage_mcps()
 
@@ -973,7 +973,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
         if self.config.query_usage_enabled:
             with self.report.query_usage_processing_time as timer:
-                self.report.report_ingestion_stage_start("generate query usage")
+                self.report.new_stage("generate query usage")
 
                 yield from self.generate_query_usage_mcps()
 
@@ -1260,7 +1260,8 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         # We drop last_modified_at to not use it in merge because we are getting last_modified_at from the usage index
         df_with_search_scores = (
             entities_lf.filter(
-                polars.col("combinedSearchRankingMultiplier").is_not_null()
+                polars.col("combinedSearchRankingMultiplier")
+                .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
@@ -1277,7 +1278,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         usages_lf = df_with_search_scores.join(
             usages_lf, on="urn", how="full", suffix="_right"
         )
-        # Merge all common fields automatically
+        ## Merge all common fields automatically
         for common_field in common_fields:
             right_col = f"{common_field}_right"
             usages_lf = usages_lf.with_columns(
@@ -1451,9 +1452,9 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         if not self.config.streaming_mode:
             return polars.LazyFrame(data, schema)
         else:
-            assert (
-                self.temp_dir is not None
-            ), "In Streaming mode temp dir should be set. Normally this should not happen..."
+            assert self.temp_dir is not None, (
+                "In Streaming mode temp dir should be set. Normally this should not happen..."
+            )
 
             with tempfile.NamedTemporaryFile(
                 delete=False,
@@ -1623,7 +1624,7 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                     time_taken = timer.elapsed_seconds()
                     processed_count += len(results["hits"]["hits"])
                     logger.info(
-                        f"Processed {len(results['hits']['hits'''])} data from {index} index in {time_taken:.3f} seconds. Total: {processed_count} processed."
+                        f"Processed {len(results['hits']['hits'])} data from {index} index in {time_taken:.3f} seconds. Total: {processed_count} processed."
                     )
                     if len(results["hits"]["hits"]) < batch_size:
                         break