acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

datahub/metadata/schema_classes.py

@@ -20,16 +20,16 @@ from datahub.utilities._custom_package_loader import get_custom_models_package
 _custom_package_path = get_custom_models_package()
 
 if TYPE_CHECKING or not _custom_package_path:
-    from ._schema_classes import *
+    from ._internal_schema_classes import *
 
     # Required explicitly because __all__ doesn't include _ prefixed names.
-    from ._schema_classes import __SCHEMA_TYPES
+    from ._internal_schema_classes import __SCHEMA_TYPES
 
     if IS_SPHINX_BUILD:
         # Set __module__ to the current module so that Sphinx will document the
         # classes as belonging to this module instead of the custom package.
         for _cls in list(globals().values()):
-            if hasattr(_cls, "__module__") and "datahub.metadata._schema_classes" in _cls.__module__:
+            if hasattr(_cls, "__module__") and "datahub.metadata._internal_schema_classes" in _cls.__module__:
                 _cls.__module__ = __name__
 else:
     _custom_package = importlib.import_module(_custom_package_path)

datahub/metadata/schemas/DataContractKey.avsc

@@ -7,7 +7,8 @@
     "entityAspects": [
       "dataContractProperties",
       "dataContractStatus",
-      "status"
+      "status",
+      "structuredProperties"
     ]
   },
   "name": "DataContractKey",

datahub/metadata/schemas/DataHubOpenAPISchemaKey.avsc

@@ -0,0 +1,22 @@
+{
+  "type": "record",
+  "Aspect": {
+    "name": "dataHubOpenAPISchemaKey",
+    "keyForEntity": "dataHubOpenAPISchema",
+    "entityCategory": "internal",
+    "entityAspects": [
+      "systemMetadata"
+    ],
+    "entityDoc": "Contains aspects which are used in OpenAPI requests/responses which are not otherwise present in the data model."
+  },
+  "name": "DataHubOpenAPISchemaKey",
+  "namespace": "com.linkedin.pegasus2avro.metadata.key",
+  "fields": [
+    {
+      "type": "string",
+      "name": "id",
+      "doc": "A unique id for the DataHub OpenAPI schema."
+    }
+  ],
+  "doc": "Key for a Query"
+}

datahub/metadata/schemas/DataTransformLogic.avsc

@@ -31,12 +31,14 @@
                       "type": {
                         "type": "enum",
                         "symbolDocs": {
-                          "SQL": "A SQL Query"
+                          "SQL": "A SQL Query",
+                          "UNKNOWN": "Unknown query language"
                         },
                         "name": "QueryLanguage",
                         "namespace": "com.linkedin.pegasus2avro.query",
                         "symbols": [
-                          "SQL"
+                          "SQL",
+                          "UNKNOWN"
                         ]
                       },
                       "name": "language",

datahub/metadata/schemas/FormInfo.avsc

