openmetadata-ingestion 1.8.0.0__py3-none-any.whl → 1.8.2.0__py3-none-any.whl

metadata/profiler/metrics/system/system.py

@@ -15,16 +15,20 @@ System Metric
 
 from abc import ABC
 from collections import defaultdict
-from typing import Callable, Generic, List, TypeVar
+from typing import Callable, Dict, Generic, List, Optional, Protocol, Type, TypeVar
 
+from sqlalchemy import text
 from sqlalchemy.orm import Session
 
 from metadata.generated.schema.configuration.profilerConfiguration import MetricType
 from metadata.generated.schema.entity.data.table import SystemProfile
 from metadata.profiler.metrics.core import SystemMetric
+from metadata.profiler.orm.registry import PythonDialects
 from metadata.utils.helpers import deep_size_of_dict
+from metadata.utils.importer import DynamicImportException, import_from_module
 from metadata.utils.logger import profiler_logger
 from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
+from metadata.utils.profiler_utils import QueryResult
 
 logger = profiler_logger()
 
@@ -45,7 +49,23 @@ class CacheProvider(ABC, Generic[T]):
     """Cache provider class to provide cache for system metrics"""
 
     def __init__(self):
-        self.cache = LRUCache[T](LRU_CACHE_SIZE)
+        self.cache = LRUCache[List[T]](LRU_CACHE_SIZE)
+
+    def __init_subclass__(cls, **kwargs):
+        """Ensure that subclasses properly initialize the cache"""
+        super().__init_subclass__(**kwargs)
+
+        # Store the original __init__ method
+        original_init = cls.__init__
+
+        def new_init(self, *args, **kwargs):
+            # Always call CacheProvider.__init__ first
+            CacheProvider.__init__(self)
+            # Then call the original __init__ if it exists and is not CacheProvider's
+            if original_init is not CacheProvider.__init__:
+                original_init(self, *args, **kwargs)
+
+        cls.__init__ = new_init
 
     def get_or_update_cache(
         self,
@@ -53,64 +73,113 @@ class CacheProvider(ABC, Generic[T]):
         get_queries_fn: Callable[..., List[T]],
         *args,
         **kwargs,
-    ):
+    ) -> List[T]:
         if cache_path in self.cache:
-            return self.cache.get(cache_path)
+            cached_result = self.cache.get(cache_path)
+            return cached_result if cached_result is not None else []
         result = get_queries_fn(*args, **kwargs)
         self.cache.put(cache_path, result)
         return result
 
 
-class EmptySystemMetricsSource:
-    """Empty system metrics source that can be used as a default. Just returns an empty list of system metrics
-    for any resource."""
+class SystemMetricsComputer(Protocol):
+    """System metrics computer class to fetch system metrics for a given table."""
+
+    def _get_query_results(
+        self,
+        session: Session,
+        query,
+        operation,
+    ) -> List[QueryResult]:
+        """get query results either from cache or from the database
+
+        Args:
+            session (Session): session
+            query (_type_): query
+            operation (_type_): operation
+
+        Returns:
+            List[QueryResult]:
+        """
+        cursor = session.execute(text(query))
+        results = [
+            QueryResult(
+                database_name=row.database,
+                schema_name=row.schema,
+                table_name=row.table,
+                query_text=None,
+                query_type=operation,
+                start_time=row.starttime,
+                rows=row.rows,
+            )
+            for row in cursor
+            if (row.rows is not None) and (row.rows > 0)
+        ]
+
+        return results
+
+    def get_system_metrics(self) -> List[SystemProfile]:
+        """Return system metrics for a given table. Actual passed object can be a variety of types based
+        on the underlying infrastructure. For example, in the case of SQLalchemy, it can be a Table object
+        and in the case of Mongo, it can be a collection object."""
+        return self.get_inserts() + self.get_deletes() + self.get_updates()
 
-    def get_inserts(self, **kwargs) -> List[SystemProfile]:
+    def get_inserts(self) -> List[SystemProfile]:
         """Get insert queries"""
         return []
 
-    def get_deletes(self, **kwargs) -> List[SystemProfile]:
+    def get_deletes(self) -> List[SystemProfile]:
         """Get delete queries"""
         return []
 
-    def get_updates(self, **kwargs) -> List[SystemProfile]:
+    def get_updates(self) -> List[SystemProfile]:
         """Get update queries"""
         return []
 
-    def get_kwargs(self, **kwargs):
-        """Get kwargs to be used in get_inserts, get_deletes, get_updates"""
-        return {}
-
 
