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

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

@@ -1,4 +1,12 @@
+from metadata.ingestion.source.database.mariadb.lineage import MariadbLineageSource
 from metadata.ingestion.source.database.mariadb.metadata import MariadbSource
+from metadata.profiler.interface.sqlalchemy.mariadb.profiler_interface import (
+    MariaDBProfilerInterface,
+)
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=MariadbSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MariadbSource,
+    lineage_source_class=MariadbLineageSource,
+    profiler_class=MariaDBProfilerInterface,
+)

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

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.pinotdb.lineage import PinotdbLineageSource
 from metadata.ingestion.source.database.pinotdb.metadata import PinotdbSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=PinotdbSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=PinotdbSource,
+    lineage_source_class=PinotdbLineageSource,
+)

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

@@ -1,3 +1,6 @@
+from metadata.ingestion.source.database.singlestore.lineage import (
+    SinglestoreLineageSource,
+)
 from metadata.ingestion.source.database.singlestore.metadata import SinglestoreSource
 from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface import (
     SingleStoreProfilerInterface,
@@ -5,5 +8,7 @@ from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface impo
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=SinglestoreSource, profiler_class=SingleStoreProfilerInterface
+    metadata_source_class=SinglestoreSource,
+    profiler_class=SingleStoreProfilerInterface,
+    lineage_source_class=SinglestoreLineageSource,
 )

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/sqlite/service_spec.py

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.sqlite.lineage import SqliteLineageSource
 from metadata.ingestion.source.database.sqlite.metadata import SqliteSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=SqliteSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=SqliteSource,
+    lineage_source_class=SqliteLineageSource,
+)

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/ingestion/source/database/teradata/service_spec.py

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.teradata.lineage import TeradataLineageSource
 from metadata.ingestion.source.database.teradata.metadata import TeradataSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=TeradataSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=TeradataSource,
+    lineage_source_class=TeradataLineageSource,
+)

metadata/utils/db_utils.py

@@ -12,7 +12,7 @@
 """
 Helpers module for db sources
 """
-
+import time
 import traceback
 from typing import Iterable
 
@@ -32,6 +32,7 @@ from metadata.ingestion.lineage.sql_lineage import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
+from metadata.utils.execution_time_tracker import calculate_execution_time_generator
 from metadata.utils.logger import utils_logger
 
 logger = utils_logger()
@@ -47,6 +48,7 @@ def get_host_from_host_port(uri: str) -> str:
     return uri.split(":")[0]
 
 
+@calculate_execution_time_generator()
 def get_view_lineage(
     view: TableView,
     metadata: OpenMetadata,
@@ -81,6 +83,8 @@ def get_view_lineage(
     try:
         connection_type = str(connection_type)
         dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
+        start_time = time.time()
+        logger.debug(f"Processing view lineage for: {table_fqn}")
         lineage_parser = LineageParser(
             view_definition, dialect, timeout_seconds=timeout_seconds
         )
@@ -89,6 +93,10 @@ def get_view_lineage(
             # For Postgres, if schema is not defined, we need to use the public schema
             schema_name = PUBLIC_SCHEMA
 
+        end_time = time.time()
+        logger.debug(
+            f"Time taken to parse view lineage for: {table_fqn} is {end_time - start_time} seconds"
+        )
         if lineage_parser.source_tables and lineage_parser.target_tables:
             yield from get_lineage_by_query(
                 metadata,
@@ -99,6 +107,7 @@ def get_view_lineage(
                 dialect=dialect,
                 timeout_seconds=timeout_seconds,
                 lineage_source=LineageSource.ViewLineage,
+                lineage_parser=lineage_parser,
             ) or []
 
         else:
@@ -112,6 +121,7 @@ def get_view_lineage(
                 dialect=dialect,
                 timeout_seconds=timeout_seconds,
                 lineage_source=LineageSource.ViewLineage,
+                lineage_parser=lineage_parser,
             ) or []
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/utils/execution_time_tracker.py

@@ -75,6 +75,14 @@ class ExecutionTimeTrackerContextMap(metaclass=Singleton):
         thread_id = thread_id or threading.get_ident()
         return self.map.get(thread_id, []).pop()
 
+    def __getstate__(self):
+        """Called when pickling the object, returns the state without thread-specific objects."""
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        """Called when unpickling the object."""
+        self.__dict__.update(state)
+
 
 class ExecutionTimeTrackerState(metaclass=Singleton):
     """Tracks the ExecutionTime State across multiple threads."""
@@ -89,6 +97,20 @@ class ExecutionTimeTrackerState(metaclass=Singleton):
         with self.lock:
             self.state[context.name] = self.state.get(context.name, 0) + elapsed
 
+    def __getstate__(self):
+        """Called when pickling the object, returns the state without the lock."""
+        state = self.__dict__.copy()
+        # Don't pickle the lock
+        if "lock" in state:
+            del state["lock"]
+        return state
+
+    def __setstate__(self, state):
+        """Called when unpickling the object, restores the lock."""
+        self.__dict__.update(state)
+        # Restore the lock
+        self.lock = threading.Lock()
+
 
 class ExecutionTimeTracker(metaclass=Singleton):
     """ExecutionTimeTracker is implemented as a Singleton in order to hold state globally.
@@ -173,6 +195,14 @@ class ExecutionTimeTracker(metaclass=Singleton):
             if context.stored:
                 self.state.add(context, elapsed)
 
+    def __getstate__(self):
+        """Called when pickling the object."""
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        """Called when unpickling the object."""
+        self.__dict__.update(state)
+
 
 def calculate_execution_time(context: Optional[str] = None, store: bool = True):
     """Utility decorator to be able to use the ExecutionTimeTracker on a function.

metadata/utils/lru_cache.py

@@ -78,6 +78,20 @@ class LRUCache(Generic[T]):
         with self.lock:
             return len(self._cache)
 
+    def __getstate__(self):
+        """Called when pickling the object, returns the state without the lock."""
+        state = self.__dict__.copy()
+        # Don't pickle the lock
+        if "lock" in state:
+            del state["lock"]
+        return state
+
+    def __setstate__(self, state):
+        """Called when unpickling the object, restores the lock."""
+        self.__dict__.update(state)
+        # Restore the lock
+        self.lock = threading.Lock()
+
     def wrap(self, key_func: Callable[..., str]):
         """Decorator to cache the result of a function based on its arguments.
 

metadata/utils/service_spec/default.py

@@ -7,6 +7,9 @@ from typing import Optional
 from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
     SQATestSuiteInterface,
 )
+from metadata.data_quality.validations.runtime_param_setter.base_diff_params_setter import (
+    BaseTableParameter,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
@@ -19,3 +22,4 @@ class DefaultDatabaseSpec(BaseSpec):
     profiler_class: Optional[str] = get_class_path(SQAProfilerInterface)
     sampler_class: Optional[str] = get_class_path(SQASampler)
     test_suite_class: Optional[str] = get_class_path(SQATestSuiteInterface)
+    data_diff: Optional[str] = get_class_path(BaseTableParameter)

metadata/utils/service_spec/service_spec.py

@@ -52,6 +52,7 @@ class BaseSpec(BaseModel):
     lineage_source_class: Optional[str] = None
     usage_source_class: Optional[str] = None
     sampler_class: Optional[str] = None
+    data_diff: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod