acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/api/entities/dataset/dataset.py

@@ -15,16 +15,23 @@ from typing import (
 
 import avro
 import yaml
-from pydantic import BaseModel, Field, root_validator, validator
+from pydantic import (
+    BaseModel,
+    Field,
+    StrictStr,
+    root_validator,
+    validator,
+)
 from ruamel.yaml import YAML
 from typing_extensions import TypeAlias
 
 import datahub.metadata.schema_classes as models
 from datahub.api.entities.structuredproperties.structuredproperties import AllowedTypes
-from datahub.configuration.common import ConfigModel
+from datahub.configuration.common import ConfigModel, LaxStr
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataset_urn,
+    make_domain_urn,
     make_schema_field_urn,
     make_tag_urn,
     make_term_urn,
@@ -37,6 +44,7 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import (
     AuditStampClass,
     DatasetPropertiesClass,
+    DomainsClass,
     GlobalTagsClass,
     GlossaryTermAssociationClass,
     GlossaryTermsClass,
@@ -90,7 +98,7 @@ class StrictModel(BaseModel):
 
 
 # Define type aliases for the complex types
-PropertyValue: TypeAlias = Union[float, str]
+PropertyValue: TypeAlias = Union[StrictStr, float]
 PropertyValueList: TypeAlias = List[PropertyValue]
 StructuredProperties: TypeAlias = Dict[str, Union[PropertyValue, PropertyValueList]]
 
@@ -128,14 +136,13 @@ class StructuredPropertiesHelper:
 
 class SchemaFieldSpecification(StrictModel):
     id: Optional[str] = None
-    urn: Optional[str] = None
+    urn: Optional[str] = Field(None, validate_default=True)
     structured_properties: Optional[StructuredProperties] = None
     type: Optional[str] = None
     nativeDataType: Optional[str] = None
     jsonPath: Union[None, str] = None
     nullable: bool = False
     description: Union[None, str] = None
-    doc: Union[None, str] = None  # doc is an alias for description
     label: Optional[str] = None
     created: Optional[dict] = None
     lastModified: Optional[dict] = None
@@ -213,14 +220,14 @@ class SchemaFieldSpecification(StrictModel):
         return v
 
     @root_validator(pre=True)
-    def sync_description_and_doc(cls, values: Dict) -> Dict:
-        """Synchronize doc and description fields if one is provided but not the other."""
+    def sync_doc_into_description(cls, values: Dict) -> Dict:
+        """Synchronize doc into description field if doc is provided."""
         description = values.get("description")
-        doc = values.get("doc")
+        doc = values.pop("doc", None)
 
-        if description is not None and doc is None:
-            values["doc"] = description
-        elif doc is not None and description is None:
+        if doc is not None:
+            if description is not None:
+                raise ValueError("doc and description cannot both be provided")
             values["description"] = doc
 
         return values
@@ -288,10 +295,6 @@ class SchemaFieldSpecification(StrictModel):
             """Custom dict method for Pydantic v1 to handle YAML serialization properly."""
             exclude = kwargs.pop("exclude", None) or set()
 
-            # If description and doc are identical, exclude doc from the output
-            if self.description == self.doc and self.description is not None:
-                exclude.add("doc")
-
             # if nativeDataType and type are identical, exclude nativeDataType from the output
             if self.nativeDataType == self.type and self.nativeDataType is not None:
                 exclude.add("nativeDataType")
@@ -319,10 +322,6 @@ class SchemaFieldSpecification(StrictModel):
             """Custom model_dump method for Pydantic v2 to handle YAML serialization properly."""
             exclude = kwargs.pop("exclude", None) or set()
 
-            # If description and doc are identical, exclude doc from the output
-            if self.description == self.doc and self.description is not None:
-                exclude.add("doc")
-
             # if nativeDataType and type are identical, exclude nativeDataType from the output
             if self.nativeDataType == self.type and self.nativeDataType is not None:
                 exclude.add("nativeDataType")
@@ -366,12 +365,6 @@ class Ownership(ConfigModel):
         return v
 
 
-class StructuredPropertyValue(ConfigModel):
-    value: Union[str, int, float, List[str], List[int], List[float]]
-    created: Optional[str] = None
-    lastModified: Optional[str] = None
-
-
 class DatasetRetrievalConfig(BaseModel):
     include_downstreams: Optional[bool] = False
 
@@ -380,12 +373,12 @@ class Dataset(StrictModel):
     id: Optional[str] = None
     platform: Optional[str] = None
     env: str = "PROD"
-    urn: Optional[str] = None
+    urn: Optional[str] = Field(None, validate_default=True)
     description: Optional[str] = None
-    name: Optional[str] = None
-    schema_metadata: Optional[SchemaSpecification] = Field(alias="schema")
+    name: Optional[str] = Field(None, validate_default=True)
+    schema_metadata: Optional[SchemaSpecification] = Field(default=None, alias="schema")
     downstreams: Optional[List[str]] = None
-    properties: Optional[Dict[str, str]] = None
+    properties: Optional[Dict[str, LaxStr]] = None
     subtype: Optional[str] = None
     subtypes: Optional[List[str]] = None
     tags: Optional[List[str]] = None
@@ -393,6 +386,7 @@ class Dataset(StrictModel):
     owners: Optional[List[Union[str, Ownership]]] = None
     structured_properties: Optional[StructuredProperties] = None
     external_url: Optional[str] = None
+    domains: Optional[List[str]] = None
 
     @property
     def platform_urn(self) -> str:
@@ -483,7 +477,7 @@ class Dataset(StrictModel):
                                 f"{urn_prefix}:{prop_key}"
                                 if not prop_key.startswith(urn_prefix)
                                 else prop_key
-                                for prop_key in field.structured_properties.keys()
+                                for prop_key in field.structured_properties
                             ]
                         )
                     if field.glossaryTerms:
