openmetadata-ingestion 1.6.6.6__py3-none-any.whl → 1.6.6.7.dev0__py3-none-any.whl

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2025-03-29T04:07:20+00:00
+#   timestamp: 2025-04-09T12:07:25+00:00
 
 from __future__ import annotations
 

metadata/ingestion/lineage/masker.py

@@ -13,12 +13,14 @@ Query masking utilities
 """
 
 import traceback
+from typing import Optional
 
 from collate_sqllineage.runner import SQLPARSE_DIALECT, LineageRunner
 from sqlparse.sql import Comparison
 from sqlparse.tokens import Literal, Number, String
 
 from metadata.ingestion.lineage.models import Dialect
+from metadata.utils.execution_time_tracker import calculate_execution_time
 
 MASK_TOKEN = "?"
 
@@ -109,11 +111,17 @@ def mask_literals_with_sqlfluff(query: str, parser: LineageRunner) -> str:
     return query
 
 
+@calculate_execution_time(context="MaskQuery")
 def mask_query(
-    query: str, dialect: str = Dialect.ANSI.value, parser: LineageRunner = None
-) -> str:
+    query: str,
+    dialect: str = Dialect.ANSI.value,
+    parser: Optional[LineageRunner] = None,
+    parser_required: bool = False,
+) -> Optional[str]:
     logger = get_logger()
     try:
+        if parser_required and not parser:
+            return None
         if not parser:
             try:
                 parser = LineageRunner(query, dialect=dialect)

metadata/ingestion/lineage/parser.py

@@ -28,6 +28,7 @@ from sqlparse.sql import Comparison, Identifier, Parenthesis, Statement
 from metadata.generated.schema.type.tableUsageCount import TableColumn, TableColumnJoin
 from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.lineage.models import Dialect
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.helpers import (
     find_in_iter,
     get_formatted_entity_name,
@@ -77,7 +78,10 @@ class LineageParser:
             self._clean_query, dialect=dialect, timeout_seconds=timeout_seconds
         )
         if self.masked_query is None:
-            self.masked_query = mask_query(self._clean_query, parser=self.parser)
+            self.masked_query = (
+                mask_query(self._clean_query, parser=self.parser, parser_required=True)
+                or self._clean_query
+            )
 
     @cached_property
     def involved_tables(self) -> Optional[List[Table]]:
@@ -409,6 +413,7 @@ class LineageParser:
 
         return clean_query.strip()
 
+    @calculate_execution_time(context="EvaluateBestParser")
     def _evaluate_best_parser(
         self, query: str, dialect: Dialect, timeout_seconds: int
     ) -> Optional[LineageRunner]:
@@ -447,8 +452,14 @@ class LineageParser:
         if lr_sqlfluff:
             return lr_sqlfluff
 
-        lr_sqlparser = LineageRunner(query)
+        @timeout(seconds=timeout_seconds)
+        def get_sqlparser_lineage_runner(qry: str) -> LineageRunner:
+            lr_sqlparser = LineageRunner(qry)
+            lr_sqlparser.get_column_lineage()
+            return lr_sqlparser
+
         try:
+            lr_sqlparser = get_sqlparser_lineage_runner(query)
             _ = len(lr_sqlparser.get_column_lineage()) + len(
                 set(lr_sqlparser.source_tables).union(
                     set(lr_sqlparser.target_tables).union(
@@ -456,12 +467,18 @@ class LineageParser:
                     )
                 )
             )
+        except TimeoutError:
+            self.query_parsing_success = False
+            self.query_parsing_failure_reason = (
+                f"Lineage with SqlParser failed for the [{dialect.value}]. "
+                f"Parser has been running for more than {timeout_seconds} seconds."
+            )
+            return None
         except Exception:
             # if both runner have failed we return the usual one
             logger.debug(f"Failed to parse query with sqlparse & sqlfluff: {query}")
             return lr_sqlfluff if lr_sqlfluff else lr_sqlparser
 
-        self.masked_query = mask_query(self._clean_query, parser=lr_sqlparser)
         logger.debug(
             f"Using sqlparse for lineage parsing for query: {self.masked_query or self.query}"
         )

metadata/ingestion/lineage/sql_lineage.py

@@ -11,12 +11,15 @@
 """
 Helper functions to handle SQL lineage operations
 """
+import functools
 import itertools
 import traceback
 from collections import defaultdict
+from copy import deepcopy
 from typing import Any, Iterable, List, Optional, Tuple, Union
 
 import networkx as nx
+from collate_sqllineage.core.holders import SQLLineageHolder
 from collate_sqllineage.core.models import Column, DataFunction
 from collate_sqllineage.core.models import Table as LineageTable
 from networkx import DiGraph
@@ -48,6 +51,10 @@ from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT, LineagePa
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils import fqn
 from metadata.utils.elasticsearch import get_entity_from_es_result
+from metadata.utils.execution_time_tracker import (
+    calculate_execution_time,
+    calculate_execution_time_generator,
+)
 from metadata.utils.fqn import build_es_fqn_search_string
 from metadata.utils.logger import utils_logger
 from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
@@ -73,6 +80,7 @@ def get_column_fqn(table_entity: Table, column: str) -> Optional[str]:
 search_cache = LRUCache(LRU_CACHE_SIZE)
 
 
+@calculate_execution_time(context="SearchTableEntities")
 def search_table_entities(
     metadata: OpenMetadata,
     service_name: str,
@@ -236,6 +244,58 @@ def handle_udf_column_lineage(
         logger.error(f"Error handling UDF column lineage: {exc}")
 
 
+@functools.lru_cache(maxsize=1000)
+def _get_udf_parser(
+    code: str, dialect: Dialect, timeout_seconds: int
+) -> Optional[LineageParser]:
+    if code:
+        return LineageParser(
+            f"create table dummy_table_name as {code}",
+            dialect=dialect,
+            timeout_seconds=timeout_seconds,
+        )
+    return None
+
+
+def _replace_target_table(
+    parser: LineageParser, expected_table_name: str
+) -> LineageParser:
+    try:
+        # Create a new target table instead of modifying the existing one
+        new_table = Table(expected_table_name.replace(DEFAULT_SCHEMA_NAME, ""))
+
+        # Create a new statement holder with the updated target table
+        stmt_holder = parser.parser._stmt_holders[0]
+        old_write = list(stmt_holder.write)[0]  # Get the original target table
+
+        # Remove old target table and add new one
+        stmt_holder.graph.remove_node(old_write)
+        stmt_holder.add_write(new_table)
+
+        # Rebuild column lineage
+        for col_lineage in parser.parser.get_column_lineage():
+            if col_lineage[-1].parent == old_write:
+                # Create new column with same name but parent is new table
+                tgt_col = col_lineage[-1]
+                new_tgt_col = Column(tgt_col.raw_name)
+                new_tgt_col.parent = new_table
+
+                # Add the column lineage from source to new target
+                stmt_holder.add_column_lineage(col_lineage[-2], new_tgt_col)
+                try:
+                    # remove the old edge
+                    stmt_holder.graph.remove_edge(col_lineage[-2], tgt_col)
+                except:
+                    # if the edge is not present, pass
+                    pass
+
+        # Rebuild the SQL holder
+        parser.parser._sql_holder = SQLLineageHolder.of(*parser.parser._stmt_holders)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.debug(f"Error replacing target table: {exc}")
+
+
 # pylint: disable=too-many-arguments
 def __process_udf_es_results(
     metadata: OpenMetadata,
@@ -255,27 +315,32 @@ def __process_udf_es_results(
             and entity.storedProcedureCode
             and entity.storedProcedureCode.language == Language.SQL
         ):
-            expected_table_name = str(source_table).replace(
-                f"{DEFAULT_SCHEMA_NAME}.", ""
-            )
-            lineage_parser = LineageParser(
-                f"create table {str(expected_table_name)} as {entity.storedProcedureCode.code}",
-                dialect=dialect,
-                timeout_seconds=timeout_seconds,
+
+            lineage_parser = _get_udf_parser(
+                entity.storedProcedureCode.code, dialect, timeout_seconds
             )
-            handle_udf_column_lineage(column_lineage, lineage_parser.column_lineage)
-            for source in lineage_parser.source_tables or []:
-                yield from get_source_table_names(
-                    metadata,
-                    dialect,
-                    source,
-                    database_name,
-                    schema_name,
-                    service_name,
-                    timeout_seconds,
-                    column_lineage,
-                    procedure or entity,
+            if lineage_parser and lineage_parser.parser:
+                expected_table_name = str(source_table).replace(
+                    f"{DEFAULT_SCHEMA_NAME}.", ""
+                )
+                lineage_parser_copy = deepcopy(lineage_parser)
+                _replace_target_table(lineage_parser_copy, expected_table_name)
+
+                handle_udf_column_lineage(
+                    column_lineage, lineage_parser_copy.column_lineage
                 )
+                for source in lineage_parser_copy.source_tables or []:
+                    yield from get_source_table_names(
+                        metadata,
+                        dialect,
+                        source,
+                        database_name,
+                        schema_name,
+                        service_name,
+                        timeout_seconds,
+                        column_lineage,
+                        procedure or entity,
+                    )
 
 
 def __process_udf_table_names(
@@ -317,6 +382,7 @@ def __process_udf_table_names(
         )
 
 
+@calculate_execution_time_generator(context="GetSourceTableNames")
 def get_source_table_names(
     metadata: OpenMetadata,
     dialect: Dialect,
@@ -617,6 +683,7 @@ def populate_column_lineage_map(raw_column_lineage):
 
 
 # pylint: disable=too-many-locals
+@calculate_execution_time_generator(context="GetLineageByQuery")
 def get_lineage_by_query(
     metadata: OpenMetadata,
     service_name: str,
@@ -627,6 +694,7 @@ def get_lineage_by_query(
     timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
     lineage_source: LineageSource = LineageSource.QueryLineage,
     graph: DiGraph = None,
+    lineage_parser: Optional[LineageParser] = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
     This method parses the query to get source, target and intermediate table names to create lineage,
@@ -636,7 +704,10 @@ def get_lineage_by_query(
     query_parsing_failures = QueryParsingFailures()
 
     try:
-        lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
+        if not lineage_parser:
+            lineage_parser = LineageParser(
+                query, dialect, timeout_seconds=timeout_seconds
+            )
         masked_query = lineage_parser.masked_query
         logger.debug(f"Running lineage with query: {masked_query or query}")
 
@@ -723,6 +794,7 @@ def get_lineage_by_query(
         )
 
 
+@calculate_execution_time_generator(context="GetLineageViaTableEntity")
 def get_lineage_via_table_entity(
     metadata: OpenMetadata,
     table_entity: Table,
@@ -734,13 +806,17 @@ def get_lineage_via_table_entity(
     timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
     lineage_source: LineageSource = LineageSource.QueryLineage,
     graph: DiGraph = None,
+    lineage_parser: Optional[LineageParser] = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """Get lineage from table entity"""
     column_lineage = {}
     query_parsing_failures = QueryParsingFailures()
 
     try:
-        lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
+        if not lineage_parser:
+            lineage_parser = LineageParser(
+                query, dialect, timeout_seconds=timeout_seconds
+            )
         masked_query = lineage_parser.masked_query
         logger.debug(
             f"Getting lineage via table entity using query: {masked_query or query}"

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -30,6 +30,7 @@ from metadata.ingestion.ometa.client import REST, APIError
 from metadata.ingestion.ometa.utils import quote
 from metadata.ingestion.source.models import TableView
 from metadata.utils.elasticsearch import ES_INDEX_MAP
+from metadata.utils.execution_time_tracker import calculate_execution_time_generator
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -388,6 +389,7 @@ class ESMixin(Generic[T]):
                     f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
                 )
 
+    @calculate_execution_time_generator(context="ES.FetchViewDefinition")
     def yield_es_view_def(
         self,
         service_name: str,

metadata/ingestion/ometa/mixins/lineage_mixin.py

@@ -14,6 +14,7 @@ Mixin class containing Lineage specific methods
 To be used by OpenMetadata class
 """
 import functools
