openmetadata-ingestion 1.3.3.0rc2__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

metadata/ingestion/models/patch_request.py

@@ -9,10 +9,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """
-Pydantic definition for storing entities for patching 
+Pydantic definition for storing entities for patching
 """
 import json
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import jsonpatch
 from pydantic import BaseModel
@@ -88,7 +88,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "tableConstraints": True,
     "tablePartition": True,
     "location": True,
-    "viewDefinition": True,
+    "schemaDefinition": True,
     "sampleData": True,
     "fileFormat": True,
     # Stored Procedure Fields
@@ -142,12 +142,176 @@ RESTRICT_UPDATE_LIST = ["description", "tags", "owner"]
 ARRAY_ENTITY_FIELDS = ["columns", "tasks", "fields"]
 
 
+PathTuple = Tuple[str]
+
+
+# For each 'replace to None' operation we will add a Remove operation at the end.
+# This helps up prevent sending to the backend lists with None values for instance.
+#
+# Obs: The `replace to None` operations are created by the jsonpatch library.
+#
+# Example
+# ----
+# Initial:
+# [
+#   {"op": "replace", "path": "/path/1", "value": None},
+#   {"op": "add",     "path": "/path/2", "value": "foo"},
+# ]
+#
+# Final:
+# [
+#   {"op": "replace", "path": "/path/1", "value": None},
+#   {"op": "add",     "path": "/path/2", "value": "foo"},
+#   {"op": "remove",  "path": "/path/1"},
+# ]
+class ReplaceWithNoneOpFixer:
+    """Responsible for creating the Remove operations that fix every
+    Replace to None ones.
+
+    It is important to keep the state of the index drift in order to
+    create a working patch.
+
+    Example:
+
+        Initial:
+        [
+            {"op": "replace", "path": "/path/1", "value": None},
+            {"op": "replace", "path": "/path/2", "value": None}
+        ]
+
+       Final:
+        [
+            {"op": "replace", "path": "/path/1", "value": None},
+            {"op": "replace", "path": "/path/2", "value": None},
+            {"op": "remove",  "path": "/path/1"},
+            {"op": "remove",  "path": "/path/1"}
+        ]
+
+        Since the first Remove operation is relative to the first replace, there is no drift.
+        When the second Remove operation happens however, the first one would have already been done.
+        This means that '/path/2' becomes '/path/1'.
+    """
+
+    def __init__(self, index_drift_map: Dict[PathTuple, int]):
+        self.index_drift_map = index_drift_map
+
+    @classmethod
+    def default(cls) -> "ReplaceWithNoneOpFixer":
+        """Instantiates the ReplaceWithNoOpFixer with an empty drift map."""
+        return cls(index_drift_map={})
+
+    def _fix_index_drift(self, path: List[str]):
+        """Modifies the incoming path depending on how many Remove operations we have already
+        registered for this path."""
+
+        # We check all the paths for which we already registered a Remove operation
+        for drifted_path, drift in self.index_drift_map.items():
+            # If any of them matches the start of the current path we update the index.
+            if path[: len(drifted_path)] == list(drifted_path):
+                try:
+                    drift_location = len(drifted_path)
+                    path[drift_location] = str(int(path[drift_location]) - drift)
+                except ValueError:
+                    # Not in a List. No need to fix the Path Index
+                    continue
+        return path
+
+    def _update_index_drift_map(self, path: List[str]):
+        """Update the dirft map with the seen path."""
+        path_tuple: PathTuple = tuple(path[:-1])
+
+        self.index_drift_map[path_tuple] = (
+            self.index_drift_map.setdefault(path_tuple, 0) + 1
+        )
+
+    def _get_remove_operation(self, path: List[str]) -> Dict:
+        """Return a JSONPatch Remove operation for the given path."""
+        return {"op": PatchOperation.REMOVE.value, "path": "/".join(path)}
+
+    def get_remove_operation(self, path: List[str]):
+        """Returns a JSONPatch Remove operation for the given path
+        while keeping in the state that we are sending a Remove operation
+        for the given path."""
+        fixed_path = self._fix_index_drift(path)
+
+        self._update_index_drift_map(fixed_path)
+
+        return self._get_remove_operation(fixed_path)
+
+
+class JsonPatchUpdater:
+    """Reponsible for applying any custom changes to the JSONPatch generated by the jsonpatch library."""
+
+    def __init__(
+        self,
+        restrict_update_fields: List,
+        replace_with_none_op_fixer: ReplaceWithNoneOpFixer,
+    ):
+        self.restrict_update_fields = restrict_update_fields
+        self.replace_with_none_op_fixer = replace_with_none_op_fixer
+
+    @classmethod
+    def from_restrict_update_fields(
+        cls, restrict_update_fields: List
+    ) -> "JsonPatchUpdater":
+        """Instantiates a JsonPatchUpdater based on the restric_update_fields"""
+        return cls(
+            restrict_update_fields=restrict_update_fields,
+            replace_with_none_op_fixer=ReplaceWithNoneOpFixer.default(),
+        )
+
+    def _determine_restricted_operation(self, patch_ops: Dict) -> 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
+        return True
+
+    def _is_replace_with_none_operation(self, patch_ops: dict) -> bool:
+        """Check if the Operation is a Replace operation to a None value."""
+        return (patch_ops.get("op") == PatchOperation.REPLACE.value) and (
+            patch_ops.get("value") is None
+        )
+
+    def _get_remove_operation_for_replace_with_none(self, path: str) -> Dict:
+        """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:
+        """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.
+        """
+        patch_ops_list = []
+        remove_ops_list = []
+
+        for patch_ops in patch.patch or []:
+            if self._determine_restricted_operation(patch_ops=patch_ops):
+                patch_ops_list.append(patch_ops)
+
+                if self._is_replace_with_none_operation(patch_ops):
+                    remove_ops_list.append(
+                        self._get_remove_operation_for_replace_with_none(
+                            patch_ops["path"]
+                        )
+                    )
+
+        patch_ops_list.extend(remove_ops_list)
+
+        return patch_ops_list
+
+
 def build_patch(
     source: T,
     destination: T,
     allowed_fields: Optional[Dict] = None,
     restrict_update_fields: Optional[List] = None,
     array_entity_fields: Optional[List] = None,
+    remove_change_description: bool = True,
 ) -> Optional[jsonpatch.JsonPatch]:
     """
     Given an Entity type and Source entity and Destination entity,
