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

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