-class SystemMetricsComputer(EmptySystemMetricsSource):
-    def __init__(self, *args, **kwargs):
-        # collaborative constructor that initalizes upstream classes
-        super().__init__(*args, **kwargs)
-
-    def get_system_metrics(self, **kwargs) -> List[SystemProfile]:
-        """Return system metrics for a given table. Actual passed object can be a variety of types based
-        on the underlying infrastructure. For example, in the case of SQLalchemy, it can be a Table object
-        and in the case of Mongo, it can be a collection object."""
-        kwargs = super().get_kwargs(**kwargs)
-        return (
-            super().get_inserts(**kwargs)
-            + super().get_deletes(**kwargs)
-            + super().get_updates(**kwargs)
-        )
+class SystemMetricsRegistry:
+    _registry: Dict[str, Type["SystemMetricsComputer"]] = {}
 
+    @classmethod
+    def register(cls, dialect: PythonDialects, implementation: Type):
+        cls._registry[dialect.name.lower()] = implementation
 
-class SQASessionProvider:
-    """SQASessionProvider class to provide session to the system metrics"""
+    @classmethod
+    def get(cls, dialect: PythonDialects) -> Optional[Type["SystemMetricsComputer"]]:
+        if dialect.name.lower() not in cls._registry:
+            cls._discover_implementation(dialect)
+        return cls._registry.get(dialect.name.lower())
 
-    def __init__(self, *args, **kwargs):
-        self.session = kwargs.pop("session")
-        super().__init__(*args, **kwargs)
+    @classmethod
+    def _discover_implementation(cls, dialect: PythonDialects):
+        """Auto-discover the implementation in the profiler metrics"""
+        try:
+            implementation = import_from_module(
+                f"metadata.profiler.metrics.system.{dialect.name.lower()}.system"
+            )
+        except DynamicImportException:
+            logger.warning(f"No implementation found for {dialect.name.lower()}")
+            return
+        cls._registry[dialect.name.lower()] = implementation
+
+
+def register_system_metrics(
+    dialect: PythonDialects,
+) -> Callable[[Type["SystemMetricsComputer"]], Type["SystemMetricsComputer"]]:
+    """Decorator to register a system metric implementation
+
+    Args:
+        dialect (PythonDialects): database type
+
+    Returns:
+        Callable: decorator function
+    """
 
-    def get_session(self):
-        return self.session
+    def decorator(cls: Type["SystemMetricsComputer"]):
+        SystemMetricsRegistry.register(dialect, cls)
+        return cls
 
-    def get_database(self) -> str:
-        return self.session.get_bind().url.database
+    return decorator
 
 
 class System(SystemMetric):

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -73,13 +73,6 @@ class BigQuerySampler(SQASampler):
             **kwargs,
         )
         self.raw_dataset_type: Optional[TableType] = entity.tableType
-        if self._table.__table__.schema and self.entity.database.name:
-            self._table.__table__.schema = (
-                f"{self.entity.database.name}.{self._table.__table__.schema}"
-            )
-            self._table.__table_args__[
-                "schema"
-            ] = f"{self.entity.database.name}.{self._table.__table_args__['schema']}"
 
     def set_tablesample(self, selectable: SqaTable):
         """Set the TABLESAMPLE clause for BigQuery

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -11,16 +11,47 @@
 """
 Helper module to handle data sampling for the profiler
 """
+from sqlalchemy import Column, text
+
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
 )
+from metadata.profiler.orm.types.custom_array import CustomArray
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class DatabricksSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     def get_client(self):
         """client is the session for SQA"""
-        self.connection = databricks_get_connection(self.service_connection_config)
         client = super().get_client()
         self.set_catalog(client)
         return client
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        return isinstance(column.type, CustomArray)
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+
+        Args:
+            column_name: Name of the column
+            max_elements: Maximum number of elements to extract
+
+        Returns:
+            SQL expression string for array slicing
+        """
+        max_elements = self._get_max_array_elements()
+        return text(
+            f"""
+        CASE 
+            WHEN `{column.name}` IS NULL THEN NULL
+            ELSE slice(`{column.name}`, 1, {max_elements})
+        END AS `{column._label}`
+        """
+        )

metadata/sampler/sqlalchemy/sampler.py

@@ -13,6 +13,7 @@ Helper module to handle data sampling
 for the profiler
 """
 import hashlib
+from contextlib import contextmanager
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
@@ -40,6 +41,9 @@ logger = profiler_interface_registry_logger()
 
 RANDOM_LABEL = "random"
 
+# Default maximum number of elements to extract from array columns to prevent OOM
+DEFAULT_MAX_ARRAY_ELEMENTS = 10
+
 
 def _object_value_for_elem(self, elem):
     """
@@ -71,6 +75,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         self._table = self.build_table_orm(
             self.entity, self.service_connection_config, self.ometa_client
         )
+        self._active_sessions = set()
 
     @property
     def raw_dataset(self):
@@ -79,7 +84,20 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
     def get_client(self):
         """Build the SQA Client"""
         session_factory = create_and_bind_thread_safe_session(self.connection)
