openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/profiler/interface/pandas/profiler_interface.py

@@ -22,17 +22,22 @@ from typing import Dict, List, Optional
 
 from sqlalchemy import Column
 
-from metadata.generated.schema.entity.data.table import TableData
+from metadata.generated.schema.entity.data.table import CustomMetricProfile, TableData
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.mixins.pandas.pandas_mixin import PandasInterfaceMixin
+from metadata.profiler.api.models import ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR, SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.datalake.datalake_utils import fetch_col_types, fetch_dataframe
+from metadata.utils.datalake.datalake_utils import (
+    GenericDataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -239,35 +244,67 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         """
         return None  # to be implemented
 
+    def _compute_custom_metrics(
+        self, metrics: List[CustomMetric], runner, *args, **kwargs
+    ):
+        """Compute custom metrics. For pandas source we expect expression
+        to be a boolean value. We'll return the length of the dataframe
+
+        Args:
+            metrics (List[Metrics]): list of customMetrics
+            runner (_type_): runner
+        """
+        if not metrics:
+            return None
+
+        custom_metrics = []
+
+        for metric in metrics:
+            try:
+                row = sum(
+                    len(df.query(metric.expression).index)
+                    for df in runner
+                    if len(df.query(metric.expression).index)
+                )
+                custom_metrics.append(
+                    CustomMetricProfile(name=metric.name.__root__, value=row)
+                )
+
+            except Exception as exc:
+                msg = f"Error trying to compute profile for custom metric: {exc}"
+                logger.debug(traceback.format_exc())
+                logger.warning(msg)
+        if custom_metrics:
+            return {"customMetrics": custom_metrics}
+        return None
+
     def compute_metrics(
         self,
-        metrics,
-        metric_type,
-        column,
-        table,
+        metric_func: ThreadPoolMetrics,
     ):
         """Run metrics in processor worker"""
-        logger.debug(f"Running profiler for {table}")
+        logger.debug(f"Running profiler for {metric_func.table}")
         try:
             row = None
             if self.complex_dataframe_sample:
-                row = self._get_metric_fn[metric_type.value](
-                    metrics,
+                row = self._get_metric_fn[metric_func.metric_type.value](
+                    metric_func.metrics,
                     self.complex_dataframe_sample,
-                    column=column,
+                    column=metric_func.column,
                 )
         except Exception as exc:
-            name = f"{column if column is not None else table}"
+            name = f"{metric_func.column if metric_func.column is not None else metric_func.table}"
             error = f"{name} metric_type.value: {exc}"
             logger.error(error)
             self.status.failed_profiler(error, traceback.format_exc())
             row = None
-        if column is not None:
-            column = column.name
-            self.status.scanned(f"{table.name.__root__}.{column}")
+        if metric_func.column is not None:
+            column = metric_func.column.name
+            self.status.scanned(f"{metric_func.table.name.__root__}.{column}")
         else:
-            self.status.scanned(table.name.__root__)
-        return row, column, metric_type.value
+            self.status.scanned(metric_func.table.name.__root__)
+            column = None
+        return row, column, metric_func.metric_type.value
 
     def fetch_sample_data(self, table, columns: SQALikeColumn) -> TableData:
         """Fetch sample data from database
@@ -329,7 +366,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
         profile_results = {"table": {}, "columns": defaultdict(dict)}
         metric_list = [
-            self.compute_metrics(*metric_func) for metric_func in metric_funcs
+            self.compute_metrics(metric_func) for metric_func in metric_funcs
         ]
         for metric_result in metric_list:
             profile, column, metric_type = metric_result
@@ -338,6 +375,8 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                     profile_results["table"].update(profile)
                 if metric_type == MetricTypes.System.value:
                     profile_results["system"] = profile
+                elif metric_type == MetricTypes.Custom.value and column is None:
+                    profile_results["table"].update(profile)
                 else:
                     if profile:
                         profile_results["columns"][column].update(
@@ -375,7 +414,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                 sqalike_columns.append(
                     SQALikeColumn(
                         column_name,
-                        fetch_col_types(self.complex_dataframe_sample[0], column_name),
+                        GenericDataFrameColumnParser.fetch_col_types(
+                            self.complex_dataframe_sample[0], column_name
+                        ),
                     )
                 )
             return sqalike_columns

metadata/profiler/interface/profiler_interface.py

@@ -18,7 +18,6 @@ from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Union
 
 from sqlalchemy import Column
-from typing_extensions import Self
 
 from metadata.generated.schema.entity.data.database import (
     Database,
@@ -43,10 +42,13 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseService,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
-from metadata.ingestion.api.models import StackTraceError
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
@@ -76,7 +78,7 @@ class ProfilerProcessorStatus(Status):
             StackTraceError(
                 name=self.entity if self.entity else "",
                 error=error,
-                stack_trace=stack_trace,
+                stackTrace=stack_trace,
             )
         )
 
@@ -130,6 +132,7 @@ class ProfilerInterface(ABC):
             MetricTypes.Query.value: self._compute_query_metrics,
             MetricTypes.Window.value: self._compute_window_metrics,
             MetricTypes.System.value: self._compute_system_metrics,
+            MetricTypes.Custom.value: self._compute_custom_metrics,
         }
 
     @abstractmethod
@@ -151,7 +154,7 @@ class ProfilerInterface(ABC):
         service_connection_config,
         ometa_client: Optional[OpenMetadata],
         **kwargs,
-    ) -> Self:
+    ) -> "ProfilerInterface":
         """create class method is used to dispatch the profiler protocol to the
         correct object based on the service connection object
 
@@ -459,6 +462,13 @@ class ProfilerInterface(ABC):
         """Get metrics"""
         raise NotImplementedError
 
+    @abstractmethod
+    def _compute_custom_metrics(
+        self, metrics: List[CustomMetric], runner, *args, **kwargs
+    ):
+        """Compute custom metrics"""
+        raise NotImplementedError
+
     @abstractmethod
     def get_all_metrics(self, metric_funcs) -> dict:
         """run profiler metrics"""

metadata/profiler/interface/profiler_interface_factory.py

@@ -18,9 +18,18 @@ from typing import cast
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Connection,
+)
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
     SingleStoreConnection,
 )
@@ -30,6 +39,9 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
     TrinoConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.profiler.interface.pandas.profiler_interface import (
     PandasProfilerInterface,
@@ -38,6 +50,15 @@ from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import (
     BigQueryProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import (
+    DatabricksProfilerInterface,
+)
+from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
+    DB2ProfilerInterface,
+)
+from metadata.profiler.interface.sqlalchemy.mariadb.profiler_interface import (
+    MariaDBProfilerInterface,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
@@ -50,6 +71,9 @@ from metadata.profiler.interface.sqlalchemy.snowflake.profiler_interface import
 from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import (
     TrinoProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
+    UnityCatalogProfilerInterface,
+)
 
 
 class ProfilerInterfaceFactory:
@@ -62,6 +86,17 @@ class ProfilerInterfaceFactory:
         """Register a new interface"""
         self._interface_type[interface_type] = interface_class
 
+    def register_many(self, interface_dict):
+        """
+        Registers multiple profiler interfaces at once.
+
+        Args:
+            interface_dict: A dictionary mapping connection class names (strings) to their
+            corresponding profiler interface classes.
+        """
+        for interface_type, interface_class in interface_dict.items():
+            self.register(interface_type, interface_class)
+
     def create(self, interface_type: str, *args, **kwargs):
         """Create interface object based on interface type"""
         interface_class = self._interface_type.get(interface_type)
@@ -72,17 +107,17 @@ class ProfilerInterfaceFactory:
 
 
 profiler_interface_factory = ProfilerInterfaceFactory()
-profiler_interface_factory.register(DatabaseConnection.__name__, SQAProfilerInterface)
-profiler_interface_factory.register(
-    BigQueryConnection.__name__, BigQueryProfilerInterface
-)
-profiler_interface_factory.register(
-    SingleStoreConnection.__name__, SingleStoreProfilerInterface
-)
-profiler_interface_factory.register(
-    DatalakeConnection.__name__, PandasProfilerInterface
-)
-profiler_interface_factory.register(
-    SnowflakeConnection.__name__, SnowflakeProfilerInterface
-)
-profiler_interface_factory.register(TrinoConnection.__name__, TrinoProfilerInterface)
+profilers = {
+    DatabaseConnection.__name__: SQAProfilerInterface,
+    BigQueryConnection.__name__: BigQueryProfilerInterface,
+    SingleStoreConnection.__name__: SingleStoreProfilerInterface,
+    DatalakeConnection.__name__: PandasProfilerInterface,
+    MariaDBConnection.__name__: MariaDBProfilerInterface,
+    SnowflakeConnection.__name__: SnowflakeProfilerInterface,
+    TrinoConnection.__name__: TrinoProfilerInterface,
+    UnityCatalogConnection.__name__: UnityCatalogProfilerInterface,
+    DatabricksConnection.__name__: DatabricksProfilerInterface,
+    Db2Connection.__name__: DB2ProfilerInterface,
+}
+
+profiler_interface_factory.register_many(profilers)

metadata/profiler/interface/sqlalchemy/bigquery/profiler_interface.py

@@ -32,6 +32,9 @@ class BigQueryProfilerInterface(SQAProfilerInterface):
         for key, value in columns.items():
             if not isinstance(value, STRUCT):
                 col = Column(f"{parent}.{key}", value)
+                # pylint: disable=protected-access
+                col._set_parent(self.table.__table__)
+                # pylint: enable=protected-access
                 columns_list.append(col)
             else:
                 col = self._get_struct_columns(

metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py

@@ -0,0 +1,26 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+
+
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+
+
+class DatabricksProfilerInterface(SQAProfilerInterface):
+    def __init__(self, service_connection_config, **kwargs):
+        super().__init__(service_connection_config=service_connection_config, **kwargs)
+        self.set_catalog(self.session)

metadata/profiler/interface/sqlalchemy/db2/profiler_interface.py

@@ -0,0 +1,38 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+
+
+class DB2ProfilerInterface(SQAProfilerInterface):
+    """
+    Interface to interact with registry supporting
+    sqlalchemy.
+    """
+
+    def _programming_error_static_metric(self, runner, column, exc, session, metrics):
+        # pylint: disable=protected-access
+        if exc.orig and "overflow" in exc.orig._message:
+            logger.info(
+                f"Computing metrics without sum for {runner.table.__tablename__}.{column.name}"
+            )
+            return self._compute_static_metrics_wo_sum(metrics, runner, session, column)
+        return None

metadata/profiler/interface/sqlalchemy/mariadb/profiler_interface.py

@@ -0,0 +1,85 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+
+from typing import List
+
+from sqlalchemy.exc import ProgrammingError
+
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+    handle_query_exception,
+)
+from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.processor.runner import QueryRunner
+from metadata.profiler.source.mariadb.metrics.window.first_quartile import (
+    MariaDBFirstQuartile,
+)
+from metadata.profiler.source.mariadb.metrics.window.median import MariaDBMedian
+from metadata.profiler.source.mariadb.metrics.window.third_quartile import (
+    MariaDBThirdQuartile,
+)
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+
+
+class MariaDBProfilerInterface(SQAProfilerInterface):
+    """
+    Interface to interact with registry supporting
+    sqlalchemy.
+    """
+
+    def _compute_window_metrics(
+        self,
+        metrics: List[Metrics],
+        runner: QueryRunner,
+        *args,
+        **kwargs,
+    ):
+        """Given a list of metrics, compute the given results
+        and returns the values
+
+        Args:
+            column: the column to compute the metrics against
+            metrics: list of metrics to compute
+        Returns:
+            dictionnary of results
+        """
+        session = kwargs.get("session")
+        column = kwargs.get("column")
+
+        if not metrics:
+            return None
+
+        try:
+            # we patch the metrics at runtime to use the MariaDB specific functions
+            # as we can't compile the query based on the dialect as it return `mysql`
+            metrics = [MariaDBFirstQuartile, MariaDBMedian, MariaDBThirdQuartile]  # type: ignore
+            row = runner.select_first_from_sample(
+                *[metric(column).fn() for metric in metrics],
+            )
+            if row:
+                return dict(row)
+        except ProgrammingError:
+            logger.info(
+                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to overflow"
+            )
+            return None
+
+        except Exception as exc:
+            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            handle_query_exception(msg, exc, session)
+        return None