openmetadata-ingestion 1.4.7.2__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/ingestion/lineage/sql_lineage.py

@@ -11,6 +11,7 @@
 """
 Helper functions to handle SQL lineage operations
 """
+import itertools
 import traceback
 from typing import Any, Iterable, List, Optional, Tuple
 
@@ -50,8 +51,8 @@ def get_column_fqn(table_entity: Table, column: str) -> Optional[str]:
     if not table_entity:
         return None
     for tbl_column in table_entity.columns:
-        if column.lower() == tbl_column.name.__root__.lower():
-            return tbl_column.fullyQualifiedName.__root__
+        if column.lower() == tbl_column.name.root.lower():
+            return tbl_column.fullyQualifiedName.root
 
     return None
 
@@ -226,7 +227,7 @@ def get_column_lineage(
         # Select all
         if "*" in column_lineage_map.get(to_table_raw_name).get(from_table_raw_name)[0]:
             column_lineage_map[to_table_raw_name][from_table_raw_name] = [
-                (c.name.__root__, c.name.__root__) for c in from_entity.columns
+                (c.name.root, c.name.root) for c in from_entity.columns
             ]
 
         # Other cases
@@ -250,11 +251,23 @@ def _build_table_lineage(
     query: str,
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
-) -> Iterable[Either[AddLineageRequest]]:
+) -> Either[AddLineageRequest]:
     """
     Prepare the lineage request generator
+
+    Args:
+        from_entity (Table): entity link comes from
+        to_entity (Table): entity to link to
+        from_table_raw_name (str): table entity raw name we link from
+        to_table_raw_name (str): table entity raw name we link to
+        query (str): query
+        column_lineage_map (dict): map of the column lineage
+        lineage_source (LineageSource): lineage source
+
+    Returns:
+        Either[AddLineageRequest] with the lineage request or an error
     """
-    if from_entity and to_entity:
+    try:
         col_lineage = get_column_lineage(
             to_entity=to_entity,
             to_table_raw_name=str(to_table_raw_name),
@@ -268,21 +281,29 @@ def _build_table_lineage(
         lineage = AddLineageRequest(
             edge=EntitiesEdge(
                 fromEntity=EntityReference(
-                    id=from_entity.id.__root__,
+                    id=from_entity.id.root,
                     type="table",
                 ),
                 toEntity=EntityReference(
-                    id=to_entity.id.__root__,
+                    id=to_entity.id.root,
                     type="table",
                 ),
             )
         )
         if lineage_details:
             lineage.edge.lineageDetails = lineage_details
-        yield Either(right=lineage)
+        return Either(right=lineage)
+    except Exception as e:
+        return Either(
+            left=StackTraceError(
+                name="Lineage",
+                error=f"Error creating lineage for tables [{from_table_raw_name}] and [{to_table_raw_name}]: {e}",
+                stackTrace=traceback.format_exc(),
+            )
+        )
 
 
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments,too-many-locals
 def _create_lineage_by_table_name(
     metadata: OpenMetadata,
     from_table: str,
@@ -315,9 +336,18 @@ def _create_lineage_by_table_name(
             table_name=to_table,
         )
 
-        for from_entity in from_table_entities or []:
-            for to_entity in to_table_entities or []:
-                yield from _build_table_lineage(
+        for table_name, entity in (
+            (from_table, from_table_entities),
+            (to_table, to_table_entities),
+        ):
+            if entity is None:
+                raise RuntimeError(f"Table entity not found: [{table_name}]")
+
+        for from_entity, to_entity in itertools.product(
+            from_table_entities, to_table_entities
+        ):
+            if to_entity and from_entity:
+                yield _build_table_lineage(
                     to_entity=to_entity,
                     from_entity=from_entity,
                     to_table_raw_name=to_table,
@@ -466,7 +496,7 @@ def get_lineage_via_table_entity(
     try:
         logger.debug(f"Getting lineage via table entity using query: {query}")
         lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
-        to_table_name = table_entity.name.__root__
+        to_table_name = table_entity.name.root
 
         for from_table_name in lineage_parser.source_tables:
             yield from _create_lineage_by_table_name(

metadata/ingestion/models/custom_properties.py

@@ -54,12 +54,12 @@ class CustomPropertyType(BaseModel):
 
     id: basic.Uuid
     name: basic.EntityName
-    displayName: Optional[str]
-    fullyQualifiedName: Optional[basic.FullyQualifiedEntityName]
-    description: Optional[basic.Markdown]
-    category: Optional[str]
-    nameSpace: Optional[str]
-    version: Optional[entityHistory.EntityVersion]
-    updatedAt: Optional[basic.Timestamp]
-    updatedBy: Optional[str]
-    href: Optional[basic.Href]
+    displayName: Optional[str] = None
+    fullyQualifiedName: Optional[basic.FullyQualifiedEntityName] = None
+    description: Optional[basic.Markdown] = None
+    category: Optional[str] = None
+    nameSpace: Optional[str] = None
+    version: Optional[entityHistory.EntityVersion] = None
+    updatedAt: Optional[basic.Timestamp] = None
+    updatedBy: Optional[str] = None
+    href: Optional[basic.Href] = None

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)
         )