@@ -53,6 +53,11 @@
           "namespace": "com.linkedin.pegasus2avro.form",
           "fields": [
             {
+              "Searchable": {
+                "fieldName": "promptId",
+                "fieldType": "KEYWORD",
+                "queryByDefault": false
+              },
               "type": "string",
               "name": "id",
               "doc": "The unique id for this prompt. This must be GLOBALLY unique."

datahub/metadata/schemas/MLModelDeploymentProperties.avsc

@@ -135,6 +135,9 @@
       "doc": "Version of the MLModelDeployment"
     },
     {
+      "Searchable": {
+        "fieldName": "deploymentStatus"
+      },
       "type": [
         "null",
         {

datahub/metadata/schemas/MetadataChangeEvent.avsc

@@ -6755,6 +6755,9 @@
                         "doc": "Version of the MLModelDeployment"
                       },
                       {
+                        "Searchable": {
+                          "fieldName": "deploymentStatus"
+                        },
                         "type": [
                           "null",
                           {
@@ -8116,6 +8119,9 @@
         "null",
         {
           "type": "record",
+          "Aspect": {
+            "name": "systemMetadata"
+          },
           "name": "SystemMetadata",
           "namespace": "com.linkedin.pegasus2avro.mxe",
           "fields": [

datahub/metadata/schemas/MetadataChangeLog.avsc

@@ -187,6 +187,9 @@
         "null",
         {
           "type": "record",
+          "Aspect": {
+            "name": "systemMetadata"
+          },
           "name": "SystemMetadata",
           "namespace": "com.linkedin.pegasus2avro.mxe",
           "fields": [

datahub/metadata/schemas/MetadataChangeProposal.avsc

@@ -187,6 +187,9 @@
         "null",
         {
           "type": "record",
+          "Aspect": {
+            "name": "systemMetadata"
+          },
           "name": "SystemMetadata",
           "namespace": "com.linkedin.pegasus2avro.mxe",
           "fields": [

datahub/metadata/schemas/QueryProperties.avsc

@@ -36,12 +36,14 @@
             "type": {
               "type": "enum",
               "symbolDocs": {
-                "SQL": "A SQL Query"
+                "SQL": "A SQL Query",
+                "UNKNOWN": "Unknown query language"
               },
               "name": "QueryLanguage",
               "namespace": "com.linkedin.pegasus2avro.query",
               "symbols": [
-                "SQL"
+                "SQL",
+                "UNKNOWN"
               ]
             },
             "name": "language",

datahub/metadata/schemas/SystemMetadata.avsc

@@ -0,0 +1,86 @@
+{
+  "type": "record",
+  "Aspect": {
+    "name": "systemMetadata"
+  },
+  "name": "SystemMetadata",
+  "namespace": "com.linkedin.pegasus2avro.mxe",
+  "fields": [
+    {
+      "type": [
+        "long",
+        "null"
+      ],
+      "name": "lastObserved",
+      "default": 0,
+      "doc": "The timestamp the metadata was observed at"
+    },
+    {
+      "type": [
+        "string",
+        "null"
+      ],
+      "name": "runId",
+      "default": "no-run-id-provided",
+      "doc": "The original run id that produced the metadata. Populated in case of batch-ingestion."
+    },
+    {
+      "type": [
+        "string",
+        "null"
+      ],
+      "name": "lastRunId",
+      "default": "no-run-id-provided",
+      "doc": "The last run id that produced the metadata. Populated in case of batch-ingestion."
+    },
+    {
+      "type": [
+        "null",
+        "string"
+      ],
+      "name": "pipelineName",
+      "default": null,
+      "doc": "The ingestion pipeline id that produced the metadata. Populated in case of batch ingestion."
+    },
+    {
+      "type": [
+        "null",
+        "string"
+      ],
+      "name": "registryName",
+      "default": null,
+      "doc": "The model registry name that was used to process this event"
+    },
+    {
+      "type": [
+        "null",
+        "string"
+      ],
+      "name": "registryVersion",
+      "default": null,
+      "doc": "The model registry version that was used to process this event"
+    },
+    {
+      "type": [
+        "null",
+        {
+          "type": "map",
+          "values": "string"
+        }
+      ],
+      "name": "properties",
+      "default": null,
+      "doc": "Additional properties"
+    },
+    {
+      "type": [
+        "null",
+        "string"
+      ],
+      "name": "version",
+      "default": null,
+      "doc": "Aspect version\n   Initial implementation will use the aspect version's number, however stored as\n   a string in the case where a different aspect versioning scheme is later adopted."
+    }
+  ],
+  "doc": "Metadata associated with each metadata change that is processed by the system"
+}

datahub/metadata/schemas/__init__.py

@@ -15,10 +15,10 @@ import pathlib
 def _load_schema(schema_name: str) -> str:
     return (pathlib.Path(__file__).parent / f"{schema_name}.avsc").read_text()
 
-def getMetadataChangeProposalSchema() -> str:
-    return _load_schema("MetadataChangeProposal")
-
 def getMetadataChangeEventSchema() -> str:
     return _load_schema("MetadataChangeEvent")
 
+def getMetadataChangeProposalSchema() -> str:
+    return _load_schema("MetadataChangeProposal")
+
 # fmt: on

datahub/sdk/_all_entities.py

@@ -3,11 +3,15 @@ from typing import Dict, List, Type
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 # TODO: Is there a better way to declare this?
 ENTITY_CLASSES_LIST: List[Type[Entity]] = [
     Container,
     Dataset,
+    MLModel,
+    MLModelGroup,
 ]
 
 ENTITY_CLASSES: Dict[str, Type[Entity]] = {

datahub/sdk/_shared.py

@@ -5,6 +5,7 @@ from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Callable,
+    Dict,
     List,
     Optional,
     Sequence,
@@ -14,6 +15,7 @@ from typing import (
 
 from typing_extensions import TypeAlias, assert_never
 
+import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import (
     make_ts_millis,
@@ -30,14 +32,16 @@ from datahub.metadata.urns import (
     DataJobUrn,
     DataPlatformInstanceUrn,
     DataPlatformUrn,
+    DataProcessInstanceUrn,
     DatasetUrn,
     DomainUrn,
     GlossaryTermUrn,
     OwnershipTypeUrn,
     TagUrn,
     Urn,
+    VersionSetUrn,
 )
-from datahub.sdk._utils import add_list_unique, remove_list_unique
+from datahub.sdk._utils import DEFAULT_ACTOR_URN, add_list_unique, remove_list_unique
 from datahub.sdk.entity import Entity
 from datahub.utilities.urns.error import InvalidUrnError
 
@@ -50,7 +54,35 @@ DatajobUrnOrStr: TypeAlias = Union[str, DataJobUrn]
 
 ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
 
-_DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
+TrainingMetricsInputType: TypeAlias = Union[
+    List[models.MLMetricClass], Dict[str, Optional[str]]
+]
+HyperParamsInputType: TypeAlias = Union[
+    List[models.MLHyperParamClass], Dict[str, Optional[str]]
+]
+MLTrainingJobInputType: TypeAlias = Union[Sequence[Union[str, DataProcessInstanceUrn]]]
+
+
+def convert_training_metrics(
+    metrics: TrainingMetricsInputType,
+) -> List[models.MLMetricClass]:
+    if isinstance(metrics, dict):
+        return [
+            models.MLMetricClass(name=name, value=str(value))
+            for name, value in metrics.items()
+        ]
+    return metrics
+
+
+def convert_hyper_params(
+    params: HyperParamsInputType,
+) -> List[models.MLHyperParamClass]:
+    if isinstance(params, dict):
+        return [
+            models.MLHyperParamClass(name=name, value=str(value))
+            for name, value in params.items()
+        ]
+    return params
 
 
 def make_time_stamp(ts: Optional[datetime]) -> Optional[models.TimeStampClass]:
@@ -441,7 +473,7 @@ class HasTerms(Entity):
     def _terms_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=0,
-            actor=_DEFAULT_ACTOR_URN,
+            actor=DEFAULT_ACTOR_URN,
         )
 
     def set_terms(self, terms: TermsInputType) -> None:
@@ -529,7 +561,7 @@ class HasInstitutionalMemory(Entity):
     def _institutional_memory_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=0,
-            actor=_DEFAULT_ACTOR_URN,
+            actor=DEFAULT_ACTOR_URN,
         )
 
     @classmethod
@@ -578,3 +610,109 @@ class HasInstitutionalMemory(Entity):
             self._link_key,
             self._parse_link_association_class(link),
         )
+
+
+class HasVersion(Entity):
+    """Mixin for entities that have version properties."""
+
+    def _get_version_props(self) -> Optional[models.VersionPropertiesClass]:
+        return self._get_aspect(models.VersionPropertiesClass)
+
+    def _ensure_version_props(self) -> models.VersionPropertiesClass:
+        version_props = self._get_version_props()
+        if version_props is None:
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+
+            version_props = models.VersionPropertiesClass(
+                versionSet=str(version_set_urn),
+                version=models.VersionTagClass(versionTag="0.1.0"),
+                sortId="0000000.1.0",
+            )
+            self._set_aspect(version_props)
+        return version_props
+
+    @property
+    def version(self) -> Optional[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.version:
+            return version_props.version.versionTag
+        return None
+
+    def set_version(self, version: str) -> None:
+        """Set the version of the entity."""
+        guid_dict = {"urn": str(self.urn)}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+        )
+
+        version_props = self._get_version_props()
+        if version_props is None:
+            # If no version properties exist, create a new one
+            version_props = models.VersionPropertiesClass(
+                version=models.VersionTagClass(versionTag=version),
+                versionSet=str(version_set_urn),
+                sortId=version.zfill(10),  # Pad with zeros for sorting
+            )
+        else:
+            # Update existing version properties
+            version_props.version = models.VersionTagClass(versionTag=version)
+            version_props.versionSet = str(version_set_urn)
+            version_props.sortId = version.zfill(10)
+
+        self._set_aspect(version_props)
+
+    @property
+    def version_aliases(self) -> List[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            return [
+                alias.versionTag
+                for alias in version_props.aliases
+                if alias.versionTag is not None
+            ]
+        return []  # Return empty list instead of None
+
+    def set_version_aliases(self, aliases: List[str]) -> None:
+        version_props = self._get_aspect(models.VersionPropertiesClass)
+        if version_props:
+            version_props.aliases = [
+                models.VersionTagClass(versionTag=alias) for alias in aliases
+            ]
+        else:
+            # If no version properties exist, we need to create one with a default version
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+            self._set_aspect(
+                models.VersionPropertiesClass(
+                    version=models.VersionTagClass(
+                        versionTag="0.1.0"
+                    ),  # Default version
+                    versionSet=str(version_set_urn),
+                    sortId="0000000.1.0",
+                    aliases=[
+                        models.VersionTagClass(versionTag=alias) for alias in aliases
+                    ],
+                )
+            )
+
+    def add_version_alias(self, alias: str) -> None:
+        if not alias:
+            raise ValueError("Alias cannot be empty")
+        version_props = self._ensure_version_props()
+        if version_props.aliases is None:
+            version_props.aliases = []
+        version_props.aliases.append(models.VersionTagClass(versionTag=alias))
+        self._set_aspect(version_props)
+
+    def remove_version_alias(self, alias: str) -> None:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            version_props.aliases = [
+                a for a in version_props.aliases if a.versionTag != alias
+            ]
+            self._set_aspect(version_props)

datahub/sdk/_utils.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, List, Protocol, TypeVar
 
 from datahub.errors import ItemNotFoundError
+from datahub.metadata.urns import CorpUserUrn
+
+# TODO: Change __ingestion to _ingestion.
+DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
 
 
 class _SupportsEq(Protocol):

datahub/sdk/dataset.py

@@ -87,7 +87,7 @@ def _parse_upstream_input(
         assert_never(upstream_input)
 
 
-def _parse_cll_mapping(
+def parse_cll_mapping(
     *,
     upstream: DatasetUrnOrStr,
     downstream: DatasetUrnOrStr,
@@ -142,7 +142,7 @@ def _parse_upstream_lineage_input(
                 )
             )
             cll.extend(
-                _parse_cll_mapping(
+                parse_cll_mapping(
                     upstream=dataset_urn,
                     downstream=downstream_urn,
                     cll_mapping=column_lineage,

datahub/sdk/entity_client.py

@@ -11,6 +11,8 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import (
     ContainerUrn,
     DatasetUrn,
+    MlModelGroupUrn,
+    MlModelUrn,
     Urn,
 )
 from datahub.sdk._all_entities import ENTITY_CLASSES
@@ -18,6 +20,8 @@ from datahub.sdk._shared import UrnOrStr
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
@@ -49,6 +53,10 @@ class EntityClient:
     @overload
     def get(self, urn: DatasetUrn) -> Dataset: ...
     @overload
+    def get(self, urn: MlModelUrn) -> MLModel: ...
+    @overload
+    def get(self, urn: MlModelGroupUrn) -> MLModelGroup: ...
+    @overload
     def get(self, urn: Union[Urn, str]) -> Entity: ...
     def get(self, urn: UrnOrStr) -> Entity:
         """Retrieve an entity by its urn.