@@ -164,8 +328,9 @@ def build_patch(
     """
 
     # remove change descriptions from entities
-    source = _remove_change_description(source)
-    destination = _remove_change_description(destination)
+    if remove_change_description:
+        source = _remove_change_description(source)
+        destination = _remove_change_description(destination)
 
     if array_entity_fields:
         _sort_array_entity_fields(
@@ -200,22 +365,14 @@ def build_patch(
     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
+    # 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:
-        patch_ops_list = []
-        for patch_ops in patch.patch or []:
-            if _determine_restricted_operation(
-                patch_ops=patch_ops, restrict_update_fields=restrict_update_fields
-            ):
-                if (
-                    patch_ops.get("op") == PatchOperation.REPLACE.value
-                    and patch_ops.get("value") is None
-                ):
-                    patch_ops["op"] = PatchOperation.REMOVE.value
-                    del patch_ops["value"]
-                patch_ops_list.append(patch_ops)
-        patch.patch = patch_ops_list
+        updated_operations = JsonPatchUpdater.from_restrict_update_fields(
+            restrict_update_fields
+        ).update(patch)
+        patch.patch = updated_operations
+
     return patch
 
 
@@ -255,21 +412,6 @@ def _sort_array_entity_fields(
             setattr(destination, field, final_attributes)
 
 
-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.

metadata/ingestion/models/topology.py

@@ -11,8 +11,10 @@
 """
 Defines the topology for ingesting sources
 """
+import queue
+import threading
 from functools import singledispatchmethod
-from typing import Any, Generic, List, Optional, Type, TypeVar
+from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
 
 from pydantic import BaseModel, Extra, Field, create_model
 
@@ -115,6 +117,10 @@ class TopologyNode(BaseModel):
     post_process: Optional[List[str]] = Field(
         None, description="Method to be run after the node has been fully processed"
     )
+    threads: bool = Field(
+        False,
+        description="Flag that defines if a node is open to MultiThreading processing.",
+    )
 
 
 class ServiceTopology(BaseModel):
@@ -245,6 +251,80 @@ class TopologyContext(BaseModel):
         )
 
 
+class TopologyContextManager:
+    """Manages the Context held for different threads."""
+
+    def __init__(self, topology: ServiceTopology):
+        # Due to our code strucutre, the first time the ContextManager is called will be within the MainThread.
+        # We can leverage this to guarantee we keep track of the MainThread ID.
+        self.main_thread = self.get_current_thread_id()
+        self.contexts: Dict[int, TopologyContext] = {
+            self.main_thread: TopologyContext.create(topology)
+        }
+
+        # Starts with the Multithreading disabled
+        self.threads = 0
+
+    def set_threads(self, threads: Optional[int]):
+        self.threads = threads or 0
+
+    def get_current_thread_id(self):
+        return threading.get_ident()
+
+    def get_global(self) -> TopologyContext:
+        return self.contexts[self.main_thread]
+
+    def get(self, thread_id: Optional[int] = None) -> TopologyContext:
+        """Returns the TopologyContext of a given thread."""
+        if thread_id:
+            return self.contexts[thread_id]
+
+        thread_id = self.get_current_thread_id()
+
+        return self.contexts[thread_id]
+
+    def pop(self, thread_id: Optional[int] = None):
+        """Cleans the TopologyContext of a given thread in order to lower the Memory Profile."""
+        if not thread_id:
+            self.contexts.pop(self.get_current_thread_id())
+        else:
+            self.contexts.pop(thread_id)
+
+    def copy_from(self, parent_thread_id: int):
+        """Copies the TopologyContext from a given Thread to the new thread TopologyContext."""
+        thread_id = self.get_current_thread_id()
+
+        # 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)
+        )
+
+
+class Queue:
+    """Small Queue wrapper"""
+
+    def __init__(self):
+        self._queue = queue.Queue()
+
+    def has_tasks(self) -> bool:
+        """Checks that the Queue is not Empty."""
+        return not self._queue.empty()
+
+    def process(self) -> Any:
+        """Yields all the items currently on the Queue."""
+        while True:
+            try:
+                item = self._queue.get_nowait()
+                yield item
+                self._queue.task_done()
+            except queue.Empty:
+                break
+
+    def put(self, item: Any):
+        """Puts new item in the Queue."""
+        self._queue.put(item)
+
+
 def get_topology_nodes(topology: ServiceTopology) -> List[TopologyNode]:
     """
     Fetch all nodes from a ServiceTopology

metadata/ingestion/ometa/client.py

@@ -201,6 +201,9 @@ class REST:
                 )
                 time.sleep(retry_wait)
                 retry -= 1
