openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/ingestion/models/custom_pydantic.py

@@ -15,69 +15,81 @@ This classes are used in the generated module, which should have NO
 dependencies against any other metadata package. This class should
 be self-sufficient with only pydantic at import time.
 """
+import json
 import logging
-import warnings
-from typing import Any, Dict
+from typing import Any, Dict, Literal, Optional, Union
 
-from pydantic.types import OptionalInt, SecretStr
-from pydantic.utils import update_not_none
-from pydantic.validators import constr_length_validator, str_validator
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic import PlainSerializer
+from pydantic.main import IncEx
+from pydantic.types import SecretStr
+from typing_extensions import Annotated
 
 logger = logging.getLogger("metadata")
 
 SECRET = "secret:"
+JSON_ENCODERS = "json_encoders"
 
 
-class CustomSecretStr(SecretStr):
+class BaseModel(PydanticBaseModel):
     """
-    Custom SecretStr class which use the configured Secrets Manager to retrieve the actual values.
-
-    If the secret string value starts with `config:` it will use the rest of the string as secret id to search for it
-    in the secrets store.
+    Base model for OpenMetadata generated models.
+    Specified as `--base-class BASE_CLASS` in the generator.
     """
 
-    min_length: OptionalInt = None
-    max_length: OptionalInt = None
-
-    @classmethod
-    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
-        update_not_none(
-            field_schema,
-            type="string",
-            writeOnly=True,
-            format="password",
-            minLength=cls.min_length,
-            maxLength=cls.max_length,
+    def model_dump_json(  # pylint: disable=too-many-arguments
+        self,
+        *,
+        indent: Optional[int] = None,
+        include: IncEx = None,
+        exclude: IncEx = None,
+        context: Optional[Dict[str, Any]] = None,
+        by_alias: bool = False,
+        exclude_unset: bool = True,
+        exclude_defaults: bool = False,
+        exclude_none: bool = True,
+        round_trip: bool = False,
+        warnings: Union[bool, Literal["none", "warn", "error"]] = True,
+        serialize_as_any: bool = False,
+    ) -> str:
+        """
+        This is needed due to https://github.com/pydantic/pydantic/issues/8825
+
+        We also tried the suggested `serialize` method but it did not
+        work well with nested models.
+
+        This solution is covered in the `test_pydantic_v2` test comparing the
+        dump results from V1 vs. V2.
+        """
+        return json.dumps(
+            self.model_dump(
+                mode="json",
+                include=include,
+                exclude=exclude,
+                context=context,
+                by_alias=by_alias,
+                exclude_unset=exclude_unset,
+                exclude_none=exclude_none,
+                exclude_defaults=exclude_defaults,
+                round_trip=round_trip,
+                warnings=warnings,
+                serialize_as_any=serialize_as_any,
+            ),
+            ensure_ascii=True,
         )
 
-    @classmethod
-    def __get_validators__(cls) -> "CallableGenerator":
-        yield cls.validate
-        yield constr_length_validator
 
-    @classmethod
-    def validate(cls, value: Any) -> "CustomSecretStr":
-        if isinstance(value, cls):
-            return value
-        value = str_validator(value)
-        return cls(value)
+class _CustomSecretStr(SecretStr):
+    """
+    Custom SecretStr class which use the configured Secrets Manager to retrieve the actual values.
 
-    def __init__(self, value: str):
-        self._secret_value = value
+    If the secret string value starts with `config:` it will use the rest of the string as secret id to search for it
+    in the secrets store.
+    """
 
     def __repr__(self) -> str:
         return f"SecretStr('{self}')"
 
-    def __len__(self) -> int:
-        return len(self._secret_value)
-
-    def display(self) -> str:
-        warnings.warn(
-            "`secret_str.display()` is deprecated, use `str(secret_str)` instead",
-            DeprecationWarning,
-        )
-        return str(self)
-
     def get_secret_value(self, skip_secret_manager: bool = False) -> str:
         """
         This function should only be called after the SecretsManager has properly
@@ -108,3 +120,17 @@ class CustomSecretStr(SecretStr):
                     f"Secret value [{secret_id}] not present in the configured secrets manager: {exc}"
                 )
         return self._secret_value
+
+
+CustomSecretStr = Annotated[
+    _CustomSecretStr, PlainSerializer(lambda secret: secret.get_secret_value())
+]
+
+
+def ignore_type_decoder(type_: Any) -> None:
+    """Given a type_, add a custom decoder to the BaseModel
+    to ignore any decoding errors for that type_."""
+    # We don't import the constants from the constants module to avoid circular imports
+    BaseModel.model_config[JSON_ENCODERS][type_] = {
+        lambda v: v.decode("utf-8", "ignore")
+    }

metadata/ingestion/models/lf_tags_model.py

@@ -28,6 +28,6 @@ class LFTagsOnColumnsItem(BaseModel):
 
 
 class LFTags(BaseModel):
-    LFTagOnDatabase: Optional[List[TagItem]]
-    LFTagsOnTable: Optional[List[TagItem]]
-    LFTagsOnColumns: Optional[List[LFTagsOnColumnsItem]]
+    LFTagOnDatabase: Optional[List[TagItem]] = None
+    LFTagsOnTable: Optional[List[TagItem]] = None
+    LFTagsOnColumns: Optional[List[LFTagsOnColumnsItem]] = None

metadata/ingestion/models/ometa_classification.py

@@ -24,6 +24,6 @@ from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 
 
 class OMetaTagAndClassification(BaseModel):
-    fqn: Optional[FullyQualifiedEntityName]
+    fqn: Optional[FullyQualifiedEntityName] = None
     classification_request: CreateClassificationRequest
     tag_request: CreateTagRequest

metadata/ingestion/models/ometa_lineage.py

@@ -0,0 +1,24 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Custom wrapper for Lineage Request
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+
+
+class OMetaLineageRequest(BaseModel):
+    override_lineage: Optional[bool] = False
+    lineage_request: AddLineageRequest

metadata/ingestion/models/patch_request.py

@@ -12,6 +12,8 @@
 Pydantic definition for storing entities for patching
 """
 import json
+import logging
+import traceback
 from typing import Dict, List, Optional, Tuple
 
 import jsonpatch
@@ -21,6 +23,8 @@ from metadata.ingestion.api.models import Entity, T
 from metadata.ingestion.ometa.mixins.patch_mixin_utils import PatchOperation
 from metadata.ingestion.ometa.utils import model_str
 
+logger = logging.getLogger("metadata")
+
 
 class PatchRequest(BaseModel):
     """
@@ -29,6 +33,7 @@ class PatchRequest(BaseModel):
 
     original_entity: Entity
     new_entity: Entity
+    override_metadata: Optional[bool] = False
 
 
 class PatchedEntity(BaseModel):
@@ -36,11 +41,12 @@ class PatchedEntity(BaseModel):
     Store the new entity after patch request
     """
 
-    new_entity: Optional[Entity]
+    new_entity: Optional[Entity] = None
 
 
 ALLOWED_COLUMN_FIELDS = {
     "name": True,
+    "displayName": True,
     "dataType": True,
     "arrayDataType": True,
     "description": True,
@@ -79,7 +85,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "displayName": True,
     "sourceUrl": True,
     "description": True,
-    "owner": True,
+    "owners": True,
     "tags": True,
     "sourceHash": True,
     # Table Entity Fields
@@ -137,7 +143,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "fileFormats": True,
 }
 
-RESTRICT_UPDATE_LIST = ["description", "tags", "owner"]
+RESTRICT_UPDATE_LIST = ["description", "tags", "owners", "displayName"]
 
 ARRAY_ENTITY_FIELDS = ["columns", "tasks", "fields"]
 
@@ -260,15 +266,26 @@ class JsonPatchUpdater:
             replace_with_none_op_fixer=ReplaceWithNoneOpFixer.default(),
         )
 
-    def _determine_restricted_operation(self, patch_ops: Dict) -> bool:
+    def _determine_restricted_operation(
+        self, patch_ops: Dict, override_metadata: bool
+    ) -> bool:
         """
         Only retain add operation for restrict_update_fields fields
         """
         path = patch_ops.get("path")
         ops = patch_ops.get("op")
         for field in self.restrict_update_fields or []:
-            if field in path and ops != PatchOperation.ADD.value:
-                return False
+            if field in path:
+                # if we have overrideMetadata enabled we will allow ADD and REPLACE operations
+                if override_metadata:
+                    # REMOVE operations will be skipped since this removes any data on the field
+                    # that is added by the user, if the source has no data on the field
+                    if ops == PatchOperation.REMOVE.value:
+                        return False
+                    return True
+                # if we have overrideMetadata disabled we will only allow ADD operations
+                if ops != PatchOperation.ADD.value:
+                    return False
         return True
 
     def _is_replace_with_none_operation(self, patch_ops: dict) -> bool:
@@ -281,7 +298,9 @@ class JsonPatchUpdater:
         """Returns the Remove operation for the given Path. Used to fix the Replace to None operations."""
         return self.replace_with_none_op_fixer.get_remove_operation(path.split("/"))
 
-    def update(self, patch: jsonpatch.JsonPatch) -> List:
+    def update(
+        self, patch: jsonpatch.JsonPatch, override_metadata: bool = False
+    ) -> List:
         """Given a JSONPatch generated by the jsonpatch library, updates it based on our custom needs.
         1. Remove any restricted operations
         2. Fix any 'Replace to None' operation by adding a 'Remove' operation at the end.
@@ -290,7 +309,9 @@ class JsonPatchUpdater:
         remove_ops_list = []
 
         for patch_ops in patch.patch or []:
-            if self._determine_restricted_operation(patch_ops=patch_ops):
+            if self._determine_restricted_operation(
+                patch_ops=patch_ops, override_metadata=override_metadata
+            ):
                 patch_ops_list.append(patch_ops)
 
                 if self._is_replace_with_none_operation(patch_ops):
@@ -312,6 +333,7 @@ def build_patch(
     restrict_update_fields: Optional[List] = None,
     array_entity_fields: Optional[List] = None,
     remove_change_description: bool = True,
+    override_metadata: Optional[bool] = False,
 ) -> Optional[jsonpatch.JsonPatch]:
     """
     Given an Entity type and Source entity and Destination entity,
@@ -326,55 +348,73 @@ def build_patch(
     Returns
         Updated Entity
     """
-
-    # remove change descriptions from entities
-    if remove_change_description:
-        source = _remove_change_description(source)
-        destination = _remove_change_description(destination)
-
-    if array_entity_fields:
-        _sort_array_entity_fields(
-            source=source,
-            destination=destination,
-            array_entity_fields=array_entity_fields,
-        )
-
-    # Get the difference between source and destination
-    if allowed_fields:
-        patch = jsonpatch.make_patch(
-            json.loads(
-                source.json(
-                    exclude_unset=True,
-                    exclude_none=True,
-                    include=allowed_fields,
-                )
-            ),
-            json.loads(
-                destination.json(
-                    exclude_unset=True,
-                    exclude_none=True,
-                    include=allowed_fields,
-                )
-            ),
-        )
-    else:
-        patch: jsonpatch.JsonPatch = jsonpatch.make_patch(
-            json.loads(source.json(exclude_unset=True, exclude_none=True)),
-            json.loads(destination.json(exclude_unset=True, exclude_none=True)),
-        )
-    if not patch:
+    try:
+        # remove change descriptions from entities
+        if remove_change_description:
+            source = _remove_change_description(source)
+            destination = _remove_change_description(destination)
+
+        if array_entity_fields:
+            _sort_array_entity_fields(
+                source=source,
+                destination=destination,
+                array_entity_fields=array_entity_fields,
+            )
+
+        # Get the difference between source and destination
+        if allowed_fields:
+            patch = jsonpatch.make_patch(
+                json.loads(
+                    source.model_dump_json(
+                        exclude_unset=True,
+                        exclude_none=True,
+                        include=allowed_fields,
+                    )
+                ),
+                json.loads(
+                    destination.model_dump_json(
+                        exclude_unset=True,
+                        exclude_none=True,
+                        include=allowed_fields,
+                    )
+                ),
+            )
+        else:
+            patch: jsonpatch.JsonPatch = jsonpatch.make_patch(
+                json.loads(
+                    source.model_dump_json(exclude_unset=True, exclude_none=True)
+                ),
+                json.loads(
+                    destination.model_dump_json(exclude_unset=True, exclude_none=True)
+                ),
+            )
+        if not patch:
+            return None
+
+        # For a user editable fields like descriptions, tags we only want to support "add" operation in patch
+        # we will remove the other operations.
+        # This will only be applicable if the override_metadata field is set to False.
+        if restrict_update_fields:
+            updated_operations = JsonPatchUpdater.from_restrict_update_fields(
+                restrict_update_fields
+            ).update(patch, override_metadata=override_metadata)
+
+            # if the only operation is to update the sourceHash, we'll skip the patch
+            # since this will only happen when we filter out the REMOVE and REPLACE ops
+            # and if the sourceHash is updated in this case further updates will not be processed
+            if (
+                len(updated_operations) == 1
+                and updated_operations[0].get("path") == "/sourceHash"
+            ):
+                return None
+            patch.patch = updated_operations
+
+        return patch
+    except Exception:
+        logger.debug(traceback.format_exc())
+        logger.warning("Couldn't build patch for Entity.")
         return None
 
