acryl-datahub-cloud 0.3.7.6__py3-none-any.whl → 0.3.7.7__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

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

acryl_datahub_cloud/datahub_usage_reporting/usage_feature_reporter.py

@@ -449,7 +449,6 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
     def queries_entities_batch(self, results: Iterable) -> Iterable[Dict]:
         with PerfTimer() as timer:
-
             for doc in results:
                 if "platform" not in doc["_source"] or not doc["_source"]["platform"]:
                     logger.warning(
@@ -473,9 +472,11 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
                         )
                     ),
                     "platform": doc["_source"]["platform"],
-                    "removed": doc["_source"]["removed"]
-                    if "removed" in doc["_source"]
-                    else False,
+                    "removed": (
+                        doc["_source"]["removed"]
+                        if "removed" in doc["_source"]
+                        else False
+                    ),
                 }
 
             time_taken = timer.elapsed_seconds()
@@ -587,7 +588,6 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
 
     def process_batch(self, results: Iterable) -> Iterable[Dict]:
         with PerfTimer() as timer:
-
             for doc in results:
                 match = re.match(platform_regexp, doc["_source"]["urn"])
                 if match:
@@ -754,7 +754,6 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
         prefix: Optional[str] = None,
         use_exp_cdf: Optional[bool] = None,
     ) -> polars.LazyFrame:
-
         logger.debug(f"Generating rank and percentile for {count_field} field")
         lf = lf.with_columns(
             polars.col(count_field)
@@ -880,7 +879,6 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
     def load_write_usage(
         self, soft_deleted_entities_df: polars.LazyFrame
     ) -> polars.LazyFrame:
-
         if self.config.streaming_mode:
             wdf = self.load_es_data_to_lf(
                 index="dataset_operationaspect_v1",
@@ -1557,10 +1555,8 @@ class DataHubUsageFeatureReportingSource(StatefulIngestionSourceBase):
             usage_with_top_users_with_ranks.join(
                 write_lf, on="urn", how="full", suffix="_write"
             )
-            .with_columns("write_count")
-            .fill_null(polars.lit(0))
-            .with_columns("totalSqlQueries")
-            .fill_null(polars.lit(0))
+            .with_columns(polars.col("write_count").fill_null(polars.lit(0)))
+            .with_columns(polars.col("totalSqlQueries").fill_null(polars.lit(0)))
         )
 
         # If we get a dataset from the operation aspect index only then we have to use its urn and platform

acryl_datahub_cloud/metadata/_urns/urn_defs.py

@@ -21,6 +21,60 @@ from datahub.utilities.urns.error import InvalidUrnError
 
 deprecated = functools.partial(_sphinx_deprecated, version="0.12.0.2")
 
+if TYPE_CHECKING:
+    from datahub.metadata.schema_classes import VersionSetKeyClass
+
+class VersionSetUrn(_SpecificUrn):
+    ENTITY_TYPE: ClassVar[str] = "versionSet"
+    URN_PARTS: ClassVar[int] = 2
+
+    def __init__(self, id: str, entity_type: str, *, _allow_coercion: bool = True) -> None:
+        if _allow_coercion:
+            # Field coercion logic (if any is required).
+            id = UrnEncoder.encode_string(id)
+            entity_type = UrnEncoder.encode_string(entity_type)
+
+        # Validation logic.
+        if not id:
+            raise InvalidUrnError("VersionSetUrn id cannot be empty")
+        if UrnEncoder.contains_reserved_char(id):
+            raise InvalidUrnError(f'VersionSetUrn id contains reserved characters')
+        if not entity_type:
+            raise InvalidUrnError("VersionSetUrn entity_type cannot be empty")
+        if UrnEncoder.contains_reserved_char(entity_type):
+            raise InvalidUrnError(f'VersionSetUrn entity_type contains reserved characters')
+
+        super().__init__(self.ENTITY_TYPE, [id, entity_type])
+
+    @classmethod
+    def _parse_ids(cls, entity_ids: List[str]) -> "VersionSetUrn":
+        if len(entity_ids) != cls.URN_PARTS:
+            raise InvalidUrnError(f"VersionSetUrn should have {cls.URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
+        return cls(id=entity_ids[0], entity_type=entity_ids[1], _allow_coercion=False)
+
+    @classmethod
+    def underlying_key_aspect_type(cls) -> Type["VersionSetKeyClass"]:
+        from datahub.metadata.schema_classes import VersionSetKeyClass
+
+        return VersionSetKeyClass
+
+    def to_key_aspect(self) -> "VersionSetKeyClass":
+        from datahub.metadata.schema_classes import VersionSetKeyClass
+
+        return VersionSetKeyClass(id=self.id, entityType=self.entity_type)
+
+    @classmethod
+    def from_key_aspect(cls, key_aspect: "VersionSetKeyClass") -> "VersionSetUrn":
+        return cls(id=key_aspect.id, entity_type=key_aspect.entityType)
+
+    @property
+    def id(self) -> str:
+        return self.entity_ids[0]
+
+    @property
+    def entity_type(self) -> str:
+        return self.entity_ids[1]
+
 if TYPE_CHECKING:
     from datahub.metadata.schema_classes import DataHubConnectionKeyClass
 

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

@@ -87,6 +87,7 @@ from .....schema_classes import SubTypesClass
 from .....schema_classes import SyncMechanismClass
 from .....schema_classes import TagAssociationClass
 from .....schema_classes import TimeStampClass
+from .....schema_classes import VersionPropertiesClass
 from .....schema_classes import VersionTagClass
 from .....schema_classes import WindowDurationClass
 
@@ -171,6 +172,7 @@ SubTypes = SubTypesClass
 SyncMechanism = SyncMechanismClass
 TagAssociation = TagAssociationClass
 TimeStamp = TimeStampClass
+VersionProperties = VersionPropertiesClass
 VersionTag = VersionTagClass
 WindowDuration = WindowDurationClass
 

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

@@ -65,6 +65,7 @@ from ......schema_classes import SubscriptionKeyClass
 from ......schema_classes import TagKeyClass
 from ......schema_classes import TelemetryKeyClass
 from ......schema_classes import TestKeyClass
+from ......schema_classes import VersionSetKeyClass
 
 
 ActionRequestKey = ActionRequestKeyClass
@@ -125,5 +126,6 @@ SubscriptionKey = SubscriptionKeyClass
 TagKey = TagKeyClass
 TelemetryKey = TelemetryKeyClass
 TestKey = TestKeyClass
+VersionSetKey = VersionSetKeyClass
 
 # fmt: on

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

@@ -0,0 +1,17 @@
+# mypy: ignore-errors
+# flake8: noqa
+
+# This file is autogenerated by /metadata-ingestion/scripts/avro_codegen.py
+# Do not modify manually!
+
+# pylint: skip-file
+# fmt: off
+# isort: skip_file
+from .....schema_classes import VersionSetPropertiesClass
+from .....schema_classes import VersioningSchemeClass
+
+
+VersionSetProperties = VersionSetPropertiesClass
+VersioningScheme = VersioningSchemeClass
+
+# fmt: on

acryl_datahub_cloud/metadata/schema.avsc

@@ -2418,6 +2418,32 @@
       }
     ]
   },
+  {
+    "type": "record",
+    "Aspect": {
+      "name": "versionSetKey",
+      "keyForEntity": "versionSet",
+      "entityCategory": "core",
+      "entityAspects": [
+        "versionSetProperties"
+      ]
+    },
+    "name": "VersionSetKey",
+    "namespace": "com.linkedin.pegasus2avro.metadata.key",
+    "fields": [
+      {
+        "type": "string",
+        "name": "id",
+        "doc": "ID of the Version Set, generated from platform + asset id / name"
+      },
+      {
+        "type": "string",
+        "name": "entityType",
+        "doc": "Type of entities included in version set, limits to a single entity type between linked versioned entities"
+      }
+    ],
+    "doc": "Key for a Version Set entity"
+  },
   {
     "type": "record",
     "Aspect": {
@@ -4090,7 +4116,8 @@
         "share",
         "origin",
         "documentation",
-        "entityInferenceMetadata"
+        "entityInferenceMetadata",
+        "versionProperties"
       ],
       "entityDoc": "Datasets represent logical or physical data assets stored or represented in various data platforms. Tables, Views, Streams are all instances of datasets."
     },
@@ -4409,7 +4436,8 @@
         "origin",
         "lineageFeatures",
         "documentation",
-        "incidentsSummary"
+        "incidentsSummary",
+        "versionProperties"
       ]
     },
     "name": "MLModelKey",