+                if retry == 0:
+                    logger.error(f"No more retries left for {url}")
+                    traceback.format_exc()
         return None
 
     def _one_request(self, method: str, url: URL, opts: dict, retry: int):

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -16,16 +16,13 @@ To be used by OpenMetadata class
 import functools
 import json
 import traceback
-from typing import Generic, List, Optional, Set, Type, TypeVar
+from typing import Generic, Iterable, List, Optional, Set, Type, TypeVar
 
 from pydantic import BaseModel
 from requests.utils import quote
 
-from metadata.generated.schema.api.createEventPublisherJob import (
-    CreateEventPublisherJob,
-)
+from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.query import Query
-from metadata.generated.schema.system.eventPublisherJob import EventPublisherResult
 from metadata.ingestion.ometa.client import REST, APIError
 from metadata.utils.elasticsearch import ES_INDEX_MAP
 from metadata.utils.logger import ometa_logger
@@ -45,7 +42,7 @@ class ESMixin(Generic[T]):
     client: REST
 
     fqdn_search = (
-        "/search/fieldQuery?fieldName=fullyQualifiedName&fieldValue={fqn}&from={from_}"
+        "/search/fieldQuery?fieldName={field_name}&fieldValue={field_value}&from={from_}"
         "&size={size}&index={index}"
     )
 
@@ -79,7 +76,7 @@ class ESMixin(Generic[T]):
 
         return None
 
