openmetadata-ingestion 1.3.3.0rc1__py3-none-any.whl → 1.4.0.0rc1__py3-none-any.whl

metadata/ingestion/api/topology_runner.py

@@ -12,8 +12,11 @@
 Mixin to be used by service sources to dynamically
 generate the _run based on their topology.
 """
+import math
+import time
 import traceback
 from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
 from functools import singledispatchmethod
 from typing import Any, Generic, Iterable, List, Type, TypeVar
 
@@ -29,14 +32,16 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.patch_request import PatchRequest
 from metadata.ingestion.models.topology import (
     NodeStage,
+    Queue,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
     get_topology_node,
     get_topology_root,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.utils import model_str
+from metadata.utils.execution_time_tracker import ExecutionTimeTrackerContextMap
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.source_hash import generate_source_hash
 
@@ -59,11 +64,78 @@ class TopologyRunnerMixin(Generic[C]):
     """
 
     topology: ServiceTopology
-    context: TopologyContext
+    context: TopologyContextManager
     metadata: OpenMetadata
 
     # The cache will have the shape {`child_stage.type_`: {`name`: `hash`}}
     cache = defaultdict(dict)
+    queue = Queue()
+
+    def _multithread_process_node(
+        self, node: TopologyNode, threads: int
+    ) -> Iterable[Entity]:
+        """Multithread Processing of a Node"""
+        node_producer = getattr(self, node.producer)
+        child_nodes = self._get_child_nodes(node)
+
+        node_entities = list(node_producer() or [])
+        node_entities_length = len(node_entities)
+
+        if node_entities_length == 0:
+            return
+        else:
+            chunksize = int(math.ceil(node_entities_length / threads))
+            chunks = [
+                node_entities[i : i + chunksize]
+                for i in range(0, node_entities_length, chunksize)
+            ]
+
+            thread_pool = ThreadPoolExecutor(max_workers=threads)
+
+            futures = [
+                thread_pool.submit(
+                    self._multithread_process_entity,
+                    node,
+                    chunk,
+                    child_nodes,
+                    self.context.get_current_thread_id(),
+                )
+                for chunk in chunks
+            ]
+
+            while True:
+                if self.queue.has_tasks():
+                    yield from self.queue.process()
+
+                else:
+                    if not futures:
+                        break
+
+                    for i, future in enumerate(futures):
+                        if future.done():
+                            future.result()
+                            futures.pop(i)
+
+                time.sleep(0.01)
+
+    def _process_node(self, node: TopologyNode) -> Iterable[Entity]:
+        """Processing of a Node in a single thread."""
+        node_producer = getattr(self, node.producer)
+        child_nodes = self._get_child_nodes(node)
+
+        for node_entity in node_producer() or []:
+            for stage in node.stages:
+                yield from self._process_stage(
+                    stage=stage, node_entity=node_entity, child_nodes=child_nodes
+                )
+
+            # Once we are done processing all the stages,
+            for stage in node.stages:
+                if stage.clear_context:
+                    self.context.get().clear_stage(stage=stage)
+
+            # process all children from the node being run
+            yield from self.process_nodes(child_nodes)
 
     def process_nodes(self, nodes: List[TopologyNode]) -> Iterable[Entity]:
         """
@@ -89,31 +161,54 @@ class TopologyRunnerMixin(Generic[C]):
         :param nodes: Topology Nodes to process
         :return: recursively build the execution tree
         """
+
         for node in nodes:
             logger.debug(f"Processing node {node}")
-            node_producer = getattr(self, node.producer)
-            child_nodes = self._get_child_nodes(node)
 
             # Each node producer will give us a list of entities that we need
             # to process. Each of the internal stages will sink result to OM API.
             # E.g., in the DB topology, at the Table TopologyNode, the node_entity
             # will be each `table`