@@ -9292,7 +9320,7 @@
                                                                     "fields": [
                                                                       {
                                                                         "Searchable": {
-                                                                          "boostScore": 5.0,
+                                                                          "boostScore": 1.0,
                                                                           "fieldName": "fieldPaths",
                                                                           "fieldType": "TEXT",
                                                                           "queryByDefault": "true"
@@ -11820,6 +11848,132 @@
     ],
     "doc": "Proposals aspect for proposed tags and glossary terms for an entity.\nThese are secondary indices and ActionRequests remain the source of truth."
   },
+  {
+    "type": "record",
+    "Aspect": {
+      "name": "versionProperties"
+    },
+    "name": "VersionProperties",
+    "namespace": "com.linkedin.pegasus2avro.common",
+    "fields": [
+      {
+        "Relationship": {
+          "entityTypes": [
+            "versionSet"
+          ],
+          "name": "VersionOf"
+        },
+        "Searchable": {
+          "queryByDefault": false
+        },
+        "java": {
+          "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+        },
+        "Urn": "Urn",
+        "entityTypes": [
+          "versionSet"
+        ],
+        "type": "string",
+        "name": "versionSet",
+        "doc": "The linked Version Set entity that ties multiple versioned assets together"
+      },
+      {
+        "Searchable": {
+          "/versionTag": {
+            "fieldName": "version",
+            "queryByDefault": false
+          }
+        },
+        "type": {
+          "type": "record",
+          "name": "VersionTag",
+          "namespace": "com.linkedin.pegasus2avro.common",
+          "fields": [
+            {
+              "type": [
+                "null",
+                "string"
+              ],
+              "name": "versionTag",
+              "default": null
+            },
+            {
+              "type": [
+                "null",
+                "com.linkedin.pegasus2avro.common.MetadataAttribution"
+              ],
+              "name": "metadataAttribution",
+              "default": null
+            }
+          ],
+          "doc": "A resource-defined string representing the resource state for the purpose of concurrency control"
+        },
+        "name": "version",
+        "doc": "Label for this versioned asset, is unique within a version set"
+      },
+      {
+        "Searchable": {
+          "/*/versionTag": {
+            "fieldName": "aliases",
+            "queryByDefault": false
+          }
+        },
+        "type": {
+          "type": "array",
+          "items": "com.linkedin.pegasus2avro.common.VersionTag"
+        },
+        "name": "aliases",
+        "default": [],
+        "doc": "Associated aliases for this versioned asset"
+      },
+      {
+        "type": [
+          "null",
+          "string"
+        ],
+        "name": "comment",
+        "default": null,
+        "doc": "Comment documenting what this version was created for, changes, or represents"
+      },
+      {
+        "Searchable": {
+          "fieldName": "versionSortId",
+          "queryByDefault": false
+        },
+        "type": "string",
+        "name": "sortId",
+        "doc": "Sort identifier that determines where a version lives in the order of the Version Set.\nWhat this looks like depends on the Version Scheme. For sort ids generated by DataHub we use an 8 character string representation."
+      },
+      {
+        "type": [
+          "null",
+          "com.linkedin.pegasus2avro.common.AuditStamp"
+        ],
+        "name": "sourceCreatedTimestamp",
+        "default": null,
+        "doc": "Timestamp reflecting when this asset version was created in the source system."
+      },
+      {
+        "type": [
+          "null",
+          "com.linkedin.pegasus2avro.common.AuditStamp"
+        ],
+        "name": "metadataCreatedTimestamp",
+        "default": null,
+        "doc": "Timestamp reflecting when the metadata for this version was created in DataHub"
+      },
+      {
+        "type": [
+          "null",
+          "boolean"
+        ],
+        "name": "isLatest",
+        "default": null,
+        "doc": "Marks whether this version is currently the latest. Set by a side effect and should not be modified by API."
+      }
+    ],
+    "doc": "Properties about a versioned asset i.e. dataset, ML Model, etc."
+  },
   {
     "type": "record",
     "Aspect": {
@@ -22048,22 +22202,7 @@
                         {
                           "type": [
                             "null",
-                            {
-                              "type": "record",
-                              "name": "VersionTag",
-                              "namespace": "com.linkedin.pegasus2avro.common",
-                              "fields": [
-                                {
-                                  "type": [
-                                    "null",
-                                    "string"
-                                  ],
-                                  "name": "versionTag",
-                                  "default": null
-                                }
-                              ],
-                              "doc": "A resource-defined string representing the resource state for the purpose of concurrency control"
-                            }
+                            "com.linkedin.pegasus2avro.common.VersionTag"
                           ],
                           "name": "version",
                           "default": null,
@@ -26385,6 +26524,55 @@
     ],
     "doc": "A request to send a notification."
   },
+  {
+    "type": "record",
+    "Aspect": {
+      "name": "versionSetProperties"
+    },
+    "name": "VersionSetProperties",
+    "namespace": "com.linkedin.pegasus2avro.versionset",
+    "fields": [
+      {
+        "Searchable": {
+          "/*": {
+            "fieldType": "TEXT",
+            "queryByDefault": true
+          }
+        },
+        "type": {
+          "type": "map",
+          "values": "string"
+        },
+        "name": "customProperties",
+        "default": {},
+        "doc": "Custom property bag."
+      },
+      {
+        "Searchable": {
+          "queryByDefault": "false"
+        },
+        "java": {
+          "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+        },
+        "Urn": "Urn",
+        "type": "string",
+        "name": "latest",
+        "doc": "The latest versioned entity linked to in this version set"
+      },
+      {
+        "type": {
+          "type": "enum",
+          "name": "VersioningScheme",
+          "namespace": "com.linkedin.pegasus2avro.versionset",
+          "symbols": [
+            "ALPHANUMERIC_GENERATED_BY_DATAHUB"
+          ]
+        },
+        "name": "versioningScheme",
+        "doc": "What versioning scheme is being utilized for the versioned entities sort criterion. Static once set"
+      }
+    ]
+  },
   {
     "type": "record",
     "Aspect": {