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

metadata/profiler/processor/core.py

@@ -33,7 +33,10 @@ from metadata.generated.schema.entity.data.table import (
     TableData,
     TableProfile,
 )
-from metadata.profiler.api.models import ProfilerResponse
+from metadata.generated.schema.tests.customMetric import (
+    CustomMetric as CustomMetricEntity,
+)
+from metadata.profiler.api.models import ProfilerResponse, ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import (
     ComposedMetric,
@@ -50,6 +53,7 @@ from metadata.profiler.metrics.static.row_count import RowCount
 from metadata.profiler.orm.registry import NOT_COMPUTE
 from metadata.profiler.processor.sample_data_handler import upload_sample_data
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -88,6 +92,7 @@ class Profiler(Generic[TMetric]):
         """
 
         self.profiler_interface = profiler_interface
+        self.source_config = self.profiler_interface.source_config
         self.include_columns = include_columns
         self.exclude_columns = exclude_columns
         self._metrics = metrics
@@ -273,6 +278,33 @@ class Profiler(Generic[TMetric]):
 
         return [metric for metric in metrics if metric.is_col_metric()]
 
+    def get_custom_metrics(
+        self, column_name: Optional[str] = None
+    ) -> Optional[List[CustomMetricEntity]]:
+        """Get custom metrics for a table or column
+
+        Args:
+            column (Optional[str]): optional column name. If None will fetch table level custom metrics
+
+        Returns:
+            List[str]
+        """
+        if column_name is None:
+            return self.profiler_interface.table_entity.customMetrics or None
+
+        # if we have a column we'll get the custom metrics for this column
+        column = next(
+            (
+                clmn
+                for clmn in self.profiler_interface.table_entity.columns
+                if clmn.name.__root__ == column_name
+            ),
+            None,
+        )
+        if column:
+            return column.customMetrics or None
+        return None
+
     @property
     def sample(self):
         """Return the sample used for the profiler"""
@@ -345,22 +377,40 @@ class Profiler(Generic[TMetric]):
 
     def _prepare_table_metrics(self) -> List:
         """prepare table metrics"""
+        metrics = []
         table_metrics = [
             metric
             for metric in self.static_metrics
             if (not metric.is_col_metric() and not metric.is_system_metrics())
         ]
 
+        custom_table_metrics = self.get_custom_metrics()
+
         if table_metrics:
-            return [
-                (
-                    table_metrics,  # metric functions
-                    MetricTypes.Table,  # metric type for function mapping
-                    None,  # column name
-                    self.table,  # table name
-                ),
-            ]
-        return []
+            metrics.extend(
+                [
+                    ThreadPoolMetrics(
+                        metrics=table_metrics,
+                        metric_type=MetricTypes.Table,
+                        column=None,
+                        table=self.table,
+                    )
+                ]
+            )
+
+        if custom_table_metrics:
+            metrics.extend(
+                [
+                    ThreadPoolMetrics(
+                        metrics=custom_table_metrics,
+                        metric_type=MetricTypes.Custom,
+                        column=None,
+                        table=self.table,
+                    )
+                ]
+            )
+
+        return metrics
 
     def _prepare_system_metrics(self) -> List:
         """prepare system metrics"""
@@ -368,11 +418,11 @@ class Profiler(Generic[TMetric]):
 
         if system_metrics:
             return [
-                (
-                    system_metric,  # metric functions
-                    MetricTypes.System,  # metric type for function mapping
-                    None,  # column name
-                    self.table,  # table name
+                ThreadPoolMetrics(
+                    metrics=system_metric,  # metric functions
+                    metric_type=MetricTypes.System,  # metric type for function mapping
+                    column=None,  # column name
+                    table=self.table,  # table name
                 )
                 for system_metric in system_metrics
             ]
@@ -386,45 +436,60 @@ class Profiler(Generic[TMetric]):
             for column in self.columns
             if column.type.__class__.__name__ not in NOT_COMPUTE
         ]
+        column_metrics_for_thread_pool = []
+        static_metrics = [
+            ThreadPoolMetrics(
+                metrics=[
+                    metric
+                    for metric in self.get_col_metrics(self.static_metrics, column)
+                    if not metric.is_window_metric()
+                ],
+                metric_type=MetricTypes.Static,
+                column=column,
+                table=self.table,
+            )
+            for column in columns
+        ]
+        query_metrics = [
+            ThreadPoolMetrics(
+                metrics=metric,
+                metric_type=MetricTypes.Query,
+                column=column,
+                table=self.table,
+            )
+            for column in columns
+            for metric in self.get_col_metrics(self.query_metrics, column)
+        ]
+        window_metrics = [
+            ThreadPoolMetrics(
+                metrics=[
+                    metric
+                    for metric in self.get_col_metrics(self.static_metrics, column)
+                    if metric.is_window_metric()
+                ],
+                metric_type=MetricTypes.Window,
+                column=column,
+                table=self.table,
+            )
+            for column in columns
+        ]
 
-        column_metrics_for_thread_pool = [
-            *[
-                (
-                    [
-                        metric
-                        for metric in self.get_col_metrics(self.static_metrics, column)
-                        if not metric.is_window_metric()
-                    ],
-                    MetricTypes.Static,
-                    column,
-                    self.table,
-                )
-                for column in columns
-            ],
-            *[
-                (
-                    metric,
-                    MetricTypes.Query,
-                    column,
-                    self.table,
-                )
-                for column in columns
-                for metric in self.get_col_metrics(self.query_metrics, column)
-            ],
-            *[
-                (
-                    [
-                        metric
-                        for metric in self.get_col_metrics(self.static_metrics, column)
-                        if metric.is_window_metric()
-                    ],
-                    MetricTypes.Window,
-                    column,
-                    self.table,
+        # we'll add the system metrics to the thread pool computation
+        for metric_type in [static_metrics, query_metrics, window_metrics]:
+            column_metrics_for_thread_pool.extend(metric_type)
+
+        # we'll add the custom metrics to the thread pool computation
+        for column in columns:
+            custom_metrics = self.get_custom_metrics(column.name)
+            if custom_metrics:
+                column_metrics_for_thread_pool.append(
+                    ThreadPoolMetrics(
+                        metrics=custom_metrics,
+                        metric_type=MetricTypes.Custom,
+                        column=column,
+                        table=self.table,
+                    )
                 )
-                for column in columns
-            ],
-        ]
 
         return column_metrics_for_thread_pool
 
@@ -455,27 +520,27 @@ class Profiler(Generic[TMetric]):
 
         return self
 
-    def process(
-        self,
-        generate_sample_data: Optional[bool],
-    ) -> ProfilerResponse:
+    def process(self) -> ProfilerResponse:
         """
         Given a table, we will prepare the profiler for
         all its columns and return all the run profilers
         in a Dict in the shape {col_name: Profiler}
         """
-        logger.debug(
-            f"Computing profile metrics for {self.profiler_interface.table_entity.fullyQualifiedName.__root__}..."
-        )
 
-        self.compute_metrics()
-        if generate_sample_data:
+        if self.source_config.computeMetrics:
+            logger.debug(
+                f"Computing profile metrics for {self.profiler_interface.table_entity.fullyQualifiedName.__root__}..."
+            )
+            self.compute_metrics()
+
+        if self.source_config.generateSampleData:
             sample_data = self.generate_sample_data()
         else:
             sample_data = None
 
         profile = self.get_profile()
-        self._check_profile_and_handle(profile)
+        if self.source_config.computeMetrics:
+            self._check_profile_and_handle(profile)
 
         table_profile = ProfilerResponse(
             table=self.profiler_interface.table_entity,
@@ -485,6 +550,7 @@ class Profiler(Generic[TMetric]):
 
         return table_profile
 
+    @calculate_execution_time(store=False)
     def generate_sample_data(self) -> Optional[TableData]:
         """Fetch and ingest sample data
 
@@ -568,6 +634,7 @@ class Profiler(Generic[TMetric]):
                 profileSampleType=self.profile_sample_config.profile_sample_type
                 if self.profile_sample_config
                 else None,
+                customMetrics=self._table_results.get("customMetrics"),
             )
 
             if self._system_results:

metadata/profiler/processor/default.py

@@ -17,19 +17,28 @@ from typing import List, Optional
 from sqlalchemy.orm import DeclarativeMeta
 
 from metadata.generated.schema.entity.data.table import ColumnProfilerConfig
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import Metric, add_props
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.core import Profiler
 
 
-def get_default_metrics(table: DeclarativeMeta) -> List[Metric]:
+def get_default_metrics(
+    table: DeclarativeMeta,
+    ometa_client: Optional[OpenMetadata] = None,
+    db_service: Optional[DatabaseService] = None,
+) -> List[Metric]:
     return [
         # Table Metrics
         Metrics.ROW_COUNT.value,
         add_props(table=table)(Metrics.COLUMN_COUNT.value),
         add_props(table=table)(Metrics.COLUMN_NAMES.value),
-        add_props(table=table)(Metrics.SYSTEM.value),
+        # We'll use the ometa_client & db_service in case we need to fetch info to ES
+        add_props(table=table, ometa_client=ometa_client, db_service=db_service)(
+            Metrics.SYSTEM.value
+        ),
         # Column Metrics
         Metrics.MEDIAN.value,
         Metrics.FIRST_QUARTILE.value,
@@ -65,7 +74,9 @@ class DefaultProfiler(Profiler):
         include_columns: Optional[List[ColumnProfilerConfig]] = None,
         exclude_columns: Optional[List[str]] = None,
     ):
-        _metrics = get_default_metrics(profiler_interface.table)
+        _metrics = get_default_metrics(
+            table=profiler_interface.table, ometa_client=profiler_interface.ometa_client
+        )
 
         super().__init__(
             *_metrics,

metadata/profiler/processor/handle_partition.py

@@ -112,8 +112,8 @@ class partition_filter_handler:
                 )
                 if self.build_sample:
                     return (
-                        _self.client.query(
-                            _self.table,
+                        _self._base_sample_query(
+                            kwargs.get("column"),
                             (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
                         )
                         .filter(partition_filter)

metadata/profiler/processor/processor.py

@@ -14,13 +14,16 @@ Profiler Processor Step
 import traceback
 from typing import cast
 
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import Processor
@@ -37,7 +40,6 @@ class ProfilerProcessor(Processor):
     """
 
     def __init__(self, config: OpenMetadataWorkflowConfig):
-
         super().__init__()
 
         self.config = config
@@ -48,22 +50,23 @@ class ProfilerProcessor(Processor):
             DatabaseServiceProfilerPipeline, self.config.source.sourceConfig.config
         )  # Used to satisfy type checked
 
-    def _run(self, record: ProfilerSourceAndEntity) -> Either[ProfilerResponse]:
+    @property
+    def name(self) -> str:
+        return "Profiler"
 
+    def _run(self, record: ProfilerSourceAndEntity) -> Either[ProfilerResponse]:
         profiler_runner: Profiler = record.profiler_source.get_profiler_runner(
             record.entity, self.profiler_config
         )
 
         try:
-            profile: ProfilerResponse = profiler_runner.process(
-                self.source_config.generateSampleData,
-            )
+            profile: ProfilerResponse = profiler_runner.process()
         except Exception as exc:
             self.status.failed(
                 StackTraceError(
                     name=record.entity.fullyQualifiedName.__root__,
                     error=f"Unexpected exception processing entity {record.entity.fullyQualifiedName.__root__}: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
             self.status.failures.extend(

metadata/profiler/processor/sample_data_handler.py

@@ -11,7 +11,7 @@
 """
 Profiler Processor Step
 """
-
+import json
 import traceback
 from datetime import datetime
 from functools import singledispatch
@@ -48,7 +48,11 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
         sample_storage_config = profiler_interface.storage_config
         if not sample_storage_config:
             return
-        df = pd.DataFrame(data=data.rows, columns=[i.__root__ for i in data.columns])
+        deserialized_data = json.loads(data.json())
+        df = pd.DataFrame(
+            data=deserialized_data.get("rows", []),
+            columns=[i.__root__ for i in data.columns],
+        )
         pq_buffer = BytesIO()
         df.to_parquet(pq_buffer)
         object_key = _get_object_key(

metadata/profiler/processor/sampler/sqlalchemy/bigquery/sampler.py

@@ -14,6 +14,7 @@ for the profiler
 """
 from typing import Dict, Optional
 
+from sqlalchemy import Column
 from sqlalchemy.orm import Query
 
 from metadata.generated.schema.entity.data.table import ProfileSampleType, TableType
@@ -50,8 +51,35 @@ class BigQuerySampler(SQASampler):
         )
         self.table_type: TableType = table_type
 
+    def _base_sample_query(self, column: Optional[Column], label=None):
+        """Base query for sampling
+
+        Args:
+            column (Optional[Column]): if computing a column metric only sample for the column
+            label (_type_, optional):
+
+        Returns:
+        """
+        # pylint: disable=import-outside-toplevel
+        from sqlalchemy_bigquery import STRUCT
+
+        if column is not None:
+            column_parts = column.name.split(".")
+            if len(column_parts) > 1:
+                # for struct columns (e.g. `foo.bar`) we need to create a new column corresponding to
+                # the struct (e.g. `foo`) and then use that in the sample query as the column that
+                # will be query is `foo.bar`.
+                # e.g. WITH sample AS (SELECT `foo` FROM table) SELECT `foo.bar`
+                # FROM sample TABLESAMPLE SYSTEM (n PERCENT)
+                column = Column(column_parts[0], STRUCT)
+                # pylint: disable=protected-access
+                column._set_parent(self.table.__table__)
+                # pylint: enable=protected-access
+
+        return super()._base_sample_query(column, label=label)
+
     @partition_filter_handler(build_sample=True)
-    def get_sample_query(self) -> Query:
+    def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
         # TABLESAMPLE SYSTEM is not supported for views
         if (
@@ -59,11 +87,11 @@ class BigQuerySampler(SQASampler):
             and self.table_type != TableType.View
         ):
             return (
-                self._base_sample_query()
+                self._base_sample_query(column)
                 .suffix_with(
                     f"TABLESAMPLE SYSTEM ({self.profile_sample or 100} PERCENT)",
                 )
                 .cte(f"{self.table.__tablename__}_sample")
             )
 
-        return super().get_sample_query()
+        return super().get_sample_query(column=column)

metadata/profiler/processor/sampler/sqlalchemy/sampler.py

@@ -31,6 +31,7 @@ from metadata.profiler.orm.functions.random_num import RandomNumFn
 from metadata.profiler.orm.registry import Dialects
 from metadata.profiler.processor.handle_partition import partition_filter_handler
 from metadata.profiler.processor.sampler.sampler_interface import SamplerInterface
+from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_utils import (
     build_query_filter,
@@ -67,17 +68,28 @@ class SQASampler(SamplerInterface):
     run the query in the whole table.
     """
 
-    def _base_sample_query(self, label=None):
+    def _base_sample_query(self, column: Optional[Column], label=None):
+        """Base query for sampling
+
+        Args:
+            column (Optional[Column]): if computing a column metric only sample for the column
+            label (_type_, optional):
+
+        Returns:
+        """
+        # only sample the column if we are computing a column metric to limit the amount of data scaned
+        entity = self.table if column is None else column
         if label is not None:
-            return self.client.query(self.table, label)
-        return self.client.query(self.table)
+            return self.client.query(entity, label)
+        return self.client.query(entity)
 
     @partition_filter_handler(build_sample=True)
-    def get_sample_query(self) -> Query:
+    def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
         if self.profile_sample_type == ProfileSampleType.PERCENTAGE:
             rnd = (
                 self._base_sample_query(
+                    column,
                     (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
                 )
                 .suffix_with(
@@ -93,6 +105,7 @@ class SQASampler(SamplerInterface):
 
         table_query = self.client.query(self.table)
         session_query = self._base_sample_query(
+            column,
             (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL),
         )
         return (
@@ -101,7 +114,7 @@ class SQASampler(SamplerInterface):
             .cte(f"{self.table.__tablename__}_rnd")
         )
 
-    def random_sample(self) -> Union[DeclarativeMeta, AliasedClass]:
+    def random_sample(self, ccolumn=None) -> Union[DeclarativeMeta, AliasedClass]:
         """
         Either return a sampled CTE of table, or
         the full table if no sampling is required.
@@ -116,7 +129,7 @@ class SQASampler(SamplerInterface):
             return self.table
 
         # Add new RandomNumFn column
-        sampled = self.get_sample_query()
+        sampled = self.get_sample_query(column=ccolumn)
 
         # Assign as an alias
         return aliased(self.table, sampled)
@@ -171,6 +184,11 @@ class SQASampler(SamplerInterface):
 
     def _fetch_sample_data_from_user_query(self) -> TableData:
         """Returns a table data object using results from query execution"""
+        if not is_safe_sql_query(self._profile_sample_query):
+            raise RuntimeError(
+                f"SQL expression is not safe\n\n{self._profile_sample_query}"
+            )
+
         rnd = self.client.execute(f"{self._profile_sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
@@ -183,6 +201,11 @@ class SQASampler(SamplerInterface):
 
     def _rdn_sample_from_user_query(self) -> Query:
         """Returns sql alchemy object to use when running profiling"""
+        if not is_safe_sql_query(self._profile_sample_query):
+            raise RuntimeError(
+                f"SQL expression is not safe\n\n{self._profile_sample_query}"
+            )
+
         return self.client.query(self.table).from_statement(
             text(f"{self._profile_sample_query}")
         )

metadata/profiler/processor/sampler/sqlalchemy/trino/sampler.py

@@ -13,7 +13,6 @@ Helper module to handle data sampling
 for the profiler
 """
 from sqlalchemy import inspect, or_, text
-from trino.sqlalchemy.dialect import TrinoDialect
 
 from metadata.profiler.orm.registry import FLOAT_SET
 from metadata.profiler.processor.handle_partition import RANDOM_LABEL
@@ -26,11 +25,18 @@ class TrinoSampler(SQASampler):
     run the query in the whole table.
     """
 
-    TrinoDialect._json_deserializer = None  # pylint: disable=protected-access
+    def __init__(self, *args, **kwargs):
+        # pylint: disable=import-outside-toplevel
+        from trino.sqlalchemy.dialect import TrinoDialect
 
-    def _base_sample_query(self, label=None):
+        TrinoDialect._json_deserializer = None
+
+        super().__init__(*args, **kwargs)
+
+    def _base_sample_query(self, column, label=None):
         sqa_columns = [col for col in inspect(self.table).c if col.name != RANDOM_LABEL]
-        return self.client.query(self.table, label).where(
+        entity = self.table if column is None else column
+        return self.client.query(entity, label).where(
             or_(
                 *[
                     text(f"is_nan({cols.name}) = False")

metadata/profiler/source/base/profiler_source.py

@@ -217,7 +217,6 @@ class ProfilerSource(ProfilerSourceInterface):
     def _get_context_entities(
         self, entity: Table
     ) -> Tuple[DatabaseSchema, Database, DatabaseService]:
-
         schema_entity = None
         database_entity = None
         db_service = None
@@ -279,7 +278,11 @@ class ProfilerSource(ProfilerSourceInterface):
         metrics = (
             [Metrics.get(name) for name in profiler_config.profiler.metrics]
             if profiler_config.profiler.metrics
-            else get_default_metrics(profiler_interface.table)
+            else get_default_metrics(
+                table=profiler_interface.table,
+                ometa_client=self.ometa_client,
+                db_service=db_service,
+            )
         )
 
         return Profiler(

metadata/profiler/source/bigquery/type_mapper.py

@@ -29,7 +29,6 @@ def bigquery_type_mapper(_type_map: dict, col: Column):
     from sqlalchemy_bigquery import STRUCT
 
     def build_struct(_type_map: dict, col: Column):
-
         structs = []
         for child in col.children:
             if child.dataType != DataType.STRUCT:

metadata/profiler/source/databricks/profiler_source.py

@@ -0,0 +1,36 @@
+"""Extend the ProfilerSource class to add support for Databricks is_disconnect SQA method"""
+
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.base.profiler_source import ProfilerSource
+
+
+def is_disconnect(self, e, connection, cursor):
+    """is_disconnect method for the Databricks dialect"""
+    if "Invalid SessionHandle: SessionHandle" in str(e):
+        return True
+    return False
+
+
+class DataBricksProfilerSource(ProfilerSource):
+    """Databricks Profiler source"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        database: DatabaseService,
+        ometa_client: OpenMetadata,
+    ):
+        super().__init__(config, database, ometa_client)
+        self.set_is_disconnect()
+
+    def set_is_disconnect(self):
+        """Set the is_disconnect method for the Databricks dialect"""
+        from databricks.sqlalchemy import (
+            DatabricksDialect,  # pylint: disable=import-outside-toplevel
+        )
+
+        DatabricksDialect.is_disconnect = is_disconnect

metadata/profiler/source/mariadb/functions/median.py

@@ -0,0 +1,20 @@
+"""Median function for MariaDB"""
+
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.functions import FunctionElement
+
+from metadata.profiler.metrics.core import CACHE
+
+
+class MariaDBMedianFn(FunctionElement):
+    inherit_cache = CACHE
+
+
+@compiles(MariaDBMedianFn)
+def _(elements, compiler, **kwargs):  # pylint: disable=unused-argument
+    col = compiler.process(elements.clauses.clauses[0])
+    percentile = elements.clauses.clauses[2].value
+    # According to the documentation available at https://mariadb.com/kb/en/median/#description,
+    # the PERCENTILE_CONT function can be utilized to calculate the median. Therefore, it is
+    # being used in this context.
+    return f"PERCENTILE_CONT({percentile:.2f}) WITHIN GROUP (ORDER BY {col}) OVER()"