+import json
 import traceback
 from copy import deepcopy
 from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
@@ -427,6 +428,7 @@ class OMetaLineageMixin(Generic[T]):
             logger.debug(traceback.format_exc())
         self.is_edit_table_allowed = False
 
+    @functools.lru_cache(maxsize=LRU_CACHE_SIZE)
     def patch_lineage_processed_flag(
         self,
         entity: Type[T],
@@ -440,15 +442,16 @@ class OMetaLineageMixin(Generic[T]):
             return
 
         try:
-            original_entity = self.get_by_name(entity=entity, fqn=fqn)
-            if not original_entity:
-                return
-
-            updated_entity = original_entity.model_copy(deep=True)
-            updated_entity.processedLineage = True
-
-            self.patch(
-                entity=entity, source=original_entity, destination=updated_entity
+            patch = [
+                {
+                    "op": "add",
+                    "path": "/processedLineage",
+                    "value": True,
+                }
+            ]
+            self.client.patch(
+                path=f"{self.get_suffix(entity)}/name/{fqn}",
+                data=json.dumps(patch),
             )
         except Exception as exc:
             logger.debug(f"Error while patching lineage processed flag: {exc}")

metadata/ingestion/source/database/lineage_source.py

@@ -13,10 +13,10 @@ Lineage Source Module
 """
 import csv
 import os
+import threading
 import time
 import traceback
 from abc import ABC
-from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from typing import Any, Callable, Iterable, Iterator, List, Union
 
@@ -42,9 +42,9 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
-CHUNK_SIZE = 200
+CHUNK_SIZE = 100
 
-THREAD_TIMEOUT = 600
+THREAD_TIMEOUT = 3 * 60 * 10  # 30 minutes in seconds
 
 
 class LineageSource(QueryParserSource, ABC):
@@ -107,61 +107,95 @@ class LineageSource(QueryParserSource, ABC):
             )
             yield from self.yield_table_query()
 
+    @staticmethod
     def generate_lineage_in_thread(
-        self,
         producer_fn: Callable[[], Iterable[Any]],
-        processor_fn: Callable[[Any], Iterable[Any]],
+        processor_fn: Callable[[Any, Queue], None],
         chunk_size: int = CHUNK_SIZE,
+        thread_timeout: int = THREAD_TIMEOUT,
+        max_threads: int = 10,  # Default maximum number of concurrent threads
     ):
         """
-        Optimized multithreaded lineage generation with improved error handling and performance.
+        Process data in separate daemon threads with timeout control.
 
         Args:
-            producer_fn: Function that yields input items
-            processor_fn: Function to process each input item
-            chunk_size: Optional batching to reduce thread creation overhead
+            producer_fn: Function that yields data chunks
+            processor_fn: Function that processes data and adds results to the queue
+            chunk_size: Size of chunks to process
+            thread_timeout: Maximum time in seconds to wait for a processor thread
+            max_threads: Maximum number of concurrent threads to run
         """
+        queue = Queue()
+        active_threads = []
 
-        def chunk_generator():
-            temp_chunk = []
-            for chunk in producer_fn():
-                temp_chunk.append(chunk)
-                if len(temp_chunk) >= chunk_size:
-                    yield temp_chunk
-                    temp_chunk = []
+        def process_chunk(chunk):
+            """Process a chunk of data in a thread."""
+            try:
+                processor_fn(chunk, queue)
+            except Exception as e:
+                logger.error(f"Error processing chunk: {e}")
+                logger.debug(traceback.format_exc())
 
-            if temp_chunk:
-                yield temp_chunk
+        # Create an iterator for the chunks but don't consume it all at once
+        chunk_iterator = iter(chunk_generator(producer_fn, chunk_size))
 
-        thread_pool = ThreadPoolExecutor(max_workers=self.source_config.threads)
-        queue = Queue()
+        # Process results from the queue and check for timed-out threads
+        chunk_processed = False  # Flag to track if all chunks have been processed
+        ignored_threads = 0
 
-        futures = [
-            thread_pool.submit(
-                processor_fn,
-                chunk,
-                queue,
-            )
-            for chunk in chunk_generator()
-        ]
         while True:
+            # Start new threads until we reach the max_threads limit
+            while (
+                len(active_threads) + ignored_threads
+            ) < max_threads and not chunk_processed:
+                try:
+                    # Only fetch a new chunk when we're ready to create a thread
+                    chunk = next(chunk_iterator)
+                    thread = threading.Thread(target=process_chunk, args=(chunk,))
+                    thread.start_time = time.time()  # Track when the thread started
+                    thread.daemon = True
+                    active_threads.append(thread)
+                    thread.start()
+                except StopIteration:
+                    # No more chunks to process
+                    chunk_processed = True
+                    break
+
+            if ignored_threads == max_threads:
+                logger.warning(f"Max threads reached, skipping remaining threads")
+                break
+
+            # Process any available results
             if queue.has_tasks():
                 yield from queue.process()
 
-            else:
-                if not futures:
-                    break
+            # Check for completed or timed-out threads
+            still_active = []
+            for thread in active_threads:
+                if thread.is_alive():
+                    # Check if the thread has timed out
+                    if time.time() - thread.start_time > thread_timeout:
+                        logger.warning(
+                            f"Thread {thread.name} timed out after {thread_timeout}s"
+                        )
+                        ignored_threads += 1
+                    else:
+                        still_active.append(thread)
+                # If thread is not alive, it has completed normally
+
+            active_threads = still_active
 
-                for i, future in enumerate(futures):
-                    if future.done():
-                        try:
-                            future.result(timeout=THREAD_TIMEOUT)
-                        except Exception as e:
-                            logger.debug(f"Error in future: {e}")
-                            logger.debug(traceback.format_exc())
-                        futures.pop(i)
+            # Exit conditions: no more active threads and no more chunks to process
+            if not active_threads and chunk_processed:
+                break
 
-            time.sleep(0.01)
+            # Small pause to prevent CPU spinning
+            if active_threads:
+                time.sleep(0.1)
+
+        # Final check for any remaining results
+        while queue.has_tasks():
+            yield from queue.process()
 
     def yield_table_query(self) -> Iterator[TableQuery]:
         """
@@ -257,7 +291,9 @@ class LineageSource(QueryParserSource, ABC):
         producer_fn = self.get_table_query
         processor_fn = self.query_lineage_generator
         yield from self.generate_lineage_in_thread(
-            producer_fn, processor_fn, CHUNK_SIZE
+            producer_fn,
+            processor_fn,
+            max_threads=self.source_config.threads,
         )
 
     def view_lineage_generator(
@@ -306,7 +342,9 @@ class LineageSource(QueryParserSource, ABC):
             self.source_config.incrementalLineageProcessing,
         )
         processor_fn = self.view_lineage_generator
-        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)
+        yield from self.generate_lineage_in_thread(
+            producer_fn, processor_fn, max_threads=self.source_config.threads
+        )
 
     def yield_procedure_lineage(
         self,
@@ -339,3 +377,18 @@ class LineageSource(QueryParserSource, ABC):
                 logger.warning(
                     f"Lineage extraction is not supported for {str(self.service_connection.type.value)} connection"
                 )
+
+
+def chunk_generator(producer_fn, chunk_size):
+    """
+    Group items from producer into chunks of specified size.
+    This is a separate function to allow for better lazy evaluation.
+    """
+    temp_chunk = []
+    for item in producer_fn():
+        temp_chunk.append(item)
+        if len(temp_chunk) >= chunk_size:
+            yield temp_chunk
+            temp_chunk = []
+    if temp_chunk:
+        yield temp_chunk

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -304,4 +304,6 @@ class StoredProcedureLineageMixin(ABC):
         logger.info("Processing Lineage for Stored Procedures")
         producer_fn = self.procedure_lineage_generator
         processor_fn = self.procedure_lineage_processor
-        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)
+        yield from self.generate_lineage_in_thread(
+            producer_fn, processor_fn, max_threads=self.source_config.threads
+        )