openmetadata-ingestion 1.6.10.0__py3-none-any.whl → 1.6.12.0__py3-none-any.whl

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
@@ -56,7 +63,7 @@ logger = utils_logger()
 DEFAULT_SCHEMA_NAME = "<default>"
 CUTOFF_NODES = 20
 
-
+# pylint: disable=too-many-function-args,protected-access
 def get_column_fqn(table_entity: Table, column: str) -> Optional[str]:
     """
     Get fqn of column if exist in table entity
@@ -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 Exception as _:
+                    # 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/models/topology.py

@@ -305,6 +305,15 @@ class TopologyContextManager:
             thread_id, self.contexts[parent_thread_id].model_copy(deep=True)
         )
 
+    def __getstate__(self):
+        """Called when pickling the object."""
+        state = self.__dict__.copy()
+        return state
+
+    def __setstate__(self, state):
+        """Called when unpickling the object."""
+        self.__dict__.update(state)
+
 
 class Queue:
     """Small Queue wrapper"""

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/ometa_api.py

@@ -529,6 +529,42 @@ class OpenMetadata(
         raw_version = self.client.get("/system/version")["version"]
         return raw_version is not None
 
+    def __getstate__(self):
+        """Called when pickling the object, returns the state without the unpicklable objects."""
+        state = self.__dict__.copy()
+        # Remove objects that might contain thread locks
+        if "_auth_provider" in state:
+            del state["_auth_provider"]
+        if "client" in state:
+            del state["client"]
+        if "secrets_manager_client" in state:
+            del state["secrets_manager_client"]
+        return state
+
+    def __setstate__(self, state):
+        """Called when unpickling the object, restores the state."""
+        self.__dict__.update(state)
+        # Restore the auth provider and client
+        self._auth_provider = OpenMetadataAuthenticationProvider.create(self.config)
+
+        get_verify_ssl = get_verify_ssl_fn(self.config.verifySSL)
+
+        client_config: ClientConfig = ClientConfig(
+            base_url=self.config.hostPort,
+            api_version=self.config.apiVersion,
+            auth_header="Authorization",
+            extra_headers=self.config.extraHeaders,
+            auth_token=self._auth_provider.get_access_token,
+            verify=get_verify_ssl(self.config.sslConfig),
+        )
+        self.client = REST(client_config)
+
+        # Restore the secrets manager client
+        self.secrets_manager_client = SecretsManagerFactory(
+            self.config.secretsManagerProvider,
+            self.config.secretsManagerLoader,
+        ).get_secrets_manager()
+
     def close(self):
         """
         Closing connection

metadata/ingestion/source/dashboard/powerbi/metadata.py

@@ -286,7 +286,7 @@ class PowerbiSource(DashboardServiceSource):
         """
         return (
             f"{clean_uri(self.service_connection.hostPort)}/groups/"
-            f"{workspace_id}/dashboards/{dashboard_id}"
+            f"{workspace_id}/dashboards/{dashboard_id}?experience=power-bi"
         )
 
     def _get_report_url(self, workspace_id: str, dashboard_id: str) -> str:

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}")