openmetadata-ingestion 1.2.5.1__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

metadata/ingestion/models/patch_request.py

@@ -11,9 +11,14 @@
 """
 Pydantic definition for storing entities for patching 
 """
+import json
+from typing import Dict, List, Optional
+
+import jsonpatch
 from pydantic import BaseModel
 
-from metadata.ingestion.api.models import Entity
+from metadata.ingestion.api.models import Entity, T
+from metadata.ingestion.ometa.mixins.patch_mixin_utils import PatchOperation
 
 
 class PatchRequest(BaseModel):
@@ -25,6 +30,14 @@ class PatchRequest(BaseModel):
     new_entity: Entity
 
 
+class PatchedEntity(BaseModel):
+    """
+    Store the new entity after patch request
+    """
+
+    new_entity: Optional[Entity]
+
+
 ALLOWED_COLUMN_FIELDS = {
     "name": True,
     "dataType": True,
@@ -65,7 +78,9 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "displayName": True,
     "sourceUrl": True,
     "description": True,
+    "owner": True,
     "tags": True,
+    "sourceHash": True,
     # Table Entity Fields
     "tableType": True,
     "columns": {"__all__": ALLOWED_COLUMN_FIELDS},
@@ -121,4 +136,95 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "fileFormats": True,
 }
 
-RESTRICT_UPDATE_LIST = ["description", "tags"]
+RESTRICT_UPDATE_LIST = ["description", "tags", "owner"]
+
+
+def build_patch(
+    source: T,
+    destination: T,
+    allowed_fields: Optional[Dict] = None,
+    restrict_update_fields: Optional[List] = None,
+) -> Optional[jsonpatch.JsonPatch]:
+    """
+    Given an Entity type and Source entity and Destination entity,
+    generate a JSON Patch and apply it.
+
+    Args
+        source: Source payload which is current state of the source in OpenMetadata
+        destination: payload with changes applied to the source.
+        allowed_fields: List of field names to filter from source and destination models
+        restrict_update_fields: List of field names which will only support add operation
+
+    Returns
+        Updated Entity
+    """
+
+    # remove change descriptions from entities
+    source = _remove_change_description(source)
+    destination = _remove_change_description(destination)
+
+    # 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:
+        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 for replace, remove from here
+    if restrict_update_fields:
+        patch.patch = [
+            patch_ops
+            for patch_ops in patch.patch
+            if _determine_restricted_operation(
+                patch_ops=patch_ops,
+                restrict_update_fields=restrict_update_fields,
+            )
+        ]
+
+    return patch
+
+
+def _determine_restricted_operation(
+    patch_ops: Dict, restrict_update_fields: Optional[List] = None
+) -> bool:
+    """
+    Only retain add operation for restrict_update_fields fields
+    """
+    path = patch_ops.get("path")
+    ops = patch_ops.get("op")
+    for field in restrict_update_fields or []:
+        if field in path and ops != PatchOperation.ADD.value:
+            return False
+    return True
+
+
+def _remove_change_description(entity: T) -> T:
+    """
+    Remove change description if applies.
+    We never want to patch that, and we won't have that information
+    from the source. It's fully handled in the server.
+    """
+    if getattr(entity, "changeDescription"):
+        entity.changeDescription = None
+
+    return entity

metadata/ingestion/models/tests_data.py

@@ -17,6 +17,9 @@ from typing import List
 from pydantic import BaseModel
 
 from metadata.generated.schema.api.tests.createTestCase import CreateTestCaseRequest
+from metadata.generated.schema.api.tests.createTestCaseResolutionStatus import (
+    CreateTestCaseResolutionStatus,
+)
 from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteRequest
 from metadata.generated.schema.tests.basic import TestCaseResult
 from metadata.generated.schema.tests.testCase import TestCase
@@ -38,3 +41,9 @@ class OMetaTestCaseSample(BaseModel):
 class OMetaTestCaseResultsSample(BaseModel):
     test_case_results: TestCaseResult
     test_case_name: str
+
+
+class OMetaTestCaseResolutionStatus(BaseModel):
+    """For sample data"""
+
+    test_case_resolution: CreateTestCaseResolutionStatus

metadata/ingestion/models/topology.py

@@ -11,12 +11,20 @@
 """
 Defines the topology for ingesting sources
 """
-
+from functools import singledispatchmethod
 from typing import Any, Generic, List, Optional, Type, TypeVar
 