-            for node_entity in node_producer() or []:
-                for stage in node.stages:
-                    yield from self._process_stage(
-                        stage=stage, node_entity=node_entity, child_nodes=child_nodes
-                    )
-
-                # Once we are done processing all the stages,
-                for stage in node.stages:
-                    if stage.clear_context:
-                        self.context.clear_stage(stage=stage)
-
-                # process all children from the node being run
-                yield from self.process_nodes(child_nodes)
+            if node.threads and self.context.threads > 1:
+                yield from self._multithread_process_node(node, self.context.threads)
+            else:
+                yield from self._process_node(node)
 
             yield from self._run_node_post_process(node=node)
 
+    def _multithread_process_entity(
+        self,
+        node: TopologyNode,
+        node_entities: List[Any],
+        child_nodes: List[TopologyNode],
+        parent_thread_id: int,
+    ):
+        """Multithread processing of a Node Entity"""
+        # Generates a new context based on the parent thread.
+        self.context.copy_from(parent_thread_id)
+        ExecutionTimeTrackerContextMap().copy_from_parent(parent_thread_id)
+
+        for node_entity in node_entities:
+            # For each stage, we get all the stage results and one by one yield them by adding them to the Queue.
+            for stage in node.stages:
+                for stage_result in self._process_stage(
+                    stage=stage, node_entity=node_entity, child_nodes=child_nodes
+                ):
+                    self.queue.put(stage_result)
+
+            # After all the stages are done, we clear the context if needed.
+            for stage in node.stages:
+                if stage.clear_context:
+                    self.context.get().clear_stage(stage=stage)
+
+            # If the Entity has child nodes that need processing we proceed to processing them with the same logic as above.
+
+            for child_result in self.process_nodes(child_nodes):
+                self.queue.put(child_result)
+
+        # Finally we pop the context and finish the thread
+        self.context.pop()
+
     def _get_child_nodes(self, node: TopologyNode) -> List[TopologyNode]:
         """Compute children nodes if any"""
         return (
@@ -158,7 +253,7 @@ class TopologyRunnerMixin(Generic[C]):
             for process in node.post_process:
                 try:
                     node_post_process = getattr(self, process)
-                    for entity_request in node_post_process():
+                    for entity_request in node_post_process() or []:
                         yield entity_request
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
@@ -177,9 +272,9 @@ class TopologyRunnerMixin(Generic[C]):
         for child_node in child_nodes or []:
             for child_stage in child_node.stages or []:
                 if child_stage.use_cache:
-                    entity_fqn = self.context.fqn_from_stage(
+                    entity_fqn = self.context.get().fqn_from_stage(
                         stage=stage,
-                        entity_name=self.context.__dict__[stage.context],
+                        entity_name=self.context.get().__dict__[stage.context],
                     )
 
                     self.get_fqn_source_hash_dict(
@@ -251,7 +346,9 @@ class TopologyRunnerMixin(Generic[C]):
         """
         entity = None
         entity_name = model_str(right.name)
-        entity_fqn = self.context.fqn_from_stage(stage=stage, entity_name=entity_name)
+        entity_fqn = self.context.get().fqn_from_stage(
+            stage=stage, entity_name=entity_name
+        )
 
         # If we don't want to write data in OM, we'll return what we fetch from the API.
         # This will be applicable for service entities since we do not want to overwrite the data
@@ -325,7 +422,7 @@ class TopologyRunnerMixin(Generic[C]):
                     "for the service connection."
                 )
 
-        self.context.update_context_name(stage=stage, right=right)
+        self.context.get().update_context_name(stage=stage, right=right)
 
     @yield_and_update_context.register
     def _(
@@ -341,7 +438,7 @@ class TopologyRunnerMixin(Generic[C]):
         lineage has been properly drawn. We'll skip the process for now.
         """
         yield entity_request
-        self.context.update_context_name(stage=stage, right=right.edge.fromEntity)
+        self.context.get().update_context_name(stage=stage, right=right.edge.fromEntity)
 
     @yield_and_update_context.register
     def _(
@@ -359,7 +456,7 @@ class TopologyRunnerMixin(Generic[C]):
         """
         yield entity_request
 
-        self.context.update_context_value(stage=stage, value=right)
+        self.context.get().update_context_value(stage=stage, value=right)
 
     @yield_and_update_context.register
     def _(
@@ -371,7 +468,7 @@ class TopologyRunnerMixin(Generic[C]):
         """Custom Property implementation for the context information"""
         yield entity_request
 
-        self.context.update_context_value(stage=stage, value=right)
+        self.context.get().update_context_value(stage=stage, value=right)
 
     def sink_request(
         self, stage: NodeStage, entity_request: Either[C]

metadata/ingestion/bulksink/metadata_usage.py

@@ -93,7 +93,12 @@ class MetadataUsageBulkSink(BulkSink):
         return "OpenMetadata"
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         config = MetadataUsageSinkConfig.parse_obj(config_dict)
         return cls(config, metadata)
 

metadata/ingestion/lineage/models.py

@@ -12,7 +12,9 @@
 Models related to lineage parsing
 """
 from enum import Enum
-from typing import Dict
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Extra, Field
 
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaType,
@@ -62,6 +64,7 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon
 from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
     SQLiteType,
 )
+from metadata.utils.singleton import Singleton
 
 
 class Dialect(Enum):
@@ -126,3 +129,33 @@ class ConnectionTypeDialectMapper:
         Returns: a dialect
         """
         return MAP_CONNECTION_TYPE_DIALECT.get(connection_type, Dialect.ANSI)
+
+
+class QueryParsingError(BaseModel):
+    """
+    Represents an error that occurs during query parsing.
+
+    Attributes:
+        query (str): The query text of the failed query.
+        error (str): The error message of the failed query.
+    """
+
+    class Config:
+        extra = Extra.forbid
+
+    query: str = Field(..., description="query text of the failed query")
+    error: Optional[str] = Field(..., description="error message of the failed query")
+
+
+class QueryParsingFailures(metaclass=Singleton):
+    """Tracks the Queries that failed to parse."""
+
+    def __init__(self):
+        """Initializes the list of parsing failures."""
+        self._query_list: List[QueryParsingError] = []
+
+    def add(self, parsing_error: QueryParsingError):
+        self._query_list.append(parsing_error)
+
+    def __iter__(self):
+        return iter(self._query_list)

metadata/ingestion/lineage/parser.py

@@ -67,6 +67,8 @@ class LineageParser:
         timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
     ):
         self.query = query
+        self.query_parsing_success = True
+        self.query_parsing_failure_reason = None
         self._clean_query = self.clean_raw_query(query)
         self.parser = self._evaluate_best_parser(
             self._clean_query, dialect=dialect, timeout_seconds=timeout_seconds
@@ -400,9 +402,8 @@ class LineageParser:
 
         return clean_query.strip()
 
-    @staticmethod
     def _evaluate_best_parser(
-        query: str, dialect: Dialect, timeout_seconds: int
+        self, query: str, dialect: Dialect, timeout_seconds: int
     ) -> Optional[LineageRunner]:
         if query is None:
             return None
@@ -424,15 +425,19 @@ class LineageParser:
                 )
             )
         except TimeoutError:
-            logger.debug(
-                f"Lineage with SqlFluff failed for the [{dialect.value}] query: [{query}]: "
+            self.query_parsing_success = False
+            self.query_parsing_failure_reason = (
+                f"Lineage with SqlFluff failed for the [{dialect.value}]. "
                 f"Parser has been running for more than {timeout_seconds} seconds."
             )
+            logger.debug(f"{self.query_parsing_failure_reason}] query: [{query}]")
             lr_sqlfluff = None
         except Exception:
-            logger.debug(
-                f"Lineage with SqlFluff failed for the [{dialect.value}] query: [{query}]"
+            self.query_parsing_success = False
+            self.query_parsing_failure_reason = (
+                f"Lineage with SqlFluff failed for the [{dialect.value}]"
             )
+            logger.debug(f"{self.query_parsing_failure_reason} query: [{query}]")
             lr_sqlfluff = None
 
         lr_sqlparser = LineageRunner(query)
@@ -451,10 +456,12 @@ class LineageParser:
         if lr_sqlfluff:
             # if sqlparser retrieve more lineage info that sqlfluff
             if sqlparser_count > sqlfluff_count:
-                logger.debug(
+                self.query_parsing_success = False
+                self.query_parsing_failure_reason = (
                     "Lineage computed with SqlFluff did not perform as expected "
-                    f"for the [{dialect.value}] query: [{query}]"
+                    f"for the [{dialect.value}]"
                 )
+                logger.debug(f"{self.query_parsing_failure_reason} query: [{query}]")
                 return lr_sqlparser
             return lr_sqlfluff
         return lr_sqlparser

metadata/ingestion/lineage/sql_lineage.py

@@ -27,16 +27,19 @@ from metadata.generated.schema.type.entityLineage import (
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
-from metadata.ingestion.lineage.models import Dialect
+from metadata.ingestion.lineage.models import (
+    Dialect,
+    QueryParsingError,
+    QueryParsingFailures,
+)
 from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT, LineageParser
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils import fqn
 from metadata.utils.fqn import build_es_fqn_search_string
 from metadata.utils.logger import utils_logger
-from metadata.utils.lru_cache import LRUCache
+from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
 
 logger = utils_logger()
-LRU_CACHE_SIZE = 4096
 DEFAULT_SCHEMA_NAME = "<default>"
 
 
@@ -365,6 +368,7 @@ def populate_column_lineage_map(raw_column_lineage):
     return lineage_map
 
 
+# pylint: disable=too-many-locals
 def get_lineage_by_query(
     metadata: OpenMetadata,
     service_name: str,
@@ -380,6 +384,7 @@ def get_lineage_by_query(
     and returns True if target table is found to create lineage otherwise returns False.
     """
     column_lineage = {}
+    query_parsing_failures = QueryParsingFailures()
 
     try:
         logger.debug(f"Running lineage with query: {query}")
@@ -427,6 +432,12 @@ def get_lineage_by_query(
                         column_lineage_map=column_lineage,
                         lineage_source=lineage_source,
                     )
+        if not lineage_parser.query_parsing_success:
+            query_parsing_failures.add(
+                QueryParsingError(
+                    query=query, error=lineage_parser.query_parsing_failure_reason
+                )
+            )
     except Exception as exc:
         yield Either(
             left=StackTraceError(
@@ -450,6 +461,7 @@ def get_lineage_via_table_entity(
 ) -> Iterable[Either[AddLineageRequest]]:
     """Get lineage from table entity"""
     column_lineage = {}
+    query_parsing_failures = QueryParsingFailures()
 
     try:
         logger.debug(f"Getting lineage via table entity using query: {query}")
@@ -468,6 +480,12 @@ def get_lineage_via_table_entity(
                 column_lineage_map=column_lineage,
                 lineage_source=lineage_source,
             ) or []
+        if not lineage_parser.query_parsing_success:
+            query_parsing_failures.add(
+                QueryParsingError(
+                    query=query, error=lineage_parser.query_parsing_failure_reason
+                )
+            )
     except Exception as exc:  # pylint: disable=broad-except
         Either(
             left=StackTraceError(

metadata/ingestion/models/custom_properties.py

@@ -37,6 +37,7 @@ class CustomPropertyDataTypes(Enum):
     TIME = "time"
     TIMEINTERVAL = "timeInterval"
     TIMESTAMP = "timestamp"
+    ENUM = "enum"
 
 
 class OMetaCustomProperties(BaseModel):

metadata/ingestion/models/life_cycle.py

@@ -11,6 +11,8 @@
 """
 Custom models for life cycle
 """
+from typing import Type
+
 from pydantic import BaseModel
 
 from metadata.generated.schema.type.lifeCycle import LifeCycle
@@ -18,5 +20,6 @@ from metadata.ingestion.api.models import Entity
 
 
 class OMetaLifeCycleData(BaseModel):
-    entity: Entity
+    entity: Type[Entity]
+    entity_fqn: str
     life_cycle: LifeCycle