-    # For a user editable fields like descriptions, tags we only want to support "add" operation in patch
-    # we will remove the other operations.
-    if restrict_update_fields:
-        updated_operations = JsonPatchUpdater.from_restrict_update_fields(
-            restrict_update_fields
-        ).update(patch)
-        patch.patch = updated_operations
-
-    return patch
-
 
 def _sort_array_entity_fields(
     source: T,
@@ -400,7 +440,7 @@ def _sort_array_entity_fields(
                 destination_attr = destination_dict.get(model_str(source_attr.name))
                 if destination_attr:
                     updated_attributes.append(
-                        source_attr.copy(update=destination_attr.__dict__)
+                        source_attr.model_copy(update=destination_attr.__dict__)
                     )
                     # Remove the updated attribute from the destination dictionary
                     del destination_dict[model_str(source_attr.name)]

metadata/ingestion/models/table_metadata.py

@@ -26,8 +26,8 @@ class OMetaTableConstraints(BaseModel):
     """
 
     table: Table
-    foreign_constraints: Optional[List[Dict]]
-    constraints: Optional[List[TableConstraint]]
+    foreign_constraints: Optional[List[Dict]] = None
+    constraints: Optional[List[TableConstraint]] = None
 
 
 class ColumnTag(BaseModel):

metadata/ingestion/models/topology.py

@@ -16,11 +16,12 @@ import threading
 from functools import singledispatchmethod
 from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
 
-from pydantic import BaseModel, Extra, Field, create_model
+from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
+from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.ingestion.ometa.utils import model_str
 from metadata.utils import fqn
@@ -37,8 +38,9 @@ class NodeStage(BaseModel, Generic[T]):
     source.
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     # Required fields to define the yielded entity type and the function processing it
     type_: Type[T] = Field(
@@ -99,8 +101,9 @@ class TopologyNode(BaseModel):
     with the updated element from the OM API.
     """
 
-    class Config:
-        extra = Extra.forbid
+    model_config = ConfigDict(
+        extra="forbid",
+    )
 
     producer: str = Field(
         ...,
@@ -128,8 +131,7 @@ class ServiceTopology(BaseModel):
     Bounds all service topologies
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
 
 class TopologyContext(BaseModel):
@@ -137,11 +139,10 @@ class TopologyContext(BaseModel):
     Bounds all topology contexts
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     def __repr__(self):
-        ctx = {key: value.name.__root__ for key, value in self.__dict__.items()}
+        ctx = {key: value.name.root for key, value in self.__dict__.items()}
         return f"TopologyContext({ctx})"
 
     @classmethod
@@ -202,6 +203,11 @@ class TopologyContext(BaseModel):
             self.__dict__[dependency]
             for dependency in stage.consumer or []  # root nodes do not have consumers
         ]
+
+        # DashboardDataModel requires extra parameter to build the correct FQN
+        if stage.type_ == DashboardDataModel:
+            context_names.append("model")
+
         return fqn._build(  # pylint: disable=protected-access
             *context_names, entity_name
         )
@@ -247,7 +253,7 @@ class TopologyContext(BaseModel):
             service_name=self.__dict__["database_service"],
             database_name=self.__dict__["database"],
             schema_name=self.__dict__["database_schema"],
-            procedure_name=right.name.__root__,
+            procedure_name=right.name.root,
         )
 
 
@@ -296,7 +302,7 @@ class TopologyContextManager:
 
         # If it does not exist yet, copies the Parent Context in order to have all context gathered until this point.
         self.contexts.setdefault(
-            thread_id, self.contexts[parent_thread_id].copy(deep=True)
+            thread_id, self.contexts[parent_thread_id].model_copy(deep=True)
         )