openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -21,28 +21,41 @@ import threading
 import traceback
 from collections import defaultdict
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Type, Union
 
 from sqlalchemy import Column, inspect, text
 from sqlalchemy.exc import DBAPIError, ProgrammingError, ResourceClosedError
 from sqlalchemy.orm import scoped_session
 
-from metadata.generated.schema.entity.data.table import CustomMetricProfile, TableData
+from metadata.generated.schema.entity.data.table import (
+    CustomMetricProfile,
+    SystemProfile,
+    Table,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    DatabaseServiceProfilerPipeline,
+)
 from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.mixins.sqalchemy.sqa_mixin import SQAInterfaceMixin
 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.core import HybridMetric, MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.metrics.static.mean import Mean
 from metadata.profiler.metrics.static.stddev import StdDev
 from metadata.profiler.metrics.static.sum import Sum
+from metadata.profiler.metrics.system.system import System, SystemMetricsComputer
 from metadata.profiler.orm.functions.table_metric_computer import TableMetricComputer
 from metadata.profiler.orm.registry import Dialects
 from metadata.profiler.processor.metric_filter import MetricFilter
 from metadata.profiler.processor.runner import QueryRunner
-from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.custom_thread_pool import CustomThreadPoolExecutor
 from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
@@ -70,41 +83,40 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
     """
 
     # pylint: disable=too-many-arguments
-
     def __init__(
         self,
-        service_connection_config,
-        ometa_client,
-        entity,
-        storage_config,
-        profile_sample_config,
-        source_config,
-        sample_query,
-        table_partition_config,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        source_config: DatabaseServiceProfilerPipeline,
+        sampler: SamplerInterface,
         thread_count: int = 5,
         timeout_seconds: int = 43200,
-        sqa_metadata=None,
-        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
         **kwargs,
     ):
         """Instantiate SQA Interface object"""
+        self.session_factory = None
+        self.session = None
 
         super().__init__(
-            service_connection_config,
-            ometa_client,
-            entity,
-            storage_config,
-            profile_sample_config,
-            source_config,
-            sample_query,
-            table_partition_config,
-            thread_count,
-            timeout_seconds,
-            sample_data_count,
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            source_config=source_config,
+            sampler=sampler,
+            thread_count=thread_count,
+            timeout_seconds=timeout_seconds,
         )
 
-        self._table = self._convert_table_to_orm_object(sqa_metadata)
+        self._table = self.sampler.raw_dataset
         self.create_session()
+        self.system_metrics_computer = self.initialize_system_metrics_computer()
+
+    def initialize_system_metrics_computer(self) -> SystemMetricsComputer:
+        """Initialize system metrics computer. Override this if you want to use a metric source with
+        state or other dependencies.
+        """
+        return SystemMetricsComputer()
 
     def create_session(self):
         self.session_factory = self._session_factory()
@@ -114,25 +126,6 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
     def table(self):
         return self._table
 
-    def _get_sampler(self, **kwargs):
-        """get sampler object"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        session = kwargs.get("session")
-        table = kwargs["table"]
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            client=session or self.session,
-            table=table,
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-            sample_data_count=self.sample_data_count,
-        )
-
     def _session_factory(self) -> scoped_session:
         """Create thread safe session that will be automatically
         garbage collected once the application thread ends
@@ -163,7 +156,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
             )
             return dict(row)
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -174,7 +167,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         session,
         *args,
         **kwargs,
-    ):
+    ) -> Optional[Dict[str, Any]]:
         """Given a list of metrics, compute the given results
         and returns the values
 
@@ -201,7 +194,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to compute profile for {runner.table.__tablename__}: {exc}"  # type: ignore
+                f"Error trying to compute profile for {runner.table_name}: {exc}"  # type: ignore
             )
             session.rollback()
             raise RuntimeError(exc)
@@ -238,7 +231,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 runner, column, exc, session, metrics
             )
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -281,10 +274,10 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 runner._session.get_bind().dialect.name
                 != Dialects.Druid
             ):
-                msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+                msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
                 handle_query_exception(msg, exc, session)
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -317,10 +310,10 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 return dict(row)
         except ProgrammingError as exc:
             logger.info(
-                f"Skipping metrics for {runner.table.__tablename__}.{column.name} due to {exc}"
+                f"Skipping metrics for {runner.table_name}.{column.name} due to {exc}"
             )
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None
 
@@ -354,7 +347,7 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                 )
 
             except Exception as exc:
-                msg = f"Error trying to compute profile for {runner.table.__tablename__}.{metric.columnName}: {exc}"
+                msg = f"Error trying to compute profile for {runner.table_name}.{metric.columnName}: {exc}"
                 logger.debug(traceback.format_exc())
                 logger.warning(msg)
         if custom_metrics:
@@ -363,12 +356,11 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
 
     def _compute_system_metrics(
         self,
-        metrics: Metrics,
+        metrics: Type[System],
         runner: QueryRunner,
-        session,
         *args,
         **kwargs,
-    ):
+    ) -> List[SystemProfile]:
         """Get system metric for tables
 
         Args:
@@ -379,44 +371,21 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         Returns:
             dictionnary of results
         """
-        try:
-            rows = metrics().sql(session, conn_config=self.service_connection_config)
-            return rows
-        except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}: {exc}"
-            handle_query_exception(msg, exc, session)
-        return None
+        logger.debug(f"Computing system metrics for {runner.table_name}")
+        return self.system_metrics_computer.get_system_metrics(runner=runner)
 
-    def _create_thread_safe_sampler(
-        self,
-        session,
-        table,
-    ):
-        """Create thread safe runner"""
-        if not hasattr(thread_local, "sampler"):
-            thread_local.sampler = self._get_sampler(
-                table=table,
-                session=session,
-            )
-        return thread_local.sampler
-
-    def _create_thread_safe_runner(
-        self,
-        session,
-        table,
-        sample,
-    ):
+    def _create_thread_safe_runner(self, session, column=None):
         """Create thread safe runner"""
         if not hasattr(thread_local, "runner"):
             thread_local.runner = QueryRunner(
                 session=session,
-                table=table,
-                sample=sample,
-                partition_details=self.partition_details,
-                profile_sample_query=self.profile_query,
+                dataset=self.sampler.get_dataset(column=column),
+                raw_dataset=self.sampler.raw_dataset,
+                partition_details=self.sampler.partition_details,
+                profile_sample_query=self.sampler.sample_query,
             )
             return thread_local.runner
-        thread_local.runner._sample = sample  # pylint: disable=protected-access
+        thread_local.runner.dataset = self.sampler.get_dataset(column=column)
         return thread_local.runner
 
     def compute_metrics_in_thread(
@@ -431,31 +400,19 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
         with Session() as session:
             self.set_session_tag(session)
             self.set_catalog(session)
-            sampler = self._create_thread_safe_sampler(
-                session,
-                metric_func.table,
-            )
-            sample = sampler.random_sample(metric_func.column)
-            runner = self._create_thread_safe_runner(
-                session,
-                metric_func.table,
-                sample,
-            )
+            runner = self._create_thread_safe_runner(session, metric_func.column)
             row = None
-
             try:
                 row = self._get_metric_fn[metric_func.metric_type.value](
                     metric_func.metrics,
                     runner=runner,
                     session=session,
                     column=metric_func.column,
-                    sample=sample,
+                    sample=runner.dataset,
                 )
-                if row and isinstance(row, dict):
+                if isinstance(row, dict):
                     row = self._validate_nulls(row)
-
-                # System metrics return a list of dictionaries, with UPDATE, INSERT or DELETE ops results
-                if row and metric_func.metric_type == MetricTypes.System:
+                if isinstance(row, list):
                     row = [
                         self._validate_nulls(r) if isinstance(r, dict) else r
                         for r in row
@@ -537,24 +494,12 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
                     logger.debug(traceback.format_exc())
                     logger.error(f"Operation was cancelled due to TimeoutError - {exc}")
                     raise concurrent.futures.TimeoutError
+                except KeyboardInterrupt:
+                    pool.shutdown39(wait=True, cancel_futures=True)
+                    raise
 
         return profile_results
 
-    def fetch_sample_data(self, table, columns) -> TableData:
-        """Fetch sample data from database
-
-        Args:
-            table: ORM declarative table
-
-        Returns:
-            TableData: sample table data
-        """
-        sampler = self._get_sampler(
-            table=table,
-        )
-
-        return sampler.fetch_sample_data(columns)
-
     def get_composed_metrics(
         self, column: Column, metric: Metrics, column_results: Dict
     ):
@@ -576,21 +521,24 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
             return None
 
     def get_hybrid_metrics(
-        self, column: Column, metric: Metrics, column_results: Dict, **kwargs
+        self,
+        column: Column,
+        metric: Type[HybridMetric],
+        column_results: Dict[str, Any],
     ):
         """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
+            metric: metric to compute
+            column_results: results of the column
         Returns:
             dictionnary of results
         """
-        sampler = self._get_sampler(table=kwargs.get("table"))
-        sample = sampler.random_sample(column)
+        dataset = self.sampler.get_dataset(column=column)
         try:
-            return metric(column).fn(sample, column_results, self.session)
+            return metric(column).fn(dataset, column_results, self.session)
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unexpected exception computing metrics: {exc}")

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

@@ -24,13 +24,13 @@ from metadata.profiler.interface.sqlalchemy.profiler_interface import (
 )
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.runner import QueryRunner
-from metadata.profiler.source.single_store.metrics.window.first_quartile import (
+from metadata.profiler.source.database.single_store.metrics.window.first_quartile import (
     SingleStoreFirstQuartile,
 )
-from metadata.profiler.source.single_store.metrics.window.median import (
+from metadata.profiler.source.database.single_store.metrics.window.median import (
     SingleStoreMedian,
 )
-from metadata.profiler.source.single_store.metrics.window.third_quartile import (
+from metadata.profiler.source.database.single_store.metrics.window.third_quartile import (
     SingleStoreThirdQuartile,
 )
 from metadata.utils.logger import profiler_interface_registry_logger
@@ -76,11 +76,11 @@ class SingleStoreProfilerInterface(SQAProfilerInterface):
                 return dict(row)
         except ProgrammingError:
             logger.info(
-                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to overflow"
+                f"Skipping window metrics for {runner.table_name}.{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}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None

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

@@ -41,7 +41,7 @@ class SnowflakeProfilerInterface(SQAProfilerInterface):
             session.bind.dialect.name
         ):
             logger.info(
-                f"Computing metrics without sum for {runner.table.__tablename__}.{column.name}"
+                f"Computing metrics without sum for {runner.table_name}.{column.name}"
             )
             return self._compute_static_metrics_wo_sum(metrics, runner, session, column)
         return None

metadata/profiler/interface/sqlalchemy/stored_statistics_profiler.py

@@ -0,0 +1,145 @@
+#  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
+"""
+
+import threading
+from typing import Any, Dict, List, Set
+
+from more_itertools import partition
+from sqlalchemy import Column
+
+from metadata.generated.schema.entity.data.table import Table
+from metadata.mixins.sqalchemy.sqa_mixin import Root
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+from metadata.profiler.metrics.core import Metric
+from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.processor.runner import QueryRunner
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+thread_local = threading.local()
+
+
+class StoredStatisticsSource(Root):
+    def get_statistics_metrics(self) -> Set[Metrics]:
+        """Statistic metrics that are found in system tables. Different for each database."""
+        return set()
+
+    def get_column_statistics(
+        self, metric: List[Metrics], schema: str, table_name: Table, column: str
+    ) -> dict:
+        raise NotImplementedError(
+            "You used a connector that does not support using statistics tables."
+        )
+
+    def get_table_statistics(
+        self, metric: List[Metrics], schema: str, table_name: Table
+    ) -> dict:
+        raise NotImplementedError(
+            "You used a connector that does not support using statistics tables."
+        )
+
+
+class ProfilerWithStatistics(SQAProfilerInterface, StoredStatisticsSource):
+    """
+    Interface to interact with registry supporting
+    sqlalchemy.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Collaborative constructor"""
+        super().__init__(*args, **kwargs)
+
+    def _compute_static_metrics(
+        self,
+        metrics: List[Metrics],
+        runner: QueryRunner,
+        column,
+        session,
+        *args,
+        **kwargs,
+    ) -> dict:
+        result = {}
+        if self.source_config.useStatistics:
+            metrics, stat_metrics = map(
+                list,
+                partition(self.is_statistic_metric, metrics),
+            )
+            schema = runner.schema_name
+            table_name = runner.table_name
+            logger.debug(
+                "Getting statistics for column: %s.%s.%s",
+                schema,
+                table_name,
+                column.name,
+            )
+            result.update(
+                super().get_column_statistics(
+                    stat_metrics, schema, table_name, column.name
+                )
+            )
+        result.update(
+            super()._compute_static_metrics(
+                metrics,
+                runner,
+                column,
+                session,
+                *args,
+                **kwargs,
+            )
+        )
+        return result
+
+    def _compute_table_metrics(
+        self,
+        metrics: List[Metric],
+        runner: QueryRunner,
+        session,
+        *args,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        result = {}
+        if self.source_config.useStatistics:
+            metrics, stat_metrics = map(
+                list,
+                partition(self.is_statistic_metric, metrics),
+            )
+            schema = runner.schema_name
+            table_name = runner.table_name
+            logger.debug("Geting statistics for table: %s.%s", schema, table_name)
+            result.update(
+                super().get_table_statistics(stat_metrics, schema, table_name)
+            )
+        super_table_metrics = super()._compute_table_metrics(
+            metrics,
+            runner,
+            session,
+            *args,
+            **kwargs,
+        )
+        if super_table_metrics is not None:
+            result.update(super_table_metrics)
+        return result
+
+    def get_hybrid_metrics(self, column: Column, metric: Metric, column_results: Dict):
+        # this metrics might have been computed in a previous step
+        return column_results.get(metric.name()) or super().get_hybrid_metrics(
+            column, metric, column_results
+        )
+
+    def is_statistic_metric(self, metric: Metric) -> bool:
+        return metric.name() in {m.name for m in super().get_statistics_metrics()}

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

@@ -19,10 +19,15 @@ from typing import List
 from sqlalchemy import func
 from sqlalchemy.exc import ProgrammingError
 
+from metadata.ingestion.source.database.trino.profiler.system_tables_profiler import (
+    TrinoStoredStatisticsSource,
+)
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
-    SQAProfilerInterface,
     handle_query_exception,
 )
+from metadata.profiler.interface.sqlalchemy.stored_statistics_profiler import (
+    ProfilerWithStatistics,
+)
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.orm.registry import FLOAT_SET
 from metadata.profiler.processor.runner import QueryRunner
@@ -31,7 +36,7 @@ from metadata.utils.logger import profiler_interface_registry_logger
 logger = profiler_interface_registry_logger()
 
 
-class TrinoProfilerInterface(SQAProfilerInterface):
+class TrinoProfilerInterface(ProfilerWithStatistics, TrinoStoredStatisticsSource):
     """
     Interface to interact with registry supporting
     sqlalchemy.
@@ -71,11 +76,11 @@ class TrinoProfilerInterface(SQAProfilerInterface):
                 return dict(row)
         except ProgrammingError as err:
             logger.info(
-                f"Skipping window metrics for {runner.table.__tablename__}.{column.name} due to {err}"
+                f"Skipping window metrics for {runner.table_name}.{column.name} due to {err}"
             )
             return None
 
         except Exception as exc:
-            msg = f"Error trying to compute profile for {runner.table.__tablename__}.{column.name}: {exc}"
+            msg = f"Error trying to compute profile for {runner.table_name}.{column.name}: {exc}"
             handle_query_exception(msg, exc, session)
         return None

metadata/profiler/metrics/static/max.py

@@ -88,10 +88,12 @@ class Max(StaticMetric):
     def fn(self):
         """sqlalchemy function"""
         if is_concatenable(self.col.type):
-            return MaxFn(LenFn(column(self.col.name, self.col.type)))
+            return MaxFn(
+                LenFn(column(self.col.name, self.col.type)), type_=self.col.type
+            )
         if (not is_quantifiable(self.col.type)) and (not is_date_time(self.col.type)):
             return None
-        return MaxFn(column(self.col.name, self.col.type))
+        return MaxFn(column(self.col.name, self.col.type), type_=self.col.type)
 
     def df_fn(self, dfs=None):
         """pandas function"""

metadata/profiler/metrics/static/min.py

@@ -88,11 +88,13 @@ class Min(StaticMetric):
     def fn(self):
         """sqlalchemy function"""
         if is_concatenable(self.col.type):
-            return MinFn(LenFn(column(self.col.name, self.col.type)))
+            return MinFn(
+                LenFn(column(self.col.name, self.col.type)), type_=self.col.type
+            )
 
         if (not is_quantifiable(self.col.type)) and (not is_date_time(self.col.type)):
             return None
-        return MinFn(column(self.col.name, self.col.type))
+        return MinFn(column(self.col.name, self.col.type), type_=self.col.type)
 
     def df_fn(self, dfs=None):
         """pandas function"""