-        return session_factory()
+        session = session_factory()
+        self._active_sessions.add(session)
+        return session
+
+    @contextmanager
+    def get_client_context(self):
+        """Get client as context manager for proper cleanup"""
+        session = self.get_client()
+        try:
+            yield session
+        finally:
+            if session in self._active_sessions:
+                self._active_sessions.remove(session)
+            session.close()
 
     def set_tablesample(self, selectable: Table):
         """Set the tablesample for the table. To be implemented by the child SQA sampler class
@@ -88,6 +106,36 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         """
         return selectable
 
+    def _get_max_array_elements(self) -> int:
+        """Get the maximum number of array elements from config or use default"""
+        if (
+            self.sample_config
+            and hasattr(self.sample_config, "maxArrayElements")
+            and self.sample_config.maxArrayElements
+        ):
+            return self.sample_config.maxArrayElements
+        return DEFAULT_MAX_ARRAY_ELEMENTS
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        # Implement this method in the child classes
+        return False
+
+    def _process_array_value(self, value):
+        """Process array values to convert numpy arrays to Python lists"""
+        import numpy as np  # pylint: disable=import-outside-toplevel
+
+        if isinstance(value, np.ndarray):
+            return value.tolist()
+        return value
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+        By default, we return the column as is.
+        Child classes can override this method to return a different expression.
+        """
+        return column
+
     def _base_sample_query(self, column: Optional[Column], label=None):
         """Base query for sampling
 
@@ -199,17 +247,50 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
+
+            # Handle array columns with special query modification
+            max_elements = self._get_max_array_elements()
+            select_columns = []
+            has_array_columns = False
+
+            for col in sqa_columns:
+                if self._handle_array_column(col):
+                    slice_expression = self._get_slice_expression(col)
+                    select_columns.append(slice_expression)
+                    logger.debug(
+                        f"Limiting array column {col.name} to {max_elements} elements to prevent OOM"
+                    )
+                    has_array_columns = True
+                else:
+                    select_columns.append(col)
+
+            # Create query with modified columns
             sqa_sample = (
-                client.query(*sqa_columns)
+                client.query(*select_columns)
                 .select_from(ds)
                 .limit(self.sample_limit)
                 .all()
             )
 
+        # Process array columns manually if we used text() expressions
+        processed_rows = []
+        if has_array_columns:
+            for row in sqa_sample:
+                processed_row = []
+                for i, col in enumerate(sqa_columns):
+                    value = row[i]
+                    if self._handle_array_column(col):
+                        processed_value = self._process_array_value(value)
+                        processed_row.append(processed_value)
+                    else:
+                        processed_row.append(value)
+                processed_rows.append(processed_row)
+        else:
+            processed_rows = [list(row) for row in sqa_sample]
         return TableData(
             columns=[column.name for column in sqa_columns],
-            rows=[list(row) for row in sqa_sample],
+            rows=processed_rows,
         )
 
     def _fetch_sample_data_from_user_query(self) -> TableData:
@@ -217,7 +298,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not is_safe_sql_query(self.sample_query):
             raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
             rnd = client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
@@ -264,5 +345,6 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def close(self):
         """Close the connection"""
-        self.get_client().close()
+        for session in self._active_sessions:
+            session.close()
         self.connection.pool.dispose()

metadata/sampler/sqlalchemy/unitycatalog/sampler.py

@@ -13,16 +13,48 @@
 Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
+
+from sqlalchemy import Column, text
+
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
 )
+from metadata.profiler.orm.types.custom_array import CustomArray
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class UnityCatalogSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     def get_client(self):
         """client is the session for SQA"""
-        self.connection = databricks_get_connection(self.service_connection_config)
         client = super().get_client()
         self.set_catalog(client)
         return client
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        return isinstance(column.type, CustomArray)
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+
+        Args:
+            column_name: Name of the column
+            max_elements: Maximum number of elements to extract
+
+        Returns:
+            SQL expression string for array slicing
+        """
+        max_elements = self._get_max_array_elements()
+        return text(
+            f"""
+        CASE 
+            WHEN `{column.name}` IS NULL THEN NULL
+            ELSE slice(`{column.name}`, 1, {max_elements})
+        END AS `{column._label}`
+        """
+        )

metadata/utils/filters.py

@@ -164,6 +164,21 @@ def filter_by_dashboard(
     return _filter(dashboard_filter_pattern, dashboard_name)
 
 
+def filter_by_stored_procedure(
+    stored_procedure_filter_pattern: Optional[FilterPattern], stored_procedure_name: str
+) -> bool:
+    """
+    Return True if the stored procedure needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param stored_procedure_filter_pattern: Model defining stored procedure filtering logic
+    :param stored_procedure_name: stored procedure name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(stored_procedure_filter_pattern, stored_procedure_name)
+
+
 def filter_by_fqn(fqn_filter_pattern: Optional[FilterPattern], fqn: str) -> bool:
     """
     Return True if the FQN needs to be filtered, False otherwise