-from pydantic import BaseModel, Extra, create_model
+from pydantic import BaseModel, Extra, Field, create_model
+
+from metadata.generated.schema.api.data.createStoredProcedure import (
+    CreateStoredProcedureRequest,
+)
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
+from metadata.ingestion.ometa.utils import model_str
+from metadata.utils import fqn
 
 T = TypeVar("T", bound=BaseModel)
+C = TypeVar("C", bound=BaseModel)
 
 
 class NodeStage(BaseModel, Generic[T]):
@@ -30,26 +38,54 @@ class NodeStage(BaseModel, Generic[T]):
     class Config:
         extra = Extra.forbid
 
-    type_: Type[T]  # Entity type
-    processor: str  # has the producer results as an argument. Here is where filters happen
-    context: Optional[
-        str
-    ] = None  # context key storing stage state, if needed. This requires us to ACK the ingestion
-    nullable: bool = False  # The yielded value can be null
-    must_return: bool = False  # The sink MUST return a value back after ack. Useful to validate services are correct.
-    cache_all: bool = (
-        False  # If we need to cache all values being yielded in the context
+    # Required fields to define the yielded entity type and the function processing it
+    type_: Type[T] = Field(
+        ..., description="Entity Type. E.g., DatabaseService, Database or Table"
+    )
+    processor: str = Field(
+        ...,
+        description="Has the producer results as an argument. Here is where filters happen. It will yield an Entity.",
+    )
+
+    # Topology behavior
+    nullable: bool = Field(False, description="Flags if the yielded value can be null")
+    must_return: bool = Field(
+        False,
+        description="The sink MUST return a value back after ack. Useful to validate if services are correct.",
+    )
+    overwrite: bool = Field(
+        True,
+        description="If we want to update existing data from OM. E.g., we don't want to overwrite services.",
+    )
+    consumer: Optional[List[str]] = Field(
+        None,
+        description="Stage dependency from parent nodes. Used to build the FQN of the processed Entity.",
+    )
+
+    # Context-related flags
+    context: Optional[str] = Field(
+        None, description="Context key storing stage state, if needed"
+    )
+    store_all_in_context: bool = Field(
+        False, description="If we need to store all values being yielded in the context"
+    )
+    clear_context: bool = Field(
+        False,
+        description="If we need to clean the values in the context for each produced element",
+    )
+    store_fqn: bool = Field(
+        False,
+        description="If true, store the entity FQN in the context instead of just the name",
     )
-    clear_cache: bool = False  # If we need to clean cache values in the context for each produced element
-    overwrite: bool = True  # If we want to overwrite existing data from OM
-    consumer: Optional[
-        List[str]
-    ] = None  # keys in the source context to fetch state from the parent's context
-    cache_entities: bool = (
-        False  # Cache all the entities which have use_cache set as True
+
+    # Used to compute the fingerprint
+    cache_entities: bool = Field(
+        False,
+        description="Cache all the entities which have use_cache set as True. Used for fingerprint comparison.",
     )
-    use_cache: bool = (
-        False  # enable this to get the entity from cached state in the context
+    use_cache: bool = Field(
+        False,
+        description="Enable this to get the entity from cached state in the context",
     )
 
 
@@ -64,18 +100,21 @@ class TopologyNode(BaseModel):
     class Config:
         extra = Extra.forbid
 
-    # method name in the source to use to generate the data to process
-    # does not accept input parameters
-    producer: str
-
-    # list of functions to execute - in order - for each element produced by the producer
-    # each stage accepts the producer results as an argument
-    stages: List[NodeStage]
-
-    children: Optional[List[str]] = None  # nodes to call execute next
-    post_process: Optional[
-        List[str]
-    ] = None  # Method to be run after the node has been fully processed
+    producer: str = Field(
+        ...,
+        description="Method name in the source called to generate the data. Does not accept input parameters",
+    )
+    stages: List[NodeStage] = Field(
+        ...,
+        description=(
+            "List of functions to execute - in order - for each element produced by the producer. "
+            "Each stage accepts the producer results as an argument"
+        ),
+    )
+    children: Optional[List[str]] = Field(None, description="Nodes to execute next")
+    post_process: Optional[List[str]] = Field(
+        None, description="Method to be run after the node has been fully processed"
+    )
 
 
 class ServiceTopology(BaseModel):
@@ -99,6 +138,112 @@ class TopologyContext(BaseModel):
         ctx = {key: value.name.__root__ for key, value in self.__dict__.items()}
         return f"TopologyContext({ctx})"
 
+    @classmethod
+    def create(cls, topology: ServiceTopology) -> "TopologyContext":
+        """
+        Dynamically build a context based on the topology nodes.
+
+        Builds a Pydantic BaseModel class.
+
+        :param topology: ServiceTopology
+        :return: TopologyContext
+        """
+        nodes = get_topology_nodes(topology)
+        ctx_fields = {
+            stage.context: (Optional[stage.type_], None)
+            for node in nodes
+            for stage in node.stages
+            if stage.context
+        }
+        return create_model(
+            "GeneratedContext", **ctx_fields, __base__=TopologyContext
+        )()
+
+    def upsert(self, key: str, value: Any) -> None:
+        """
+        Update the key of the context with the given value
+        :param key: element to update from the source context
+        :param value: value to use for the update
+        """
+        self.__dict__[key] = value
+
+    def append(self, key: str, value: Any) -> None:
+        """
+        Update the key of the context with the given value
+        :param key: element to update from the source context
+        :param value: value to use for the update
+        """
+        if self.__dict__.get(key):
+            self.__dict__[key].append(value)
+        else:
+            self.__dict__[key] = [value]
+
+    def clear_stage(self, stage: NodeStage) -> None:
+        """
+        Clear the available context
+        :param stage: Update stage context to the default values
+        """
+        self.__dict__[stage.context] = None
+
+    def fqn_from_stage(self, stage: NodeStage, entity_name: str) -> str:
+        """
+        Read the context
+        :param stage: Topology node being processed
+        :param entity_name: name being stored
+        :return: Entity FQN derived from context
+        """
+        context_names = [
+            self.__dict__[dependency]
+            for dependency in stage.consumer or []  # root nodes do not have consumers
+        ]
+        return fqn._build(  # pylint: disable=protected-access
+            *context_names, entity_name
+        )
+
+    def update_context_name(self, stage: NodeStage, right: C) -> None:
+        """
+        Append or update context
+
+        We'll store the entity name or FQN in the topology context.
+        If we store the name, the FQN will be built in the source itself when needed.
+        """
+
+        if stage.store_fqn:
+            new_context = self._build_new_context_fqn(right)
+        else:
+            new_context = model_str(right.name)
+
+        self.update_context_value(stage=stage, value=new_context)
+
+    def update_context_value(self, stage: NodeStage, value: Any) -> None:
+        if stage.context and not stage.store_all_in_context:
+            self.upsert(key=stage.context, value=value)
+        if stage.context and stage.store_all_in_context:
+            self.append(key=stage.context, value=value)
+
+    @singledispatchmethod
+    def _build_new_context_fqn(self, right: C) -> str:
+        """Build context fqn string"""
+        raise NotImplementedError(f"Missing implementation for [{type(C)}]")
+
+    @_build_new_context_fqn.register
+    def _(self, right: CreateStoredProcedureRequest) -> str:
+        """
+        Implement FQN context building for Stored Procedures.
+
+        We process the Stored Procedures lineage at the very end of the service. If we
+        just store the SP name, we lose the information of which db/schema the SP belongs to.
+        """
+
+        return fqn.build(
+            metadata=None,
+            entity_type=StoredProcedure,
+            service_name=self.__dict__["database_service"],
+            database_name=self.__dict__["database"],
+            schema_name=self.__dict__["database_schema"],
+            procedure_name=right.name.__root__,
+        )
+
 
 def get_topology_nodes(topology: ServiceTopology) -> List[TopologyNode]:
     """
@@ -132,34 +277,6 @@ def get_topology_root(topology: ServiceTopology) -> List[TopologyNode]:
     return [node for node in nodes if node_has_no_consumers(node)]
 
 
-def get_ctx_default(stage: NodeStage) -> Optional[List[Any]]:
-    """
-    If we cache all, default value is an empty list
-    :param stage: Node Stage
-    :return: None or []
-    """
-    return [] if stage.cache_all else None
-
-
-def create_source_context(topology: ServiceTopology) -> TopologyContext:
-    """
-    Dynamically build a context based on the topology nodes.
-
-    Builds a Pydantic BaseModel class.
-
-    :param topology: ServiceTopology
-    :return: TopologyContext
-    """
-    nodes = get_topology_nodes(topology)
-    ctx_fields = {
-        stage.context: (Optional[stage.type_], get_ctx_default(stage))
-        for node in nodes
-        for stage in node.stages
-        if stage.context
-    }
-    return create_model("GeneratedContext", **ctx_fields, __base__=TopologyContext)()
-
-
 def get_topology_node(name: str, topology: ServiceTopology) -> TopologyNode:
     """
     Fetch a topology node by name