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

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

@@ -3,16 +3,17 @@
 from typing import List
 
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 
 from metadata.generated.schema.entity.data.table import DmlOperationType, SystemProfile
 from metadata.ingestion.source.database.bigquery.queries import BigQueryQueryResult
 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.time_utils import datetime_to_timestamp
@@ -20,36 +21,32 @@ from metadata.utils.time_utils import datetime_to_timestamp
 logger = profiler_logger()
 
 
-class BigQuerySystemMetricsSource(
-    SQASessionProvider, EmptySystemMetricsSource, CacheProvider
-):
+@register_system_metrics(PythonDialects.BigQuery)
+class BigQuerySystemMetricsComputer(SystemMetricsComputer, CacheProvider):
     """BigQuery system metrics source class"""
 
-    def get_kwargs(self, **kwargs):
-        runner: QueryRunner = kwargs.get("runner")
-        return {
-            "table": runner.table_name,
-            "project_id": runner.session.get_bind().url.host,
-            "dataset_id": runner.schema_name,
-            "usage_location": kwargs.get("usage_location"),
-        }
+    def __init__(
+        self,
+        session: Session,
+        runner: QueryRunner,
+        usage_location: str,
+    ):
+        self.session = session
+        self.table = runner.table_name
+        self.project_id = runner.session.get_bind().url.host
+        self.dataset_id = runner.schema_name
+        self.usage_location = usage_location
 
-    def get_deletes(self, **kwargs) -> List[SystemProfile]:
-        table, project_id, usage_location, dataset_id = (
-            kwargs.get("table"),
-            kwargs.get("project_id"),
-            kwargs.get("usage_location"),
-            kwargs.get("dataset_id"),
-        )
+    def get_deletes(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            project_id,
-            dataset_id,
-            table,
+            self.project_id,
+            self.dataset_id,
+            self.table,
             list(
                 self.get_queries_by_operation(
-                    usage_location,
-                    project_id,
-                    dataset_id,
+                    self.usage_location,
+                    self.project_id,
+                    self.dataset_id,
                     [
                         DatabaseDMLOperations.DELETE,
                     ],
@@ -59,21 +56,15 @@ class BigQuerySystemMetricsSource(
             DmlOperationType.DELETE,
         )
 
-    def get_updates(self, **kwargs) -> List[SystemProfile]:
-        table, project_id, usage_location, dataset_id = (
-            kwargs.get("table"),
-            kwargs.get("project_id"),
-            kwargs.get("usage_location"),
-            kwargs.get("dataset_id"),
-        )
+    def get_updates(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            project_id,
-            dataset_id,
-            table,
+            self.project_id,
+            self.dataset_id,
+            self.table,
             self.get_queries_by_operation(
-                usage_location,
-                project_id,
-                dataset_id,
+                self.usage_location,
+                self.project_id,
+                self.dataset_id,
                 [
                     DatabaseDMLOperations.UPDATE,
                     DatabaseDMLOperations.MERGE,
@@ -83,21 +74,15 @@ class BigQuerySystemMetricsSource(
             DmlOperationType.UPDATE,
         )
 
-    def get_inserts(self, **kwargs) -> List[SystemProfile]:
-        table, project_id, usage_location, dataset_id = (
-            kwargs.get("table"),
-            kwargs.get("project_id"),
-            kwargs.get("usage_location"),
-            kwargs.get("dataset_id"),
-        )
+    def get_inserts(self) -> List[SystemProfile]:
         return self.get_system_profile(
-            project_id,
-            dataset_id,
-            table,
+            self.project_id,
+            self.dataset_id,
+            self.table,
             self.get_queries_by_operation(
-                usage_location,
-                project_id,
-                dataset_id,
+                self.usage_location,
+                self.project_id,
+                self.dataset_id,
                 [
                     DatabaseDMLOperations.INSERT,
                     DatabaseDMLOperations.MERGE,
@@ -127,7 +112,7 @@ class BigQuerySystemMetricsSource(
         return self.get_or_update_cache(
             f"{project_id}.{dataset_id}",
             BigQueryQueryResult.get_for_table,
-            session=super().get_session(),
+            session=self.session,
             usage_location=usage_location,
             project_id=project_id,
             dataset_id=dataset_id,
@@ -154,13 +139,10 @@ class BigQuerySystemMetricsSource(
                     "rowsAffected": getattr(q, rows_affected_field),
                 }
                 for q in query_results
-                if getattr(q, rows_affected_field, -1) > 0
+                if getattr(q, rows_affected_field)
+                or -1 > 0
                 and q.project_id == project_id
                 and q.dataset_id == dataset_id
                 and q.table_name == table
             ]
         )
-
-
-class BigQuerySystemMetricsComputer(SystemMetricsComputer, BigQuerySystemMetricsSource):
-    pass

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