@@ -497,7 +491,7 @@ class Dataset(StrictModel):
                     f"{urn_prefix}:{prop_key}"
                     if not prop_key.startswith(urn_prefix)
                     else prop_key
-                    for prop_key in self.structured_properties.keys()
+                    for prop_key in self.structured_properties
                 ]
             )
         if self.glossary_terms:
@@ -506,19 +500,17 @@ class Dataset(StrictModel):
         # We don't check references for tags
         return list(set(references))
 
-    def generate_mcp(  # noqa: C901
+    def generate_mcp(
         self,
     ) -> Iterable[Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper]]:
-        mcp = MetadataChangeProposalWrapper(
-            entityUrn=self.urn,
-            aspect=DatasetPropertiesClass(
-                description=self.description,
-                name=self.name,
-                customProperties=self.properties,
-                externalUrl=self.external_url,
-            ),
-        )
-        yield mcp
+        patch_builder = self.patch_builder()
+
+        patch_builder.set_custom_properties(self.properties or {})
+        patch_builder.set_description(self.description)
+        patch_builder.set_display_name(self.name)
+        patch_builder.set_external_url(self.external_url)
+
+        yield from patch_builder.build()
 
         if self.schema_metadata:
             schema_fields = set()
@@ -604,7 +596,7 @@ class Dataset(StrictModel):
                         ],
                         platformSchema=OtherSchemaClass(
                             rawSchema=yaml.dump(
-                                self.schema_metadata.dict(
+                                self.schema_metadata.model_dump(
                                     exclude_none=True, exclude_unset=True
                                 )
                             )
@@ -643,33 +635,6 @@ class Dataset(StrictModel):
                     )
                     assert field_urn.startswith("urn:li:schemaField:")
 
-                    if field.globalTags:
-                        mcp = MetadataChangeProposalWrapper(
-                            entityUrn=field_urn,
-                            aspect=GlobalTagsClass(
-                                tags=[
-                                    TagAssociationClass(tag=make_tag_urn(tag))
-                                    for tag in field.globalTags
-                                ]
-                            ),
-                        )
-                        yield mcp
-
-                    if field.glossaryTerms:
-                        mcp = MetadataChangeProposalWrapper(
-                            entityUrn=field_urn,
-                            aspect=GlossaryTermsClass(
-                                terms=[
-                                    GlossaryTermAssociationClass(
-                                        urn=make_term_urn(term)
-                                    )
-                                    for term in field.glossaryTerms
-                                ],
-                                auditStamp=self._mint_auditstamp("yaml"),
-                            ),
-                        )
-                        yield mcp
-
                     if field.structured_properties:
                         urn_prefix = f"{StructuredPropertyUrn.URN_PREFIX}:{StructuredPropertyUrn.LI_DOMAIN}:{StructuredPropertyUrn.ENTITY_TYPE}"
                         mcp = MetadataChangeProposalWrapper(
@@ -764,7 +729,14 @@ class Dataset(StrictModel):
                     )
                 )
                 yield from patch_builder.build()
-
+        if self.domains:
+            mcp = MetadataChangeProposalWrapper(
+                entityUrn=self.urn,
+                aspect=DomainsClass(
+                    [make_domain_urn(domain) for domain in self.domains]
+                ),
+            )
+            yield mcp
         logger.info(f"Created dataset {self.urn}")
 
     @staticmethod
@@ -815,6 +787,7 @@ class Dataset(StrictModel):
         if schema_metadata:
             # If the schema is built off of an avro schema, we only extract the fields if they have structured properties
             # Otherwise, we extract all fields
+            schema_fields = []
             if (
                 schema_metadata.platformSchema
                 and isinstance(schema_metadata.platformSchema, models.OtherSchemaClass)
@@ -925,6 +898,7 @@ class Dataset(StrictModel):
                     structured_properties_map[sp.propertyUrn].extend(sp.values)  # type: ignore[arg-type,union-attr]
                 else:
                     structured_properties_map[sp.propertyUrn] = sp.values
+        domains: Optional[DomainsClass] = graph.get_aspect(urn, DomainsClass)
 
         if config.include_downstreams:
             related_downstreams = graph.get_related_entities(
@@ -965,6 +939,7 @@ class Dataset(StrictModel):
             structured_properties=(
                 structured_properties_map if structured_properties else None
             ),
+            domains=[domain for domain in domains.domains] if domains else None,
             downstreams=downstreams if config.include_downstreams else None,
         )
 
@@ -1008,7 +983,7 @@ class Dataset(StrictModel):
 
         def model_dump(self, **kwargs):
             """Custom model_dump method for Pydantic v2 to handle YAML serialization properly."""
-            exclude = kwargs.pop("exclude", set())
+            exclude = kwargs.pop("exclude", None) or set()
 
             # If id and name are identical, exclude name from the output
             if self.id == self.name and self.id is not None: