openmetadata-ingestion 1.6.9.0__py3-none-any.whl → 1.6.11.0__py3-none-any.whl

metadata/ingestion/source/database/lineage_processors.py

@@ -0,0 +1,336 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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.
+"""
+Mixin class with common Stored Procedures logic aimed at lineage.
+"""
+import re
+import time
+import traceback
+from copy import deepcopy
+from datetime import datetime
+from multiprocessing import Queue
+from typing import Iterable, List, Optional, Union
+
+import networkx as nx
+from pydantic import BaseModel, ConfigDict, Field
+
+from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.type.basic import SqlQuery, Timestamp
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.tableQuery import TableQuery
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.lineage.models import Dialect
+from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.models import TableView
+from metadata.utils import fqn
+from metadata.utils.db_utils import get_view_lineage
+from metadata.utils.logger import ingestion_logger
+from metadata.utils.time_utils import datetime_to_timestamp
+
+logger = ingestion_logger()
+
+# pylint: disable=invalid-name
+
+
+class QueryByProcedure(BaseModel):
+    """
+    Query(ies) executed by each stored procedure
+    """
+
+    procedure_name: str = Field(None, alias="PROCEDURE_NAME")
+    query_type: str = Field(..., alias="QUERY_TYPE")
+    query_database_name: Optional[str] = Field(None, alias="QUERY_DATABASE_NAME")
+    query_schema_name: Optional[str] = Field(None, alias="QUERY_SCHEMA_NAME")
+    procedure_text: str = Field(..., alias="PROCEDURE_TEXT")
+    procedure_start_time: datetime = Field(..., alias="PROCEDURE_START_TIME")
+    procedure_end_time: datetime = Field(..., alias="PROCEDURE_END_TIME")
+    query_start_time: Optional[datetime] = Field(None, alias="QUERY_START_TIME")
+    query_duration: Optional[float] = Field(None, alias="QUERY_DURATION")
+    query_text: str = Field(..., alias="QUERY_TEXT")
+    query_user_name: Optional[str] = Field(None, alias="QUERY_USER_NAME")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class ProcedureAndQuery(BaseModel):
+    """
+    Model to hold the procedure and its queries
+    """
+
+    procedure: StoredProcedure
+    query_by_procedure: QueryByProcedure
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+def is_lineage_query(query_type: str, query_text: str) -> bool:
+    """Check if it's worth it to parse the query for lineage"""
+
+    logger.debug(
+        f"Validating query lineage for type [{query_type}] and text [{query_text}]"
+    )
+
+    if query_type in ("MERGE", "UPDATE", "CREATE_TABLE_AS_SELECT"):
+        return True
+
+    if query_type == "INSERT" and re.search(
+        "^.*insert.*into.*select.*$", query_text.replace("\n", " "), re.IGNORECASE
+    ):
+        return True
+
+    return False
+
+
+def _yield_procedure_lineage(
+    metadata: OpenMetadata,
+    service_name: str,
+    dialect: Dialect,
+    parsingTimeoutLimit: int,
+    query_by_procedure: QueryByProcedure,
+    procedure: StoredProcedure,
+) -> Iterable[Either[AddLineageRequest]]:
+    """Add procedure lineage from its query"""
+    if is_lineage_query(
+        query_type=query_by_procedure.query_type,
+        query_text=query_by_procedure.query_text,
+    ):
+        for either_lineage in get_lineage_by_query(
+            metadata,
+            query=query_by_procedure.query_text,
+            service_name=service_name,
+            database_name=query_by_procedure.query_database_name,
+            schema_name=query_by_procedure.query_schema_name,
+            dialect=dialect,
+            timeout_seconds=parsingTimeoutLimit,
+            lineage_source=LineageSource.QueryLineage,
+        ):
+            if either_lineage.left is None and either_lineage.right.edge.lineageDetails:
+                either_lineage.right.edge.lineageDetails.pipeline = EntityReference(
+                    id=procedure.id,
+                    type="storedProcedure",
+                )
+
+            yield either_lineage
+
+
+def procedure_lineage_processor(
+    procedure_and_queries: List[ProcedureAndQuery],
+    queue: Queue,
+    metadata: OpenMetadata,
+    service_name: str,
+    dialect: Dialect,
+    parsingTimeoutLimit: int,
+) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+    """
+    Process the procedure and its queries to add lineage
+    """
+    for procedure_and_query in procedure_and_queries:
+        try:
+            for lineage in _yield_procedure_lineage(
+                query_by_procedure=procedure_and_query.query_by_procedure,
+                procedure=procedure_and_query.procedure,
+                metadata=deepcopy(metadata),
+                service_name=service_name,
+                dialect=dialect,
+                parsingTimeoutLimit=parsingTimeoutLimit,
+            ):
+                if lineage and lineage.right is not None:
+                    queue.put(
+                        Either(
+                            right=OMetaLineageRequest(
+                                override_lineage=False,
+                                lineage_request=lineage.right,
+                                entity=StoredProcedure,
+                                entity_fqn=procedure_and_query.procedure.fullyQualifiedName.root,
+                            )
+                        )
+                    )
+                else:
+                    queue.put(lineage)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Could not get lineage for store procedure "
+                f"'{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
+            )
+        try:
+            for lineage in yield_procedure_query(
+                query_by_procedure=procedure_and_query.query_by_procedure,
+                procedure=procedure_and_query.procedure,
+                service_name=service_name,
+            ):
+                queue.put(lineage)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Could not get query for store procedure "
+                f"'{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
+            )
+
+
+def yield_procedure_query(
+    query_by_procedure: QueryByProcedure, procedure: StoredProcedure, service_name: str
+) -> Iterable[Either[CreateQueryRequest]]:
+    """Check the queries triggered by the procedure and add their lineage, if any"""
+    stored_procedure_query_lineage = is_lineage_query(
+        query_type=query_by_procedure.query_type,
+        query_text=query_by_procedure.query_text,
+    )
+
+    yield Either(
+        right=CreateQueryRequest(
+            query=SqlQuery(query_by_procedure.query_text),
+            query_type=query_by_procedure.query_type,
+            duration=query_by_procedure.query_duration,
+            queryDate=Timestamp(
+                root=datetime_to_timestamp(query_by_procedure.query_start_time, True)
+            ),
+            triggeredBy=EntityReference(
+                id=procedure.id,
+                type="storedProcedure",
+            ),
+            processedLineage=bool(stored_procedure_query_lineage),
+            service=service_name,
+        )
+    )
+
+
+# Function that will run in separate processes - defined at module level for pickling
+def _process_chunk_in_subprocess(chunk, processor_fn, queue, *args):
+    """
+    Process a chunk of data in a subprocess.
+
+    Args:
+        chunk_and_processor_fn: Tuple containing (chunk, processor_fn, queue, *args)
+
+    Returns:
+        True if processing succeeded, False otherwise
+    """
+    try:
+        # Process each item in the chunk
+        processor_fn(chunk, queue, *args)
+        time.sleep(0.1)
+        return True
+    except Exception as e:
+        logger.error(f"Error processing chunk in subprocess: {e}")
+        logger.error(traceback.format_exc())
+        return False
+
+
+def _query_already_processed(metadata: OpenMetadata, table_query: TableQuery) -> bool:
+    """
+    Check if a query has already been processed by validating if exists
+    in ES with lineageProcessed as True
+    """
+    checksums = metadata.es_get_queries_with_lineage(
+        service_name=table_query.serviceName,
+    )
+    return fqn.get_query_checksum(table_query.query) in checksums or {}
+
+
+def query_lineage_generator(
+    table_queries: List[TableQuery],
+    queue: Queue,
+    metadata: OpenMetadata,
+    dialect: Dialect,
+    graph: nx.DiGraph,
+    parsingTimeoutLimit: int,
+    serviceName: str,
+) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+    """
+    Generate lineage for a list of table queries
+    """
+
+    for table_query in table_queries or []:
+        if not _query_already_processed(
+            metadata=deepcopy(metadata), table_query=table_query
+        ):
+            lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
+                metadata=deepcopy(metadata),
+                query=table_query.query,
+                service_name=table_query.serviceName,
+                database_name=table_query.databaseName,
+                schema_name=table_query.databaseSchema,
+                dialect=dialect,
+                timeout_seconds=parsingTimeoutLimit,
+                graph=graph,
+            )
+
+            for lineage_request in lineages or []:
+                queue.put(lineage_request)
+
+                # If we identified lineage properly, ingest the original query
+                if lineage_request.right:
+                    queue.put(
+                        Either(
+                            right=CreateQueryRequest(
+                                query=SqlQuery(table_query.query),
+                                query_type=table_query.query_type,
+                                duration=table_query.duration,
+                                processedLineage=True,
+                                service=serviceName,
+                            )
+                        )
+                    )
+
+
+def view_lineage_generator(
+    views: List[TableView],
+    queue: Queue,
+    metadata: OpenMetadata,
+    serviceName: str,
+    connectionType: str,
+    parsingTimeoutLimit: int,
+    overrideViewLineage: bool,
+) -> Iterable[Either[AddLineageRequest]]:
+    """
+    Generate lineage for a list of views
+    """
+    try:
+        for view in views:
+            for lineage in get_view_lineage(
+                view=view,
+                metadata=deepcopy(metadata),
+                service_name=serviceName,
+                connection_type=connectionType,
+                timeout_seconds=parsingTimeoutLimit,
+            ):
+                if lineage.right is not None:
+                    view_fqn = fqn.build(
+                        metadata=deepcopy(metadata),
+                        entity_type=Table,
+                        service_name=serviceName,
+                        database_name=view.db_name,
+                        schema_name=view.schema_name,
+                        table_name=view.table_name,
+                        skip_es_search=True,
+                    )
+                    queue.put(
+                        Either(
+                            right=OMetaLineageRequest(
+                                lineage_request=lineage.right,
+                                override_lineage=overrideViewLineage,
+                                entity_fqn=view_fqn,
+                                entity=Table,
+                            )
+                        )
+                    )
+                else:
+                    queue.put(lineage)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Error processing view {view}: {exc}")

metadata/ingestion/source/database/lineage_source.py

@@ -16,35 +16,40 @@ import os
 import time
 import traceback
 from abc import ABC
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor
 from functools import partial
-from typing import Any, Callable, Iterable, Iterator, List, Union
+from multiprocessing import Queue
+from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple, Union
+
+import networkx as nx
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.table import Table
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName, SqlQuery
+from metadata.generated.schema.type.basic import Uuid
+from metadata.generated.schema.type.entityLineage import (
+    ColumnLineage,
+    EntitiesEdge,
+    LineageDetails,
+    Source,
+)
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Dialect
-from metadata.ingestion.lineage.sql_lineage import (
-    get_lineage_by_graph,
-    get_lineage_by_query,
+from metadata.ingestion.lineage.sql_lineage import get_column_fqn, get_lineage_by_graph
+from metadata.ingestion.source.database.lineage_processors import (
+    _process_chunk_in_subprocess,
+    query_lineage_generator,
+    view_lineage_generator,
 )
-from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
-from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
-from metadata.ingestion.source.models import TableView
-from metadata.utils import fqn
-from metadata.utils.db_utils import get_view_lineage
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-CHUNK_SIZE = 200
-
-THREAD_TIMEOUT = 600
+CHUNK_SIZE = 100
 
 
 class LineageSource(QueryParserSource, ABC):
@@ -107,10 +112,11 @@ class LineageSource(QueryParserSource, ABC):
             )
             yield from self.yield_table_query()
 
