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

metadata/profiler/metrics/system/databricks/system.py

@@ -0,0 +1,132 @@
+import textwrap
+from typing import List
+
+from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
+
+from metadata.generated.schema.entity.data.table import SystemProfile
+from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations
+from metadata.profiler.metrics.system.system import (
+    CacheProvider,
+    SystemMetricsComputer,
+    register_system_metrics,
+)
+from metadata.profiler.orm.registry import PythonDialects
+from metadata.profiler.processor.runner import QueryRunner
+from metadata.utils.profiler_utils import QueryResult
+from metadata.utils.time_utils import datetime_to_timestamp
+
+SYSTEM_QUERY = textwrap.dedent(
+    """
+    SELECT
+        timestamp AS starttime,
+        COALESCE(CAST({column1} AS BIGINT), 0) + COALESCE(CAST({column2} AS BIGINT), 0) AS rows,
+        '{database}' AS database,
+        '{schema}' AS schema,
+        '{table}' AS table
+    FROM (DESCRIBE HISTORY {database}.{schema}.{table})
+    WHERE operation IN ({operations}) AND timestamp > DATEADD(day, -1, CURRENT_TIMESTAMP())
+    """
+)
+
+
+@register_system_metrics(PythonDialects.Databricks)
+class DatabricksSystemMetricsComputer(SystemMetricsComputer, CacheProvider):
+    """Databricks system metrics computer"""
+
+    def __init__(self, session: Session, runner: QueryRunner, catalog: str):
+        self.session = session
+        self.table = runner.table_name
+        self.database = catalog
+        self.schema = runner.schema_name
+
+    def _get_metrics_from_queries(
+        self, ddls: List[QueryResult], operation: str
+    ) -> List[SystemProfile]:
+        return TypeAdapter(List[SystemProfile]).validate_python(
+            [
+                {
+                    "timestamp": datetime_to_timestamp(
+                        ddl.start_time, milliseconds=True
+                    ),
+                    "operation": operation,
+                    "rowsAffected": ddl.rows,
+                }
+                for ddl in ddls
+            ]
+        )
+
+    def get_inserts(self) -> List[SystemProfile]:
+        operations = ", ".join(
+            [
+                f"'{DatabaseDMLOperations.WRITE.value}'",
+                f"'{DatabaseDMLOperations.MERGE.value}'",
+            ]
+        )
+        queries = self.get_or_update_cache(
+            f"{self.database}.{self.schema}.{self.table}.{DatabaseDMLOperations.INSERT.value}",
+            self._get_query_results,
+            self.session,
+            SYSTEM_QUERY.format(
+                column1="operationMetrics.numOutputRows",
+                column2="operationMetrics.numTargetRowsInserted",
+                database=self.database,
+                schema=self.schema,
+                table=self.table,
+                operations=operations,
+            ),
+            DatabaseDMLOperations.INSERT.value,
+        )
+        return self._get_metrics_from_queries(
+            queries, DatabaseDMLOperations.INSERT.value
+        )
+
+    def get_deletes(self) -> List[SystemProfile]:
+        operations = ", ".join(
+            [
+                f"'{DatabaseDMLOperations.DELETE.value}'",
+                f"'{DatabaseDMLOperations.MERGE.value}'",
+            ]
+        )
+        queries = self.get_or_update_cache(
+            f"{self.database}.{self.schema}.{self.table}.{DatabaseDMLOperations.DELETE.value}",
+            self._get_query_results,
+            self.session,
+            SYSTEM_QUERY.format(
+                column1="operationMetrics.numDeletedRows",
+                column2="operationMetrics.numTargetRowsDeleted",
+                database=self.database,
+                schema=self.schema,
+                table=self.table,
+                operations=operations,
+            ),
+            DatabaseDMLOperations.DELETE.value,
+        )
+        return self._get_metrics_from_queries(
+            queries, DatabaseDMLOperations.DELETE.value
+        )
+
+    def get_updates(self) -> List[SystemProfile]:
+        operations = ", ".join(
+            [
+                f"'{DatabaseDMLOperations.UPDATE.value}'",
+                f"'{DatabaseDMLOperations.MERGE.value}'",
+            ]
+        )
+        queries = self.get_or_update_cache(
+            f"{self.database}.{self.schema}.{self.table}.{DatabaseDMLOperations.UPDATE.value}",
+            self._get_query_results,
+            self.session,
+            SYSTEM_QUERY.format(
+                column1="operationMetrics.numUpdatedRows",
+                column2="operationMetrics.numTargetRowsUpdated",
+                database=self.database,
+                schema=self.schema,
+                table=self.table,
+                operations=operations,
+            ),
+            DatabaseDMLOperations.UPDATE.value,
+        )
+        return self._get_metrics_from_queries(
+            queries, DatabaseDMLOperations.UPDATE.value
+        )

metadata/profiler/metrics/system/dml_operation.py

@@ -25,6 +25,7 @@ class DatabaseDMLOperations(Enum):
     UPDATE = "UPDATE"
     DELETE = "DELETE"
     MERGE = "MERGE"
