acryl-datahub-cloud 0.3.8rc11__py3-none-any.whl → 0.3.8rc13__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

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

acryl_datahub_cloud/datahub_metadata_sharing/metadata_sharing_source.py

@@ -112,6 +112,8 @@ class DataHubMetadataSharingSource(Source):
 
                 if next_scroll_id is None:
                     break
+                else:
+                    scroll_id = next_scroll_id
 
                 time.sleep(self.config.batch_delay_ms / 1000.0)
 
@@ -193,8 +195,11 @@ class DataHubMetadataSharingSource(Source):
                 self.report.entities_failed += 1
 
     def _determine_lineage_direction(
-        self, share_config: Dict[str, Any]
+        self, share_config: Optional[Dict[str, Any]]
     ) -> Optional[str]:
+        if share_config is None:
+            return None
+
         """Determine lineage direction based on share config"""
         include_upstreams = share_config.get("enableUpstreamLineage", False)
         include_downstreams = share_config.get(

acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -396,59 +396,30 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                 continue
 
             yield {
-                "timestampMillis": doc["_source"]["timestampMillis"],
-                "lastObserved": doc["_source"]["systemMetadata"]["lastObserved"],
-                "urn": doc["_source"]["urn"],
-                "eventGranularity": (
-                    doc["_source"]["eventGranularity"]
-                    if "eventGranularity" in doc["_source"]
-                    else None
-                ),
-                "viewsCount": (
-                    doc["_source"]["viewsCount"]
-                    if "viewsCount" in doc["_source"]
-                    else 0
-                ),
-                "uniqueUserCount": (
-                    doc["_source"]["uniqueUserCount"]
-                    if "uniqueUserCount" in doc["_source"]
-                    else None
-                ),
-                "userCounts": (
-                    doc["_source"]["event"]["userCounts"]
-                    if "event" in doc["_source"]
-                    and "userCounts" in doc["_source"]["event"]
-                    else []
-                ),
+                "timestampMillis": doc["_source"].get("timestampMillis"),
+                "lastObserved": doc["_source"]
+                .get("systemMetadata", {})
+                .get("lastObserved"),
+                "urn": doc["_source"].get("urn"),
+                "eventGranularity": doc["_source"].get("eventGranularity"),
+                "viewsCount": doc["_source"].get("viewsCount", 0),
+                "uniqueUserCount": doc["_source"].get("uniqueUserCount"),
+                "userCounts": doc["_source"].get("event", {}).get("userCounts", []),
                 "platform": platform,
             }
 
     def process_query_usage(self, results: Iterable) -> Iterable[Dict]:
         for doc in results:
             yield {
-                "timestampMillis": doc["_source"]["timestampMillis"],
-                "lastObserved": doc["_source"]["systemMetadata"]["lastObserved"],
-                "urn": doc["_source"]["urn"],
-                "eventGranularity": (
-                    doc["_source"]["eventGranularity"]
-                    if "eventGranularity" in doc["_source"]
-                    else None
-                ),
-                "queryCount": (
-                    doc["_source"]["queryCount"]
-                    if "queryCount" in doc["_source"]
-                    else 0
-                ),
-                "uniqueUserCount": (
-                    doc["_source"]["uniqueUserCount"]
-                    if "uniqueUserCount" in doc["_source"]
-                    else None
-                ),
-                "userCounts": (
-                    doc["_source"]["event"]["userCounts"]
-                    if "userCounts" in doc["_source"]["event"]
-                    else []
-                ),
+                "timestampMillis": doc["_source"].get("timestampMillis"),
+                "lastObserved": doc["_source"]
+                .get("systemMetadata", {})
+                .get("lastObserved"),
+                "urn": doc["_source"].get("urn"),
+                "eventGranularity": doc["_source"].get("eventGranularity"),
+                "queryCount": doc["_source"].get("queryCount", 0),
+                "uniqueUserCount": doc["_source"].get("uniqueUserCount"),
+                "userCounts": doc["_source"].get("event", {}).get("userCounts", []),
             }
 
     def upstream_lineage_batch(self, results: Iterable) -> Iterable[Dict]:

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/ml/metadata/__init__.py

@@ -30,6 +30,7 @@ from ......schema_classes import MLModelFactorsClass
 from ......schema_classes import MLModelGroupPropertiesClass
 from ......schema_classes import MLModelPropertiesClass
 from ......schema_classes import MLPrimaryKeyPropertiesClass
+from ......schema_classes import MLTrainingRunPropertiesClass
 from ......schema_classes import MetricsClass
 from ......schema_classes import QuantitativeAnalysesClass
 from ......schema_classes import SourceCodeClass
@@ -61,6 +62,7 @@ MLModelFactors = MLModelFactorsClass
 MLModelGroupProperties = MLModelGroupPropertiesClass
 MLModelProperties = MLModelPropertiesClass
 MLPrimaryKeyProperties = MLPrimaryKeyPropertiesClass
+MLTrainingRunProperties = MLTrainingRunPropertiesClass
 Metrics = MetricsClass
 QuantitativeAnalyses = QuantitativeAnalysesClass
 SourceCode = SourceCodeClass