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

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:

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

@@ -0,0 +1,37 @@
+"""Snowflake spec for data diff"""
+
+from typing import Optional, cast
+
+from metadata.data_quality.validations.models import TableParameter
+from metadata.data_quality.validations.runtime_param_setter.base_diff_params_setter import (
+    BaseTableParameter,
+)
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.ingestion.source.database.snowflake.metadata import SnowflakeConnection
+
+
+class SnowflakeTableParameter(BaseTableParameter):
+    """SnowflakeTableParameter class for setting runtime parameters for data diff"""
+
+    def get(
+        self,
+        service: DatabaseService,
+        entity: Table,
+        key_columns,
+        extra_columns,
+        case_sensitive_columns,
+        service_url: Optional[str],
+    ) -> TableParameter:
+        table_param: TableParameter = super().get(
+            service,
+            entity,
+            key_columns,
+            extra_columns,
+            case_sensitive_columns,
+            service_url,
+        )
+        connection_config = cast(SnowflakeConnection, service.connection.config)
+        table_param.privateKey = connection_config.privateKey
+        table_param.passPhrase = connection_config.snowflakePrivatekeyPassphrase
+        return table_param

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

@@ -1,6 +1,9 @@
 from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import (
     SnowflakeTestSuiteInterface,
 )
+from metadata.ingestion.source.database.snowflake.data_diff.data_diff import (
+    SnowflakeTableParameter,
+)
 from metadata.ingestion.source.database.snowflake.lineage import SnowflakeLineageSource
 from metadata.ingestion.source.database.snowflake.metadata import SnowflakeSource
 from metadata.ingestion.source.database.snowflake.profiler.profiler import (
@@ -17,4 +20,5 @@ ServiceSpec = DefaultDatabaseSpec(
     profiler_class=SnowflakeProfiler,
     test_suite_class=SnowflakeTestSuiteInterface,
     sampler_class=SnowflakeSampler,
+    data_diff=SnowflakeTableParameter,
 )

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -12,14 +12,12 @@
 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 copy import deepcopy
+from typing import Dict, Iterable, List, Union
 
-from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -31,54 +29,21 @@ 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.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 StoredProcedureLineageMixin(ABC):
     """
     The full flow is:
@@ -138,121 +103,10 @@ class StoredProcedureLineageMixin(ABC):
 
         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,
-        ):
-            self.stored_procedure_query_lineage = True
-            for either_lineage in get_lineage_by_query(
-                self.metadata,
-                query=query_by_procedure.query_text,
-                service_name=self.service_name,
-                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,
-            ):
-                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]:
+        """
+        Generate lineage for a list of stored procedures
+        """
         query = {
             "query": {
                 "bool": {
@@ -303,5 +157,14 @@ class StoredProcedureLineageMixin(ABC):
         """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)
+        processor_fn = procedure_lineage_processor
+        dialect = ConnectionTypeDialectMapper.dialect_of(
+            self.service_connection.type.value
+        )
+        args = (
+            deepcopy(self.metadata),
+            self.service_name,
+            dialect,
+            self.source_config.parsingTimeoutLimit,
+        )
+        yield from self.generate_lineage_with_processes(producer_fn, processor_fn, args)

metadata/readers/dataframe/parquet.py

@@ -68,8 +68,16 @@ class ParquetDataFrameReader(DataFrameReader):
         from pyarrow.parquet import ParquetDataset
 
         client_kwargs = {
-            "endpoint_override": str(self.config_source.securityConfig.endPointURL),
-            "region": self.config_source.securityConfig.awsRegion,
+            "endpoint_override": (
+                str(self.config_source.securityConfig.endPointURL)
+                if self.config_source.securityConfig.endPointURL
+                else None
+            ),
+            "region": (
+                self.config_source.securityConfig.awsRegion
+                if self.config_source.securityConfig.awsRegion
+                else None
+            ),
             "access_key": self.config_source.securityConfig.awsAccessKeyId,
             "session_token": self.config_source.securityConfig.awsSessionToken,
             "role_arn": self.config_source.securityConfig.assumeRoleArn,