+    WRITE = "WRITE"
 
 
 DML_OPERATION_MAP = {
@@ -32,4 +33,5 @@ DML_OPERATION_MAP = {
     DatabaseDMLOperations.MERGE.value: DmlOperationType.UPDATE.value,
     DatabaseDMLOperations.UPDATE.value: DmlOperationType.UPDATE.value,
     DatabaseDMLOperations.DELETE.value: DmlOperationType.DELETE.value,
+    DatabaseDMLOperations.WRITE.value: DmlOperationType.WRITE.value,
 }

metadata/{ingestion/source/database/redshift/profiler → profiler/metrics/system/redshift}/system.py

@@ -1,23 +1,21 @@
 """
-Imeplemetation for the redshift system metrics source
+Implemetation for the redshift system metrics source
 """
 
 from typing import List
 
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 
 from metadata.generated.schema.entity.data.table import SystemProfile
-from metadata.ingestion.source.database.redshift.queries import (
-    STL_QUERY,
-    get_query_results,
-)
+from metadata.ingestion.source.database.redshift.queries import STL_QUERY
 from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations
 from metadata.profiler.metrics.system.system import (
     CacheProvider,
-    EmptySystemMetricsSource,
-    SQASessionProvider,
     SystemMetricsComputer,
+    register_system_metrics,
 )
+from metadata.profiler.orm.registry import PythonDialects
 from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.logger import profiler_logger
 from metadata.utils.profiler_utils import QueryResult
@@ -26,58 +24,46 @@ from metadata.utils.time_utils import datetime_to_timestamp
 logger = profiler_logger()
 
 
-class RedshiftSystemMetricsSource(
-    SQASessionProvider, EmptySystemMetricsSource, CacheProvider
-):
+@register_system_metrics(PythonDialects.Redshift)
+class RedshiftSystemMetricsComputer(SystemMetricsComputer, CacheProvider):
     """Redshift system metrics source class"""
 
-    def get_inserts(self, **kwargs) -> List[SystemProfile]:
-        database, schema, table = (
-            kwargs.get("database"),
-            kwargs.get("schema"),
-            kwargs.get("table"),
-        )
+    def __init__(
+        self,
+        session: Session,
+        runner: QueryRunner,
+    ):
+        self.session = session
+        self.table = runner.table_name
+        self.database = runner.session.get_bind().url.database
+        self.schema = runner.schema_name
+
+    def get_inserts(self) -> List[SystemProfile]:
         queries = self.get_or_update_cache(
-            f"{database}.{schema}",
+            f"{self.database}.{self.schema}.{DatabaseDMLOperations.INSERT.value}",
             self._get_insert_queries,
-            database=database,
-            schema=schema,
-        )
-        return get_metric_result(queries, table)
-
-    def get_kwargs(self, **kwargs):
-        runner: QueryRunner = kwargs.get("runner")
-        return {
-            "table": runner.table_name,
-            "database": runner.session.get_bind().url.database,
-            "schema": runner.schema_name,
-        }
-
-    def get_deletes(self, **kwargs) -> List[SystemProfile]:
-        database, schema, table = (
-            kwargs.get("database"),
-            kwargs.get("schema"),
-            kwargs.get("table"),
+            database=self.database,
+            schema=self.schema,
         )
+        return get_metric_result(queries, self.table)
+
+    def get_deletes(self) -> List[SystemProfile]:
         queries = self.get_or_update_cache(
-            f"{database}.{schema}",
+            f"{self.database}.{self.schema}.{DatabaseDMLOperations.DELETE.value}",
             self._get_delete_queries,
-            database=database,
-            schema=schema,
+            database=self.database,
+            schema=self.schema,
         )
-        return get_metric_result(queries, table)
+        return get_metric_result(queries, self.table)
 
-    def get_updates(self, **kwargs) -> List[SystemProfile]:
-        database = kwargs.get("database")
-        schema = kwargs.get("schema")
-        table = kwargs.get("table")
+    def get_updates(self) -> List[SystemProfile]:
         queries = self.get_or_update_cache(
-            f"{database}.{schema}",
+            f"{self.database}.{self.schema}.{DatabaseDMLOperations.UPDATE.value}",
             self._get_update_queries,
-            database=database,
-            schema=schema,
+            database=self.database,
+            schema=self.schema,
         )
-        return get_metric_result(queries, table)
+        return get_metric_result(queries, self.table)
 
     def _get_insert_queries(self, database: str, schema: str) -> List[QueryResult]:
         insert_query = STL_QUERY.format(
@@ -87,8 +73,8 @@ class RedshiftSystemMetricsSource(
             database=database,
             schema=schema,
         )
-        return get_query_results(
-            super().get_session(),
+        return self._get_query_results(
+            self.session,
             insert_query,
             DatabaseDMLOperations.INSERT.value,
         )
@@ -101,8 +87,8 @@ class RedshiftSystemMetricsSource(
             database=database,
             schema=schema,
         )
-        return get_query_results(
-            super().get_session(),
+        return self._get_query_results(
+            self.session,
             delete_query,
             DatabaseDMLOperations.DELETE.value,
         )
@@ -115,8 +101,8 @@ class RedshiftSystemMetricsSource(
             database=database,
             schema=schema,
         )
-        return get_query_results(
-            super().get_session(),
+        return self._get_query_results(
+            self.session,
             update_query,
             DatabaseDMLOperations.UPDATE.value,
         )
@@ -143,7 +129,3 @@ def get_metric_result(ddls: List[QueryResult], table_name: str) -> List[SystemPr
             if ddl.table_name == table_name
         ]
     )
-
-
-class RedshiftSystemMetricsComputer(SystemMetricsComputer, RedshiftSystemMetricsSource):
-    pass

metadata/{ingestion/source/database/snowflake/profiler → profiler/metrics/system/snowflake}/system.py

@@ -7,6 +7,7 @@ from typing import List, Optional, Tuple
 
 import sqlalchemy.orm
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 
 from metadata.generated.schema.entity.data.table import DmlOperationType, SystemProfile
 from metadata.ingestion.source.database.snowflake.models import (
@@ -16,10 +17,10 @@ from metadata.ingestion.source.database.snowflake.models import (
 from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations
 from metadata.profiler.metrics.system.system import (
     CacheProvider,
-    EmptySystemMetricsSource,
-    SQASessionProvider,
     SystemMetricsComputer,
+    register_system_metrics,
 )
+from metadata.profiler.orm.registry import PythonDialects
 from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.collections import CaseInsensitiveString
 from metadata.utils.logger import profiler_logger
@@ -283,38 +284,30 @@ def get_snowflake_system_queries(
     return None
 
 
-class SnowflakeSystemMetricsSource(
-    SQASessionProvider, EmptySystemMetricsSource, CacheProvider[SnowflakeQueryLogEntry]
+@register_system_metrics(PythonDialects.Snowflake)
+class SnowflakeSystemMetricsComputer(
+    SystemMetricsComputer, CacheProvider[SnowflakeQueryLogEntry]
 ):
     """Snowflake system metrics source"""
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, session: Session, runner: QueryRunner):
+        self.session = session
+        self.runner = runner
+        self.table = runner.table_name
+        self.database = runner.session.get_bind().url.database
+        self.schema = runner.schema_name
         self.resolver = SnowflakeTableResovler(
-            session=super().get_session(),
+            session=session,
         )
 
-    def get_kwargs(self, **kwargs):
-        runner: QueryRunner = kwargs.get("runner")
-        return {
-            "table": runner.table_name,
-            "database": runner.session.get_bind().url.database,
-            "schema": runner.schema_name,
-        }
-
-    def get_inserts(self, **kwargs) -> List[SystemProfile]:
-        database, schema, table = (
-            kwargs.get("database"),
-            kwargs.get("schema"),
-            kwargs.get("table"),
-        )
+    def get_inserts(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            database,
-            schema,
-            table,
+            self.database,
+            self.schema,
+            self.table,
             list(
                 self.get_queries_by_operation(
-                    table,
+                    self.table,
                     [
                         DatabaseDMLOperations.INSERT,
                         DatabaseDMLOperations.MERGE,
@@ -325,19 +318,14 @@ class SnowflakeSystemMetricsSource(
             DmlOperationType.INSERT,
         )
 
-    def get_updates(self, **kwargs) -> List[SystemProfile]:
-        database, schema, table = (
-            kwargs.get("database"),
-            kwargs.get("schema"),
-            kwargs.get("table"),
-        )
+    def get_updates(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            database,
-            schema,
-            table,
+            self.database,
+            self.schema,
+            self.table,
             list(
                 self.get_queries_by_operation(
-                    table,
+                    self.table,
                     [
                         DatabaseDMLOperations.UPDATE,
                         DatabaseDMLOperations.MERGE,
@@ -348,19 +336,14 @@ class SnowflakeSystemMetricsSource(
             DmlOperationType.UPDATE,
         )
 
-    def get_deletes(self, **kwargs) -> List[SystemProfile]:
-        database, schema, table = (
-            kwargs.get("database"),
-            kwargs.get("schema"),
-            kwargs.get("table"),
-        )
+    def get_deletes(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            database,
-            schema,
-            table,
+            self.database,
+            self.schema,
+            self.table,
             list(
                 self.get_queries_by_operation(
-                    table,
+                    self.table,
                     [
                         DatabaseDMLOperations.DELETE,
                     ],
@@ -418,7 +401,7 @@ class SnowflakeSystemMetricsSource(
         queries = self.get_or_update_cache(
             table,
             SnowflakeQueryLogEntry.get_for_table,
-            session=super().get_session(),
+            session=self.session,
             tablename=table,
         )
         results = [
@@ -429,9 +412,3 @@ class SnowflakeSystemMetricsSource(
             for row in queries
         ]
         return [result for result in results if result is not None]
-
-
-class SnowflakeSystemMetricsComputer(
-    SystemMetricsComputer, SnowflakeSystemMetricsSource
-):
-    pass

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