openmetadata-ingestion 1.9.5.0__py3-none-any.whl → 1.9.7.0__py3-none-any.whl

metadata/ingestion/source/database/mssql/lineage.py

@@ -11,13 +11,8 @@
 """
 MSSQL lineage module
 """
-import traceback
 from datetime import datetime
-from typing import Dict, List
 
-from metadata.generated.schema.entity.services.ingestionPipelines.status import (
-    StackTraceError,
-)
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.mssql.constants import (
     DEFAULT_DATETIME_FORMAT,
@@ -32,13 +27,9 @@ from metadata.ingestion.source.database.mssql.utils import (
     get_sqlalchemy_engine_dateformat,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
     StoredProcedureLineageMixin,
 )
 from metadata.utils.helpers import get_start_and_end
-from metadata.utils.logger import ingestion_logger
-
-logger = ingestion_logger()
 
 
 class MssqlLineageSource(
@@ -74,10 +65,9 @@ class MssqlLineageSource(
             result_limit=self.source_config.resultLimit,
         )
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         server_date_format = get_sqlalchemy_engine_dateformat(self.engine)
@@ -88,19 +78,5 @@ class MssqlLineageSource(
         query = MSSQL_GET_STORED_PROCEDURE_QUERIES.format(
             start_date=start,
         )
-        try:
-            queries_dict = self.procedure_queries_dict(
-                query=query,
-            )
-        except Exception as ex:  # pylint: disable=broad-except
-            logger.debug(f"Error runnning query:\n{query}")
-            self.status.failed(
-                StackTraceError(
-                    name="Stored Procedure",
-                    error=f"Error trying to get stored procedure queries: {ex}",
-                    stackTrace=traceback.format_exc(),
-                )
-            )
-            return {}
 
-        return queries_dict
+        return query

metadata/ingestion/source/database/oracle/lineage.py

@@ -30,8 +30,6 @@ workflowConfig:
       jwtToken: "token"
 """
 
-from typing import Dict, List
-
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.oracle.queries import (
     ORACLE_GET_STORED_PROCEDURE_QUERIES,
@@ -41,7 +39,6 @@ from metadata.ingestion.source.database.oracle.query_parser import (
     OracleQueryParserSource,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
     StoredProcedureLineageMixin,
 )
 from metadata.utils.helpers import get_start_and_end
@@ -63,17 +60,13 @@ class OracleLineageSource(
 
     stored_procedure_query = ORACLE_GET_STORED_PROCEDURE_QUERIES
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = self.stored_procedure_query.format(
             start_date=start,
         )
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
 
-        return queries_dict
+        return query

metadata/ingestion/source/database/redshift/lineage.py

@@ -31,7 +31,7 @@ workflowConfig:
 """
 
 import traceback
-from typing import Dict, Iterator, List
+from typing import Iterator
 
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.source.database.lineage_source import LineageSource
@@ -43,7 +43,6 @@ from metadata.ingestion.source.database.redshift.query_parser import (
     RedshiftQueryParserSource,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
     StoredProcedureLineageMixin,
 )
 from metadata.utils.helpers import get_start_and_end
@@ -97,16 +96,11 @@ class RedshiftLineageSource(
                             f"Error processing query_dict {query_dict}: {exc}"
                         )
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES.format(start_date=start)
 
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
-
-        return queries_dict
+        return query

metadata/ingestion/source/database/snowflake/lineage.py

@@ -13,7 +13,7 @@ Snowflake lineage module
 """
 
 import traceback
-from typing import Dict, Iterator, List
+from typing import Iterator
 
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.source.database.lineage_source import LineageSource
@@ -25,7 +25,6 @@ from metadata.ingestion.source.database.snowflake.query_parser import (
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
     StoredProcedureLineageMixin,
 )
 from metadata.utils.helpers import get_start_and_end
@@ -55,21 +54,17 @@ class SnowflakeLineageSource(
 
     stored_procedure_query = SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = self.stored_procedure_query.format(
             start_date=start,
             account_usage=self.service_connection.accountUsageSchema,
         )
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
 
-        return queries_dict
+        return query
 
     def yield_table_query(self) -> Iterator[TableQuery]:
         """

metadata/ingestion/source/database/snowflake/metadata.py

@@ -672,7 +672,7 @@ class SnowflakeSource(
         self, database_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> str:
         url = (
-            f"https://app.snowflake.com/{self.org_name.lower()}"
+            f"https://{self.service_connection.snowflakeSourceHost}/{self.org_name.lower()}"
             f"/{self.account.lower()}/#/data/databases/{database_name}"
         )
         if schema_name:
@@ -898,6 +898,8 @@ class SnowflakeSource(
         """
         Get columns of table/view/stream
         """
+        # For streams, we will use source table/view's columns
+        # since stream does not define columns separately in Snowflake
         if table_type == TableType.Stream:
             cursor = self.connection.execute(
                 SNOWFLAKE_GET_STREAM.format(stream_name=table_name, schema=schema_name)
@@ -906,6 +908,13 @@ class SnowflakeSource(
                 result = cursor.fetchone()
                 if result:
                     table_name = result[6].split(".")[-1]
+                    # Can't fetch source of stream is source is dropped or no priviledge
+                    if table_name == "No privilege or table dropped":
+                        logger.debug(
+                            f"Couldn't fetch columns of stream [{result and result[1]}],"
+                            f" due to error on source: {table_name}. Result: {result}"
+                        )
+                        return []
             except Exception:
                 pass
 

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -12,15 +12,11 @@
 Mixin class with common Stored Procedures logic aimed at lineage.
 """
 import json
-import re
 import traceback
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from datetime import datetime
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Iterator, Union
 
-from networkx import DiGraph
-from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -32,16 +28,15 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.databaseServiceQueryLineagePipeline import (
     DatabaseServiceQueryLineagePipeline,
 )
-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.ingestion.api.models import Either
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
-from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
-from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
-from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.lineage_processors import (
+    ProcedureAndQuery,
+    QueryByProcedure,
+    procedure_lineage_processor,
+)
 from metadata.utils.filters import (
     filter_by_database,
     filter_by_schema,
@@ -50,53 +45,10 @@ from metadata.utils.filters import (
 from metadata.utils.helpers import pprint_format_object
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.stored_procedures import get_procedure_name_from_call
-from metadata.utils.time_utils import datetime_to_timestamp
 
 logger = ingestion_logger()
 
 
-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)
-
-
-class ProcedureAndProcedureGraph(BaseModel):
-    """
-    Model to hold the procedure and its graph
-    """
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    procedure: StoredProcedure
-    graph: DiGraph
-
-
 class StoredProcedureLineageMixin(ABC):
     """
     The full flow is:
@@ -118,33 +70,28 @@ class StoredProcedureLineageMixin(ABC):
     metadata: OpenMetadata
 
     @abstractmethod
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
 
-    def procedure_queries_dict(self, query: str) -> Dict[str, List[QueryByProcedure]]:
+    def yield_stored_procedure_queries(self) -> Iterator[QueryByProcedure]:
         """
-        Cache the queries ran for the stored procedures in the last `queryLogDuration` days.
-
-        We will run this for each different and db name.
-
-        The dictionary key will be the case-insensitive procedure name.
+        Yield query and stored procedure object for lineage processing.
         """
+        query = self.get_stored_procedure_sql_statement()
         results = self.engine.execute(query).all()
-        queries_dict = defaultdict(list)
 
         for row in results:
             try:
                 query_by_procedure = QueryByProcedure.model_validate(dict(row))
-                procedure_name = (
+                query_by_procedure.procedure_name = (
                     query_by_procedure.procedure_name
                     or get_procedure_name_from_call(
                         query_text=query_by_procedure.procedure_text,
                     )
                 )
-                queries_dict[procedure_name].append(query_by_procedure)
+                yield query_by_procedure
             except Exception as exc:
                 self.status.failed(
                     StackTraceError(
@@ -154,155 +101,10 @@ class StoredProcedureLineageMixin(ABC):
                     )
                 )
 
-        queries_count_per_procedure = {
-            procedure_name: len(queries)
-            for procedure_name, queries in queries_dict.items()
-        }
-        logger.info(
-            f"Count of queries executed for stored procedures: {sum(queries_count_per_procedure.values())}"
-        )
-        logger.info(
-            f"Count of queries per stored procedure: {pprint_format_object(queries_count_per_procedure)}"
-        )
-
-        return queries_dict
-
-    @staticmethod
-    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(
-        self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
-    ) -> Iterable[Either[AddLineageRequest]]:
-        """Add procedure lineage from its query"""
-        self.stored_procedure_query_lineage = False
-        if self.is_lineage_query(
-            query_type=query_by_procedure.query_type,
-            query_text=query_by_procedure.query_text,
-        ):
-            graph = None
-            if self.source_config.enableTempTableLineage:
-                if not self.procedure_graph_map.get(procedure.fullyQualifiedName.root):
-                    # Map to store the directed graph for each procedure with its FQN as key
-                    self.procedure_graph_map[
-                        procedure.fullyQualifiedName.root
-                    ] = ProcedureAndProcedureGraph(procedure=procedure, graph=DiGraph())
-
-                graph = self.procedure_graph_map.get(
-                    procedure.fullyQualifiedName.root
-                ).graph
-
-            self.stored_procedure_query_lineage = True
-            # Prepare service names for lineage processing
-            service_names = [self.service_name]
-            if (
-                self.source_config.processCrossDatabaseLineage
-                and self.source_config.crossDatabaseServiceNames
-            ):
-                service_names.extend(self.source_config.crossDatabaseServiceNames)
-
-            for either_lineage in get_lineage_by_query(
-                self.metadata,
-                query=query_by_procedure.query_text,
-                service_names=service_names,
-                database_name=query_by_procedure.query_database_name,
-                schema_name=query_by_procedure.query_schema_name,
-                dialect=ConnectionTypeDialectMapper.dialect_of(
-                    self.service_connection.type.value
-                ),
-                timeout_seconds=self.source_config.parsingTimeoutLimit,
-                lineage_source=LineageSource.QueryLineage,
-                graph=graph,
-            ):
-                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 yield_procedure_query(
-        self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
-    ) -> Iterable[Either[CreateQueryRequest]]:
-        """Check the queries triggered by the procedure and add their lineage, if any"""
-
-        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(self.stored_procedure_query_lineage),
-                service=self.service_name,
-            )
-        )
-
-    def procedure_lineage_processor(
-        self, procedure_and_queries: List[ProcedureAndQuery], queue: Queue
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        for procedure_and_query in procedure_and_queries:
-            try:
-                for lineage in self._yield_procedure_lineage(
-                    query_by_procedure=procedure_and_query.query_by_procedure,
-                    procedure=procedure_and_query.procedure,
-                ):
-                    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 '{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
-                )
-            try:
-                for lineage in self.yield_procedure_query(
-                    query_by_procedure=procedure_and_query.query_by_procedure,
-                    procedure=procedure_and_query.procedure,
-                ):
-                    queue.put(lineage)
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Could not get query for store procedure '{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
-                )
-
-    def procedure_lineage_generator(self) -> Iterable[ProcedureAndQuery]:
+    def procedure_lineage_producer(self) -> Iterator[ProcedureAndQuery]:
+        """
+        Generate lineage for a list of stored procedures
+        """
         query = {
             "query": {
                 "bool": {
@@ -329,9 +131,12 @@ class StoredProcedureLineageMixin(ABC):
             )
         query_filter = json.dumps(query)
         logger.info("Processing Lineage for Stored Procedures")
-        # First, get all the query history
-        queries_dict = self.get_stored_procedure_queries_dict()
-        # Then for each procedure, iterate over all its queries
+
+        procedures_dict = {}
+        queries = self.yield_stored_procedure_queries()
+        queries_count_per_procedure = defaultdict(int)
+
+        # Get the filtered list of stored procedure to process
         for procedure in (
             self.metadata.paginate_es(
                 entity=StoredProcedure, query_filter=query_filter, size=10
@@ -359,20 +164,49 @@ class StoredProcedureLineageMixin(ABC):
                     )
                     continue
                 logger.debug(f"Processing Lineage for [{procedure.name}]")
-                for query_by_procedure in (
-                    queries_dict.get(procedure.name.root.lower()) or []
-                ):
-                    yield ProcedureAndQuery(
-                        procedure=procedure, query_by_procedure=query_by_procedure
-                    )
+                procedures_dict[procedure.name.root.lower()] = procedure
+
+        # Yield the ProcedureAndQuery for filtered stored procedure
+        for query_by_procedure in queries:
+            procedure_name = query_by_procedure.procedure_name.lower()
+            queries_count_per_procedure[procedure_name] += 1
+
+            if procedure_name in procedures_dict:
+                yield ProcedureAndQuery(
+                    procedure=procedures_dict[procedure_name],
+                    query_by_procedure=query_by_procedure,
+                )
+
+        logger.info(
+            f"Count of queries executed for stored procedures: {sum(queries_count_per_procedure.values())}"
+        )
+        logger.info(
+            f"Count of queries per stored procedure: {pprint_format_object(dict(queries_count_per_procedure))}"
+        )
 
     def yield_procedure_lineage(
         self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+    ) -> Iterator[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """Get all the queries and procedures list and yield them"""
         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, max_threads=self.source_config.threads
+        producer_fn = self.procedure_lineage_producer
+        processor_fn = procedure_lineage_processor
+        dialect = ConnectionTypeDialectMapper.dialect_of(
+            self.service_connection.type.value
+        )
+        args = (
+            self.metadata,
+            self.service_name,
+            dialect,
+            self.source_config.processCrossDatabaseLineage,
+            self.source_config.crossDatabaseServiceNames,
+            self.source_config.parsingTimeoutLimit,
+            self.procedure_graph_map,
+            self.source_config.enableTempTableLineage,
+        )
+        yield from self.generate_lineage_with_processes(
+            producer_fn,
+            processor_fn,
+            args,
+            max_threads=self.source_config.threads,
         )