-    def _get_entity_from_es(
+    def get_entity_from_es(
         self, entity: Type[T], query_string: str, fields: Optional[list] = None
     ) -> Optional[T]:
         """Fetch an entity instance from ES"""
@@ -96,6 +93,23 @@ class ESMixin(Generic[T]):
 
         return None
 
+    def yield_entities_from_es(
+        self, entity: Type[T], query_string: str, fields: Optional[list] = None
+    ) -> Iterable[T]:
+        """Fetch an entity instance from ES"""
+
+        try:
+            entity_list = self._search_es_entity(
+                entity_type=entity, query_string=query_string, fields=fields
+            )
+            for instance in entity_list or []:
+                yield instance
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not get {entity.__name__} info from ES due to {err}")
+
+        return None
+
     def es_search_from_fqn(
         self,
         entity_type: Type[T],
@@ -105,17 +119,83 @@ class ESMixin(Generic[T]):
         fields: Optional[str] = None,
     ) -> Optional[List[T]]:
         """
-        Given a service_name and some filters, search for entities using ES
+        Given a service name and filters, search for entities using Elasticsearch.
+
+        Args:
+            entity_type (Type[T]): The type of entity to look for.
+            fqn_search_string (str): The string used to search by fully qualified name (FQN).
+                Example: "service.*.schema.table".
+            from_count (int): The starting index of the search results.
+            size (int): The maximum number of records to return.
+            fields (Optional[str]): Comma-separated list of fields to be returned.
+
+        Returns:
+            Optional[List[T]]: A list of entities that match the search criteria, or None if no entities are found.
+        """
+        return self._es_search_entity(
+            entity_type=entity_type,
+            field_value=fqn_search_string,
+            field_name="fullyQualifiedName",
+            from_count=from_count,
+            size=size,
+            fields=fields,
+        )
 
-        :param entity_type: Entity to look for
-        :param fqn_search_string: string used to search by FQN. E.g., service.*.schema.table
-        :param from_count: Records to expect
-        :param size: Number of records
-        :param fields: Comma separated list of fields to be returned
-        :return: List of entities
+    def es_search_container_by_path(
+        self,
+        full_path: str,
+        from_count: int = 0,
+        size: int = 10,
+        fields: Optional[str] = None,
+    ) -> Optional[List[Container]]:
+        """
+        Given a service name and filters, search for containers using Elasticsearch.
+
+        Args:
+            entity_type (Type[T]): The type of entity to look for.
+            full_path (str): The string used to search by full path.
+            from_count (int): The starting index of the search results.
+            size (int): The maximum number of records to return.
+            fields (Optional[str]): Comma-separated list of fields to be returned.
+
+        Returns:
+            Optional[List[Container]]: A list of containers that match the search criteria, or None if no entities are found.
+        """
+        return self._es_search_entity(
+            entity_type=Container,
+            field_value=full_path,
+            field_name="fullPath",
+            from_count=from_count,
+            size=size,
+            fields=fields,
+        )
+
+    def _es_search_entity(
+        self,
+        entity_type: Type[T],
+        field_value: str,
+        field_name: str,
+        from_count: int = 0,
+        size: int = 10,
+        fields: Optional[str] = None,
+    ) -> Optional[List[T]]:
+        """
+        Search for entities using Elasticsearch.
+
+        Args:
+            entity_type (Type[T]): The type of entity to look for.
+            field_value (str): The value to search for in the specified field.
+            field_name (str): The name of the field to search in.
+            from_count (int, optional): The starting index of the search results. Defaults to 0.
+            size (int, optional): The maximum number of search results to return. Defaults to 10.
+            fields (Optional[str], optional): Comma-separated list of fields to be returned. Defaults to None.
+
+        Returns:
+            Optional[List[T]]: A list of entities that match the search criteria, or None if no entities are found.
         """
         query_string = self.fqdn_search.format(
-            fqn=fqn_search_string,
+            field_name=field_name,
+            field_value=field_value,
             from_=from_count,
             size=size,
             index=ES_INDEX_MAP[entity_type.__name__],  # Fail if not exists
@@ -138,33 +218,6 @@ class ESMixin(Generic[T]):
             )
         return None
 
-    def reindex_es(
-        self,
-        config: CreateEventPublisherJob,
-    ) -> Optional[EventPublisherResult]:
-        """
-        Method to trigger elasticsearch reindex
-        """
-        try:
-            resp = self.client.post(path="/search/reindex", data=config.json())
-            return EventPublisherResult(**resp)
-        except APIError as err:
-            logger.debug(traceback.format_exc())
-            logger.debug(f"Failed to trigger es reindex job due to {err}")
-            return None
-
-    def get_reindex_job_status(self, job_id: str) -> Optional[EventPublisherResult]:
-        """
-        Method to fetch the elasticsearch reindex job status
-        """
-        try:
-            resp = self.client.get(path=f"/search/reindex/{job_id}")
-            return EventPublisherResult(**resp)
-        except APIError as err:
-            logger.debug(traceback.format_exc())
-            logger.debug(f"Failed to fetch reindex job status due to {err}")
-            return None
-
     @staticmethod
     def get_query_with_lineage_filter(service_name: str) -> str:
         query_lineage_filter = {