-    def generate_lineage_in_thread(
+    def generate_lineage_with_processes(
         self,
         producer_fn: Callable[[], Iterable[Any]],
-        processor_fn: Callable[[Any], Iterable[Any]],
+        processor_fn: Callable[[Any, Queue], None],
+        args: Tuple[Any, ...],
         chunk_size: int = CHUNK_SIZE,
     ):
         """
@@ -123,6 +129,7 @@ class LineageSource(QueryParserSource, ABC):
         """
 
         def chunk_generator():
+
             temp_chunk = []
             for chunk in producer_fn():
                 temp_chunk.append(chunk)
@@ -133,34 +140,38 @@ class LineageSource(QueryParserSource, ABC):
             if temp_chunk:
                 yield temp_chunk
 
-        thread_pool = ThreadPoolExecutor(max_workers=self.source_config.threads)
-        queue = Queue()
+        from multiprocessing import Manager
+
+        manager = Manager()
+        queue = manager.Queue()
+
+        process_pool = ProcessPoolExecutor(max_workers=self.source_config.threads)
 
         futures = [
-            thread_pool.submit(
-                processor_fn,
-                chunk,
-                queue,
+            process_pool.submit(
+                _process_chunk_in_subprocess, chunk, processor_fn, queue, *args
             )
             for chunk in chunk_generator()
         ]
         while True:
-            if queue.has_tasks():
-                yield from queue.process()
-
-            else:
-                if not futures:
-                    break
-
-                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)
-
+            try:
+                while not queue.empty():
+                    yield queue.get_nowait()
+            except Exception as exc:
+                logger.warning(f"Error processing queue: {exc}")
+                logger.debug(traceback.format_exc())
+
+            if not futures:
+                break
+
+            for i, future in enumerate(futures):
+                if future.done():
+                    try:
+                        future.result(timeout=0)
+                    except Exception as e:
+                        logger.debug(f"Error in future: {e}")
+                        logger.debug(traceback.format_exc())
+                    futures.pop(i)
             time.sleep(0.01)
 
     def yield_table_query(self) -> Iterator[TableQuery]:
@@ -193,57 +204,6 @@ class LineageSource(QueryParserSource, ABC):
                             f"Error processing query_dict {query_dict}: {exc}"
                         )
 
