acryl-datahub 1.2.0.7rc1__py3-none-any.whl → 1.2.0.7rc3__py3-none-any.whl

datahub/ingestion/source/unity/source.py

@@ -12,6 +12,7 @@ from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
     make_domain_urn,
     make_group_urn,
+    make_ml_model_group_urn,
     make_schema_field_urn,
     make_ts_millis,
     make_user_urn,
@@ -26,6 +27,7 @@ from datahub.emitter.mcp_builder import (
     UnitySchemaKey,
     UnitySchemaKeyWithMetastore,
     add_dataset_to_container,
+    add_entity_to_container,
     gen_containers,
 )
 from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
@@ -87,6 +89,8 @@ from datahub.ingestion.source.unity.proxy_types import (
     CustomCatalogType,
     HiveTableType,
     Metastore,
+    Model,
+    ModelVersion,
     Notebook,
     NotebookId,
     Schema,
@@ -121,6 +125,7 @@ from datahub.metadata.schema_classes import (
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
     DomainsClass,
+    MLModelPropertiesClass,
     MySqlDDLClass,
     NullTypeClass,
     OwnerClass,
@@ -134,7 +139,8 @@ from datahub.metadata.schema_classes import (
     UpstreamClass,
     UpstreamLineageClass,
 )
-from datahub.metadata.urns import TagUrn
+from datahub.metadata.urns import MlModelGroupUrn, MlModelUrn, TagUrn
+from datahub.sdk import MLModel, MLModelGroup
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import (
     SqlParsingResult,
@@ -182,6 +188,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     - metastores
     - schemas
     - tables and column lineage
+    - model and model versions
     """
 
     config: UnityCatalogSourceConfig
@@ -512,6 +519,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 yield from self.gen_schema_containers(schema)
                 try:
                     yield from self.process_tables(schema)
+                    yield from self.process_ml_models(schema)
                 except Exception as e:
                     logger.exception(f"Error parsing schema {schema}")
                     self.report.report_warning(
@@ -665,6 +673,69 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             )
         ]
 
+    def process_ml_models(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
+        for ml_model in self.unity_catalog_api_proxy.ml_models(
+            schema=schema, max_results=self.config.ml_model_max_results
+        ):
+            yield from self.process_ml_model(ml_model, schema)
+            ml_model_urn = self.gen_ml_model_urn(ml_model.id)
+            for ml_model_version in self.unity_catalog_api_proxy.ml_model_versions(
+                ml_model, include_aliases=self.config.include_ml_model_aliases
+            ):
+                yield from self.process_ml_model_version(
+                    ml_model_urn, ml_model_version, schema
+                )
+
+    def process_ml_model(
+        self, ml_model: Model, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        ml_model_group = MLModelGroup(
+            id=ml_model.id,
+            name=ml_model.name,
+            platform=self.platform,
+            platform_instance=schema.name,
+            env=self.config.env,
+            description=ml_model.description,
+            created=ml_model.created_at,
+            last_modified=ml_model.updated_at,
+        )
+        yield from ml_model_group.as_workunits()
+        yield from self.add_model_to_schema_container(str(ml_model_group.urn), schema)
+        self.report.ml_models.processed(ml_model.id)
+
+    def process_ml_model_version(
+        self, ml_model_urn: str, ml_model_version: ModelVersion, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        extra_aspects = []
+        if ml_model_version.created_at is not None:
+            created_time = int(ml_model_version.created_at.timestamp() * 1000)
+            created_actor = (
+                f"urn:li:platformResource:{ml_model_version.created_by}"
+                if ml_model_version.created_by
+                else None
+            )
+            extra_aspects.append(
+                MLModelPropertiesClass(
+                    created=TimeStampClass(time=created_time, actor=created_actor),
+                )
+            )
+
+        ml_model = MLModel(
+            id=ml_model_version.id,
+            name=ml_model_version.name,
+            version=str(ml_model_version.version),
+            aliases=ml_model_version.aliases,
+            description=ml_model_version.description,
+            model_group=ml_model_urn,
+            platform=self.platform,
+            last_modified=ml_model_version.updated_at,
+            extra_aspects=extra_aspects,
+        )
+
+        yield from ml_model.as_workunits()
+        yield from self.add_model_version_to_schema_container(str(ml_model.urn), schema)
+        self.report.ml_model_versions.processed(ml_model_version.id)
+
     def ingest_lineage(self, table: Table) -> Optional[UpstreamLineageClass]:
         # Calculate datetime filters for lineage
         lineage_start_time = None
@@ -802,6 +873,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             env=self.config.env,
         )
 
+    def gen_ml_model_urn(self, name: str) -> str:
+        return make_ml_model_group_urn(
+            platform=self.platform,
+            group_name=name,
+            env=self.config.env,
+        )
+
     def gen_notebook_urn(self, notebook: Union[Notebook, NotebookId]) -> str:
         notebook_id = notebook.id if isinstance(notebook, Notebook) else notebook
         return NotebookKey(
@@ -973,6 +1051,26 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             dataset_urn=dataset_urn,
         )
 
+    def add_model_to_schema_container(
+        self, model_urn: str, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        schema_container_key = self.gen_schema_key(schema)
+        yield from add_entity_to_container(
+            container_key=schema_container_key,
+            entity_type=MlModelGroupUrn.ENTITY_TYPE,
+            entity_urn=model_urn,
+        )
+
+    def add_model_version_to_schema_container(
+        self, model_version_urn: str, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        schema_container_key = self.gen_schema_key(schema)
+        yield from add_entity_to_container(
+            container_key=schema_container_key,
+            entity_type=MlModelUrn.ENTITY_TYPE,
+            entity_urn=model_version_urn,
+        )
+
     def _get_catalog_tags(
         self, catalog: str, schema: str, table: str
     ) -> List[UnityCatalogTag]:

datahub/metadata/_internal_schema_classes.py

@@ -15283,23 +15283,23 @@ class LogicalParentClass(_Aspect):
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.logical.LogicalParent")
 
     def __init__(self,
-        parent: "EdgeClass",
+        parent: Union[None, "EdgeClass"]=None,
     ):
         super().__init__()
         
         self.parent = parent
     
     def _restore_defaults(self) -> None:
-        self.parent = EdgeClass._construct_with_defaults()
+        self.parent = self.RECORD_SCHEMA.fields_dict["parent"].default
     
     
     @property
-    def parent(self) -> "EdgeClass":
+    def parent(self) -> Union[None, "EdgeClass"]:
         # No docs available.
         return self._inner_dict.get('parent')  # type: ignore
     
     @parent.setter
-    def parent(self, value: "EdgeClass") -> None:
+    def parent(self, value: Union[None, "EdgeClass"]) -> None:
         self._inner_dict['parent'] = value
     
     
@@ -17004,7 +17004,7 @@ class SchemaFieldKeyClass(_Aspect):
 
 
     ASPECT_NAME = 'schemaFieldKey'
-    ASPECT_INFO = {'keyForEntity': 'schemaField', 'entityCategory': 'core', 'entityAspects': ['schemafieldInfo', 'structuredProperties', 'forms', 'businessAttributes', 'status', 'schemaFieldAliases', 'documentation', 'testResults', 'deprecation', 'subTypes', 'logicalParent']}
+    ASPECT_INFO = {'keyForEntity': 'schemaField', 'entityCategory': 'core', 'entityAspects': ['schemafieldInfo', 'structuredProperties', 'forms', 'businessAttributes', 'status', 'schemaFieldAliases', 'documentation', 'testResults', 'deprecation', 'subTypes', 'logicalParent', 'globalTags', 'glossaryTerms']}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.SchemaFieldKey")
 
     def __init__(self,

datahub/metadata/schema.avsc

@@ -4410,7 +4410,9 @@
         "testResults",
         "deprecation",
         "subTypes",
-        "logicalParent"
+        "logicalParent",
+        "globalTags",
+        "glossaryTerms"
       ]
     },
     "name": "SchemaFieldKey",
@@ -10423,67 +10425,71 @@
             "queryByDefault": false
           }
         },
-        "type": {
-          "type": "record",
-          "name": "Edge",
-          "namespace": "com.linkedin.pegasus2avro.common",
-          "fields": [
-            {
-              "java": {
-                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+        "type": [
+          "null",
+          {
+            "type": "record",
+            "name": "Edge",
+            "namespace": "com.linkedin.pegasus2avro.common",
+            "fields": [
+              {
+                "java": {
+                  "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                },
+                "Urn": "Urn",
+                "type": [
+                  "null",
+                  "string"
+                ],
+                "name": "sourceUrn",
+                "default": null,
+                "doc": "Urn of the source of this relationship edge.\nIf not specified, assumed to be the entity that this aspect belongs to."
               },
-              "Urn": "Urn",
-              "type": [
-                "null",
-                "string"
-              ],
-              "name": "sourceUrn",
-              "default": null,
-              "doc": "Urn of the source of this relationship edge.\nIf not specified, assumed to be the entity that this aspect belongs to."
-            },
-            {
-              "java": {
-                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+              {
+                "java": {
+                  "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                },
+                "Urn": "Urn",
+                "type": "string",
+                "name": "destinationUrn",
+                "doc": "Urn of the destination of this relationship edge."
               },
-              "Urn": "Urn",
-              "type": "string",
-              "name": "destinationUrn",
-              "doc": "Urn of the destination of this relationship edge."
-            },
-            {
-              "type": [
-                "null",
-                "com.linkedin.pegasus2avro.common.AuditStamp"
-              ],
-              "name": "created",
-              "default": null,
-              "doc": "Audit stamp containing who created this relationship edge and when"
-            },
-            {
-              "type": [
-                "null",
-                "com.linkedin.pegasus2avro.common.AuditStamp"
-              ],
-              "name": "lastModified",
-              "default": null,
-              "doc": "Audit stamp containing who last modified this relationship edge and when"
-            },
-            {
-              "type": [
-                "null",
-                {
-                  "type": "map",
-                  "values": "string"
-                }
-              ],
-              "name": "properties",
-              "default": null,
-              "doc": "A generic properties bag that allows us to store specific information on this graph edge."
-            }
-          ],
-          "doc": "A common structure to represent all edges to entities when used inside aspects as collections\nThis ensures that all edges have common structure around audit-stamps and will support PATCH, time-travel automatically."
-        },
-        "name": "parent"
+              {
+                "type": [
+                  "null",
+                  "com.linkedin.pegasus2avro.common.AuditStamp"
+                ],
+                "name": "created",
+                "default": null,
+                "doc": "Audit stamp containing who created this relationship edge and when"
+              },
+              {
+                "type": [
+                  "null",
+                  "com.linkedin.pegasus2avro.common.AuditStamp"
+                ],
+                "name": "lastModified",
+                "default": null,
+                "doc": "Audit stamp containing who last modified this relationship edge and when"
+              },
+              {
+                "type": [
+                  "null",
+                  {
+                    "type": "map",
+                    "values": "string"
+                  }
+                ],
+                "name": "properties",
+                "default": null,
+                "doc": "A generic properties bag that allows us to store specific information on this graph edge."
+              }
+            ],
+            "doc": "A common structure to represent all edges to entities when used inside aspects as collections\nThis ensures that all edges have common structure around audit-stamps and will support PATCH, time-travel automatically."
+          }
+        ],
+        "name": "parent",
+        "default": null
       }
     ]
   },

datahub/metadata/schemas/LogicalParent.avsc

@@ -31,110 +31,114 @@
           "queryByDefault": false
         }
       },
-      "type": {
-        "type": "record",
-        "name": "Edge",
-        "namespace": "com.linkedin.pegasus2avro.common",
-        "fields": [
-          {
-            "java": {
-              "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+      "type": [
+        "null",
+        {
+          "type": "record",
+          "name": "Edge",
+          "namespace": "com.linkedin.pegasus2avro.common",
+          "fields": [
+            {
+              "java": {
+                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+              },
+              "type": [
+                "null",
+                "string"
+              ],
+              "name": "sourceUrn",
+              "default": null,
+              "doc": "Urn of the source of this relationship edge.\nIf not specified, assumed to be the entity that this aspect belongs to.",
+              "Urn": "Urn"
             },
-            "type": [
-              "null",
-              "string"
-            ],
-            "name": "sourceUrn",
-            "default": null,
-            "doc": "Urn of the source of this relationship edge.\nIf not specified, assumed to be the entity that this aspect belongs to.",
-            "Urn": "Urn"
-          },
-          {
-            "java": {
-              "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+            {
+              "java": {
+                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+              },
+              "type": "string",
+              "name": "destinationUrn",
+              "doc": "Urn of the destination of this relationship edge.",
+              "Urn": "Urn"
             },
-            "type": "string",
-            "name": "destinationUrn",
-            "doc": "Urn of the destination of this relationship edge.",
-            "Urn": "Urn"
-          },
-          {
-            "type": [
-              "null",
-              {
-                "type": "record",
-                "name": "AuditStamp",
-                "namespace": "com.linkedin.pegasus2avro.common",
-                "fields": [
-                  {
-                    "type": "long",
-                    "name": "time",
-                    "doc": "When did the resource/association/sub-resource move into the specific lifecycle stage represented by this AuditEvent."
-                  },
-                  {
-                    "java": {
-                      "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+            {
+              "type": [
+                "null",
+                {
+                  "type": "record",
+                  "name": "AuditStamp",
+                  "namespace": "com.linkedin.pegasus2avro.common",
+                  "fields": [
+                    {
+                      "type": "long",
+                      "name": "time",
+                      "doc": "When did the resource/association/sub-resource move into the specific lifecycle stage represented by this AuditEvent."
                     },
-                    "type": "string",
-                    "name": "actor",
-                    "doc": "The entity (e.g. a member URN) which will be credited for moving the resource/association/sub-resource into the specific lifecycle stage. It is also the one used to authorize the change.",
-                    "Urn": "Urn"
-                  },
-                  {
-                    "java": {
-                      "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                    {
+                      "java": {
+                        "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                      },
+                      "type": "string",
+                      "name": "actor",
+                      "doc": "The entity (e.g. a member URN) which will be credited for moving the resource/association/sub-resource into the specific lifecycle stage. It is also the one used to authorize the change.",
+                      "Urn": "Urn"
                     },
-                    "type": [
-                      "null",
-                      "string"
-                    ],
-                    "name": "impersonator",
-                    "default": null,
-                    "doc": "The entity (e.g. a service URN) which performs the change on behalf of the Actor and must be authorized to act as the Actor.",
-                    "Urn": "Urn"
-                  },
-                  {
-                    "type": [
-                      "null",
-                      "string"
-                    ],
-                    "name": "message",
-                    "default": null,
-                    "doc": "Additional context around how DataHub was informed of the particular change. For example: was the change created by an automated process, or manually."
-                  }
-                ],
-                "doc": "Data captured on a resource/association/sub-resource level giving insight into when that resource/association/sub-resource moved into a particular lifecycle stage, and who acted to move it into that specific lifecycle stage."
-              }
-            ],
-            "name": "created",
-            "default": null,
-            "doc": "Audit stamp containing who created this relationship edge and when"
-          },
-          {
-            "type": [
-              "null",
-              "com.linkedin.pegasus2avro.common.AuditStamp"
-            ],
-            "name": "lastModified",
-            "default": null,
-            "doc": "Audit stamp containing who last modified this relationship edge and when"
-          },
-          {
-            "type": [
-              "null",
-              {
-                "type": "map",
-                "values": "string"
-              }
-            ],
-            "name": "properties",
-            "default": null,
-            "doc": "A generic properties bag that allows us to store specific information on this graph edge."
-          }
-        ],
-        "doc": "A common structure to represent all edges to entities when used inside aspects as collections\nThis ensures that all edges have common structure around audit-stamps and will support PATCH, time-travel automatically."
-      },
-      "name": "parent"
+                    {
+                      "java": {
+                        "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+                      },
+                      "type": [
+                        "null",
+                        "string"
+                      ],
+                      "name": "impersonator",
+                      "default": null,
+                      "doc": "The entity (e.g. a service URN) which performs the change on behalf of the Actor and must be authorized to act as the Actor.",
+                      "Urn": "Urn"
+                    },
+                    {
+                      "type": [
+                        "null",
+                        "string"
+                      ],
+                      "name": "message",
+                      "default": null,
+                      "doc": "Additional context around how DataHub was informed of the particular change. For example: was the change created by an automated process, or manually."
+                    }
+                  ],
+                  "doc": "Data captured on a resource/association/sub-resource level giving insight into when that resource/association/sub-resource moved into a particular lifecycle stage, and who acted to move it into that specific lifecycle stage."
+                }
+              ],
+              "name": "created",
+              "default": null,
+              "doc": "Audit stamp containing who created this relationship edge and when"
+            },
+            {
+              "type": [
+                "null",
+                "com.linkedin.pegasus2avro.common.AuditStamp"
+              ],
+              "name": "lastModified",
+              "default": null,
+              "doc": "Audit stamp containing who last modified this relationship edge and when"
+            },
+            {
+              "type": [
+                "null",
+                {
+                  "type": "map",
+                  "values": "string"
+                }
+              ],
+              "name": "properties",
+              "default": null,
+              "doc": "A generic properties bag that allows us to store specific information on this graph edge."
+            }
+          ],
+          "doc": "A common structure to represent all edges to entities when used inside aspects as collections\nThis ensures that all edges have common structure around audit-stamps and will support PATCH, time-travel automatically."
+        }
+      ],
+      "name": "parent",
+      "default": null
     }
   ]
 }

datahub/metadata/schemas/SchemaFieldKey.avsc

@@ -15,7 +15,9 @@
       "testResults",
       "deprecation",
       "subTypes",
-      "logicalParent"
+      "logicalParent",
+      "globalTags",
+      "glossaryTerms"
     ]
   },
   "name": "SchemaFieldKey",