-    def _query_already_processed(self, table_query: TableQuery) -> bool:
-        """
-        Check if a query has already been processed by validating if exists
-        in ES with lineageProcessed as True
-        """
-        checksums = self.metadata.es_get_queries_with_lineage(
-            service_name=table_query.serviceName,
-        )
-        return fqn.get_query_checksum(table_query.query) in checksums or {}
-
-    def query_lineage_generator(
-        self, table_queries: List[TableQuery], queue: Queue
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        if self.graph is None and self.source_config.enableTempTableLineage:
-            import networkx as nx
-
-            # Create a directed graph
-            self.graph = nx.DiGraph()
-
-        for table_query in table_queries or []:
-            if not self._query_already_processed(table_query):
-                lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
-                    self.metadata,
-                    query=table_query.query,
-                    service_name=table_query.serviceName,
-                    database_name=table_query.databaseName,
-                    schema_name=table_query.databaseSchema,
-                    dialect=self.dialect,
-                    timeout_seconds=self.source_config.parsingTimeoutLimit,
-                    graph=self.graph,
-                )
-
-                for lineage_request in lineages or []:
-                    queue.put(lineage_request)
-
-                    # If we identified lineage properly, ingest the original query
-                    if lineage_request.right:
-                        queue.put(
-                            Either(
-                                right=CreateQueryRequest(
-                                    query=SqlQuery(table_query.query),
-                                    query_type=table_query.query_type,
-                                    duration=table_query.duration,
-                                    processedLineage=True,
-                                    service=FullyQualifiedEntityName(
-                                        self.config.serviceName
-                                    ),
-                                )
-                            )
-                        )
-
     def yield_query_lineage(
         self,
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
@@ -255,48 +215,19 @@ class LineageSource(QueryParserSource, ABC):
         connection_type = str(self.service_connection.type.value)
         self.dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
         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
+        processor_fn = query_lineage_generator
+        args = (
+            self.metadata,
+            self.dialect,
+            self.graph,
+            self.source_config.parsingTimeoutLimit,
+            self.config.serviceName,
+        )
+        yield from self.generate_lineage_with_processes(
+            producer_fn,
+            processor_fn,
+            args,
         )
-
-    def view_lineage_generator(
-        self, views: List[TableView], queue: Queue
-    ) -> Iterable[Either[AddLineageRequest]]:
-        try:
-            for view in views:
-                for lineage in get_view_lineage(
-                    view=view,
-                    metadata=self.metadata,
-                    service_name=self.config.serviceName,
-                    connection_type=self.service_connection.type.value,
-                    timeout_seconds=self.source_config.parsingTimeoutLimit,
-                ):
-                    if lineage.right is not None:
-                        view_fqn = fqn.build(
-                            metadata=self.metadata,
-                            entity_type=Table,
-                            service_name=self.service_name,
-                            database_name=view.db_name,
-                            schema_name=view.schema_name,
-                            table_name=view.table_name,
-                            skip_es_search=True,
-                        )
-                        queue.put(
-                            Either(
-                                right=OMetaLineageRequest(
-                                    lineage_request=lineage.right,
-                                    override_lineage=self.source_config.overrideViewLineage,
-                                    entity_fqn=view_fqn,
-                                    entity=Table,
-                                )
-                            )
-                        )
-                    else:
-                        queue.put(lineage)
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Error processing view {view}: {exc}")
 
     def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
         logger.info("Processing View Lineage")
@@ -305,19 +236,86 @@ class LineageSource(QueryParserSource, ABC):
             self.config.serviceName,
             self.source_config.incrementalLineageProcessing,
         )
-        processor_fn = self.view_lineage_generator
-        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)
+        processor_fn = view_lineage_generator
+        args = (
+            self.metadata,
+            self.config.serviceName,
+            self.service_connection.type.value,
+            self.source_config.parsingTimeoutLimit,
+            self.source_config.overrideViewLineage,
+        )
+        yield from self.generate_lineage_with_processes(producer_fn, processor_fn, args)
 
     def yield_procedure_lineage(
         self,
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """
-        By default stored procedure lineage is not supported.
+        By default stored   procedure lineage is not supported.
         """
         logger.info(
             f"Processing Procedure Lineage not supported for {str(self.service_connection.type.value)}"
         )
 
+    def get_column_lineage(
+        self, from_table: Table, to_table: Table
+    ) -> List[ColumnLineage]:
+        """
+        Get the column lineage from the fields
+        """
+        try:
+            column_lineage = []
+            for column in from_table.columns:
+                field = column.name.root
+                from_column = get_column_fqn(table_entity=from_table, column=field)
+                to_column = get_column_fqn(table_entity=to_table, column=field)
+                if from_column and to_column:
+                    column_lineage.append(
+                        ColumnLineage(fromColumns=[from_column], toColumn=to_column)
+                    )
+
+            return column_lineage
+        except Exception as exc:
+            logger.debug(f"Error to get column lineage: {exc}")
+            logger.debug(traceback.format_exc())
+        return []
+
+    def get_add_cross_database_lineage_request(
+        self,
+        from_entity: Table,
+        to_entity: Table,
+        column_lineage: List[ColumnLineage] = None,
+    ) -> Optional[Either[AddLineageRequest]]:
+        """
+        Get the add cross database lineage request
+        """
+        if from_entity and to_entity:
+            return Either(
+                right=AddLineageRequest(
+                    edge=EntitiesEdge(
+                        fromEntity=EntityReference(
+                            id=Uuid(from_entity.id.root), type="table"
+                        ),
+                        toEntity=EntityReference(
+                            id=Uuid(to_entity.id.root), type="table"
+                        ),
+                        lineageDetails=LineageDetails(
+                            source=Source.CrossDatabaseLineage,
+                            columnsLineage=column_lineage,
+                        ),
+                    )
+                )
+            )
+
+        return None
+
+    def yield_cross_database_lineage(self) -> Iterable[Either[AddLineageRequest]]:
+        """
+        By default cross database lineage is not supported.
+        """
+        logger.info(
+            f"Processing Cross Database Lineage not supported for {str(self.service_connection.type.value)}"
+        )
+
     def _iter(
         self, *_, **__
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
@@ -325,6 +323,9 @@ class LineageSource(QueryParserSource, ABC):
         Based on the query logs, prepare the lineage
         and send it to the sink
         """
+        if self.graph is None and self.source_config.enableTempTableLineage:
+            # Create a directed graph
+            self.graph = nx.DiGraph()
         if self.source_config.processViewLineage:
             yield from self.yield_view_lineage() or []
         if self.source_config.processStoredProcedureLineage: