openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/profiler/orm/functions/table_metric_computer.py

@@ -53,7 +53,7 @@ class AbstractTableMetricComputer(ABC):
         self._metrics = metrics
         self._conn_config = conn_config
         self._database = self._runner._session.get_bind().url.database
-        self._table = self._runner.table
+        self._table = self._runner.dataset
         self._entity = entity
 
     @property
@@ -91,8 +91,8 @@ class AbstractTableMetricComputer(ABC):
             table (DeclarativeMeta): _description_
         """
         try:
-            self._schema_name = self.table.__table_args__.get("schema")
-            self._table_name = self.table.__tablename__
+            self._schema_name = self.table.schema
+            self._table_name = self.table.name
         except AttributeError:
             raise AttributeError(ERROR_MSG)
 
@@ -147,7 +147,8 @@ class BaseTableMetricComputer(AbstractTableMetricComputer):
     """Base table computer"""
 
     def compute(self):
-        """Default compute behavior for table metrics"""
+        """Default compute behavior for table metrics. This method will use the raw table
+        to compute metrics and omit any sampling or partitioning logic."""
         return self.runner.select_first_from_table(
             *[metric().fn() for metric in self.metrics]
         )
@@ -412,7 +413,7 @@ class RedshiftTableMetricComputer(BaseTableMetricComputer):
         )
         res = self.runner._session.execute(query).first()
         if not res:
-            return None
+            return super().compute()
         if res.rowCount is None or (
             res.rowCount == 0 and self._entity.tableType == TableType.View
         ):

metadata/profiler/orm/functions/unique_count.py

@@ -53,5 +53,11 @@ def _unique_count_query_mssql(col, session, sample):
     )
 
 
+def _unique_count_query_oracle(col, session, sample):
+    count_fn = CountFn(col)
+    return _unique_count_query(count_fn, session, sample)
+
+
 _unique_count_query_mapper = defaultdict(lambda: _unique_count_query)
 _unique_count_query_mapper[Dialects.MSSQL] = _unique_count_query_mssql
+_unique_count_query_mapper[Dialects.Oracle] = _unique_count_query_oracle

metadata/profiler/orm/registry.py

@@ -14,10 +14,11 @@ Custom types' registry for easy access
 without having an import mess
 """
 import math
+from enum import Enum
 
 import sqlalchemy
 from sqlalchemy import Date, DateTime, Integer, Numeric, Time
-from sqlalchemy.sql.sqltypes import Concatenable, Enum
+from sqlalchemy.sql.sqltypes import Concatenable
 
 from metadata.generated.schema.entity.data.table import DataType
 from metadata.ingestion.source import sqa_types
@@ -45,7 +46,7 @@ class CustomTypes(TypeRegistry):
     UNDETERMINED = UndeterminedType
 
 
-class Dialects(Enum):
+class PythonDialects(Enum):
     """
     Map the service types from DatabaseServiceType
     to the dialect scheme name used for ingesting
@@ -54,6 +55,8 @@ class Dialects(Enum):
     Keep this alphabetically ordered
     """
 
+    # pylint: disable=invalid-name
+
     Athena = "awsathena"
     AzureSQL = "azuresql"
     BigQuery = "bigquery"
@@ -83,6 +86,28 @@ class Dialects(Enum):
     Vertica = "vertica"
 
 
+class EnumAdapter(type):
+    """A hack to use the Dialects string values can be accesses
+    without using the value attribute.
+
+    Example:
+        Dialets.MySQL == "mysql"
+
+    Instead of:
+        Dialects.MySQL.value == "mysql"
+
+    We use this functionality when registring sqlalchemy custom functions. But we should
+    avoid using this pattern as it can be confusing.
+    """
+
+    def __getattr__(cls, item):
+        return PythonDialects[item].value
+
+
+class Dialects(metaclass=EnumAdapter):
+    pass
+
+
 # Sometimes we want to skip certain types for computing metrics.
 # If the type is NULL, then we won't run the metric execution
 # in the profiler.

metadata/profiler/processor/core.py

@@ -40,7 +40,7 @@ from metadata.generated.schema.tests.customMetric import (
     CustomMetric as CustomMetricEntity,
 )
 from metadata.generated.schema.type.basic import Timestamp
-from metadata.profiler.api.models import ProfilerResponse, SampleData, ThreadPoolMetrics
+from metadata.profiler.api.models import ProfilerResponse, ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import (
     ComposedMetric,
@@ -54,9 +54,6 @@ from metadata.profiler.metrics.static.row_count import RowCount
 from metadata.profiler.orm.functions.table_metric_computer import CREATE_DATETIME
 from metadata.profiler.orm.registry import NOT_COMPUTE
 from metadata.profiler.processor.metric_filter import MetricFilter
-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()
@@ -107,7 +104,6 @@ class Profiler(Generic[TMetric]):
         self.exclude_columns = exclude_columns
         self._metrics = metrics
         self._profile_ts = Timestamp(int(datetime.now().timestamp() * 1000))
-        self.profile_sample_config = self.profiler_interface.profile_sample_config
 
         self.metric_filter = MetricFilter(
             metrics=self.metrics,
@@ -249,11 +245,6 @@ class Profiler(Generic[TMetric]):
             return column.customMetrics or None
         return None
 
-    @property
-    def sample(self):
-        """Return the sample used for the profiler"""
-        return self.profiler_interface.sample
-
     def validate_composed_metric(self) -> None:
         """
         Make sure that all composed metrics have
@@ -273,9 +264,6 @@ class Profiler(Generic[TMetric]):
 
         Data should be saved under self.results
         """
-
-        logger.debug("Running post Profiler...")
-
         current_col_results: Dict[str, Any] = self._column_results.get(col.name)
         if not current_col_results:
             logger.debug(
@@ -320,12 +308,15 @@ class Profiler(Generic[TMetric]):
                 col,
                 metric,
                 current_col_results,
-                table=self.table,
             )
 
     def _prepare_table_metrics(self) -> List:
         """prepare table metrics"""
         metrics = []
+
+        if self.source_config and not self.source_config.computeTableMetrics:
+            return metrics
+
         table_metrics = [
             metric
             for metric in self.metric_filter.static_metrics
@@ -380,6 +371,9 @@ class Profiler(Generic[TMetric]):
     def _prepare_column_metrics(self) -> List:
         """prepare column metrics"""
         column_metrics_for_thread_pool = []
+        if self.source_config and not self.source_config.computeColumnMetrics:
+            return column_metrics_for_thread_pool
+
         columns = [
             column
             for column in self.columns
@@ -491,16 +485,6 @@ class Profiler(Generic[TMetric]):
             )
             self.compute_metrics()
 
-        # We need the sample data for Sample Data or PII Sensitive processing.
-        # We'll nullify the Sample Data after the PII processing so that it's not stored.
-        if (
-            self.source_config.generateSampleData
-            or self.source_config.processPiiSensitive
-        ):
-            sample_data = self.generate_sample_data()
-        else:
-            sample_data = None
-
         profile = self.get_profile()
         if self.source_config.computeMetrics:
             self._check_profile_and_handle(profile)
@@ -508,43 +492,10 @@ class Profiler(Generic[TMetric]):
         table_profile = ProfilerResponse(
             table=self.profiler_interface.table_entity,
             profile=profile,
-            sample_data=sample_data,
         )
 
         return table_profile
 
-    @calculate_execution_time(store=False)
-    def generate_sample_data(self) -> Optional[SampleData]:
-        """Fetch and ingest sample data
-
-        Returns:
-            TableData: sample data
-        """
-        try:
-            logger.debug(
-                "Fetching sample data for "
-                f"{self.profiler_interface.table_entity.fullyQualifiedName.root}..."  # type: ignore
-            )
-            table_data = self.profiler_interface.fetch_sample_data(
-                self.table, self.columns
-            )
-            upload_sample_data(
-                data=table_data, profiler_interface=self.profiler_interface
-            )
-            table_data.rows = table_data.rows[
-                : min(
-                    SAMPLE_DATA_DEFAULT_COUNT, self.profiler_interface.sample_data_count
-                )
-            ]
-            return SampleData(
-                data=table_data, store=self.source_config.generateSampleData
-            )
-
-        except Exception as err:
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Error fetching sample data: {err}")
-            return None
-
     def get_profile(self) -> CreateTableProfileRequest:
         """
         After executing the profiler, get all results
@@ -599,13 +550,13 @@ class Profiler(Generic[TMetric]):
                 createDateTime=raw_create_date,
                 sizeInByte=self._table_results.get("sizeInBytes"),
                 profileSample=(
-                    self.profile_sample_config.profile_sample
-                    if self.profile_sample_config
+                    self.profiler_interface.sampler.sample_config.profile_sample
+                    if self.profiler_interface.sampler.sample_config
                     else None
                 ),
                 profileSampleType=(
-                    self.profile_sample_config.profile_sample_type
-                    if self.profile_sample_config
+                    self.profiler_interface.sampler.sample_config.profile_sample_type
+                    if self.profiler_interface.sampler.sample_config
                     else None
                 ),
                 customMetrics=self._table_results.get("customMetrics"),
@@ -613,7 +564,7 @@ class Profiler(Generic[TMetric]):
 
             if self._system_results:
                 system_profile = [
-                    SystemProfile(**system_result)
+                    SystemProfile.model_validate(system_result)
                     for system_result in self._system_results
                 ]
             else:

metadata/profiler/processor/handle_partition.py

@@ -22,8 +22,6 @@ from metadata.generated.schema.entity.data.table import (
     PartitionIntervalTypes,
     PartitionProfilerConfig,
 )
-from metadata.profiler.orm.functions.modulo import ModuloFn
-from metadata.profiler.orm.functions.random_num import RandomNumFn
 from metadata.utils.logger import profiler_logger
 from metadata.utils.sqa_utils import (
     build_query_filter,
@@ -80,49 +78,3 @@ def build_partition_predicate(
         [(Column(partition_field), "ge", date_or_datetime_fn)],
         False,
     )
-
-
-# pylint: disable=invalid-name,protected-access
-class partition_filter_handler:
-    """Decorator to handle partioned queries (hence lowercase class name)
-
-    Attributes:
-        first (bool): whether to return just the first row
-        sampled (bool): whether data should be sampled
-        build_sample (bool): whether to build the sample data set
-    """
-
-    def __init__(
-        self,
-        first: bool = True,
-        sampled: bool = False,
-        build_sample: bool = False,
-    ):
-        self.first = first
-        self.sampled = sampled
-        self.build_sample = build_sample
-
-    def __call__(self, func):
-        def handle_and_execute(_self, *args, **kwargs):
-            """Handle partitioned queries"""
-            if _self._partition_details:
-                partition_filter = build_partition_predicate(
-                    _self._partition_details,
-                    _self.table.__table__.c,
-                )
-                if self.build_sample:
-                    return (
-                        _self._base_sample_query(
-                            kwargs.get("column"),
-                            (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
-                        )
-                        .filter(partition_filter)
-                        .cte(f"{_self.table.__tablename__}_rnd")
-                    )
-                query_results = _self._build_query(*args, **kwargs).select_from(
-                    _self._sample if self.sampled else _self.table
-                )
-                return query_results.first() if self.first else query_results.all()
-            return func(_self, *args, **kwargs)
-
-        return handle_and_execute

metadata/profiler/processor/processor.py

@@ -30,7 +30,7 @@ from metadata.ingestion.api.steps import Processor
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.api.models import ProfilerProcessorConfig, ProfilerResponse
 from metadata.profiler.processor.core import Profiler
-from metadata.profiler.source.metadata import ProfilerSourceAndEntity
+from metadata.profiler.source.model import ProfilerSourceAndEntity
 
 
 class ProfilerProcessor(Processor):

metadata/profiler/processor/runner.py

@@ -18,14 +18,10 @@ and manage behavior such as timeouts.
 """
 from typing import Dict, Optional, Union
 
-from sqlalchemy import text
+from sqlalchemy import Table, text
 from sqlalchemy.orm import DeclarativeMeta, Query, Session
 from sqlalchemy.orm.util import AliasedClass
 
-from metadata.profiler.processor.handle_partition import (
-    build_partition_predicate,
-    partition_filter_handler,
-)
 from metadata.utils.logger import query_runner_logger
 from metadata.utils.sqa_utils import get_query_filter_for_runner
 
@@ -47,25 +43,66 @@ class QueryRunner:
     def __init__(
         self,
         session: Session,
-        table: DeclarativeMeta,
-        sample: Union[DeclarativeMeta, AliasedClass],
+        dataset: Union[DeclarativeMeta, AliasedClass],
+        raw_dataset: Table,
         partition_details: Optional[Dict] = None,
         profile_sample_query: Optional[str] = None,
     ):
         self._session = session
-        self.table = table
-        self._sample = sample
-        self._partition_details = partition_details
-        self._profile_sample_query = profile_sample_query
+        self._dataset = dataset
+        self.partition_details = partition_details
+        self.profile_sample_query = profile_sample_query
+        self.raw_dataset = raw_dataset
+
+    @property
+    def table(self) -> Table:
+        """Backward compatibility table attribute access"""
+        return self.raw_dataset
+
+    @property
+    def _sample(self):
+        """Backward compatibility _sample attribute access"""
+        return self._dataset
+
+    @property
+    def dataset(self):
+        """Dataset attribute access"""
+        return self._dataset
+
+    @dataset.setter
+    def dataset(self, dataset):
+        self._dataset = dataset
+
+    @property
+    def table_name(self):
+        """Table name attribute access"""
+        return self.raw_dataset.__table__.name
+
+    @property
+    def schema_name(self):
+        """Table name attribute access"""
+        return self.raw_dataset.__table__.schema
 
     def _build_query(self, *entities, **kwargs) -> Query:
+        """Build query object
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         return self._session.query(*entities, **kwargs)
 
     def _select_from_sample(self, *entities, **kwargs):
-        """Run select statement against sample data"""
+        """This method will use the sample data
+        and the partitioning logic if available otherwise it will use the raw table.
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         filter_ = get_query_filter_for_runner(kwargs)
 
-        query = self._build_query(*entities, **kwargs).select_from(self._sample)
+        query = self._build_query(*entities, **kwargs).select_from(self._dataset)
 
         if filter_ is not None:
             return query.filter(filter_)
@@ -73,11 +110,16 @@ class QueryRunner:
         return query
 
     def _select_from_user_query(self, *entities, **kwargs):
-        """Run select statement against user defined query"""
+        """Use the user query to select data from the table
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         filter_ = get_query_filter_for_runner(kwargs)
 
-        user_query = self._session.query(self.table).from_statement(
-            text(f"{self._profile_sample_query}")
+        user_query = self._session.query(self._dataset).from_statement(
+            text(f"{self.profile_sample_query}")
         )
 
         query = self._build_query(*entities, **kwargs).select_from(user_query)
@@ -87,12 +129,17 @@ class QueryRunner:
 
         return query
 
-    @partition_filter_handler()
     def select_first_from_table(self, *entities, **kwargs):
-        """Select first row from the table"""
+        """Select first row from the table. This method will use the raw table and
+        omit any sampling or partitioning logic.
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         filter_ = get_query_filter_for_runner(kwargs)
 
-        if self._profile_sample_query:
+        if self.profile_sample_query:
             return self._select_from_user_query(*entities, **kwargs).first()
         query = self._build_query(*entities, **kwargs).select_from(self.table)
 
@@ -101,12 +148,17 @@ class QueryRunner:
 
         return query.first()
 
-    @partition_filter_handler(first=False)
     def select_all_from_table(self, *entities, **kwargs):
-        """Select all rows from the table"""
+        """Select all rows from the table. This method will use the raw table and
+        omit any sampling or partitioning logic.
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         filter_ = get_query_filter_for_runner(kwargs)
 
-        if self._profile_sample_query:
+        if self.profile_sample_query:
             return self._select_from_user_query(*entities, **kwargs).all()
 
         query = self._build_query(*entities, **kwargs).select_from(self.table)
@@ -116,23 +168,33 @@ class QueryRunner:
 
         return query.all()
 
-    @partition_filter_handler(sampled=True)
     def select_first_from_sample(self, *entities, **kwargs):
+        """Select first row from the sample data. This method will use the sample data
+        and the partitioning logic if available otherwise it will use the raw table.
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         return self._select_from_sample(*entities, **kwargs).first()
 
-    @partition_filter_handler(first=False, sampled=True)
     def select_all_from_sample(self, *entities, **kwargs):
+        """Select all rows from the sample data. This method will use the sample data
+        and the partitioning logic if available otherwise it will use the raw table.
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         return self._select_from_sample(*entities, **kwargs).all()
 
     def yield_from_sample(self, *entities, **kwargs):
-        query = self._select_from_sample(*entities, **kwargs)
-        if self._partition_details:
-            partition_filter = build_partition_predicate(
-                self._partition_details,
-                self.table.__table__.c,
-            )
-            query.filter(partition_filter)
+        """Yield rows from the sample data
 
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
         result = self._session.execute(self._select_from_sample(*entities, **kwargs))
         while True:
             rows = result.fetchmany(1000)
@@ -141,15 +203,29 @@ class QueryRunner:
             yield from rows
 
     def dispatch_query_select_first(self, *entities, **kwargs):
-        """dispatch query to sample or all table"""
-        if isinstance(self._sample, AliasedClass):
-            return self.select_first_from_sample(*entities, **kwargs)
-        return self.select_first_from_table(*entities, **kwargs)
+        """Dispatch query to sample or all table.
+        Note: Kept for backward compatibility
+
+        Args:
+            *entities: entities to select
+            **kwargs: kwargs to pass to the query
+        """
+        return self.select_first_from_sample(*entities, **kwargs)
 
     @staticmethod
     def select_first_from_query(query: Query):
+        """Given a query object, return the first row
+
+        Args:
+            query (Query): query object
+        """
         return query.first()
 
     @staticmethod
     def select_all_from_query(query: Query):
+        """Given a query object, return all the rows
+
+        Args:
+            query (Query): query object
+        """
         return query.all()

metadata/profiler/processor/sample_data_handler.py

@@ -16,6 +16,7 @@ import traceback
 from datetime import datetime
 from functools import singledispatch
 from io import BytesIO
+from typing import Optional
 
 from metadata.clients.aws_client import AWSClient
 from metadata.generated.schema.entity.data.table import Table, TableData
@@ -24,7 +25,6 @@ from metadata.generated.schema.entity.services.connections.connectionBasicType i
 )
 from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
 from metadata.ingestion.models.custom_pydantic import ignore_type_decoder
-from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import profiler_logger
 
@@ -72,14 +72,17 @@ def _get_object_key(
     return file_name
 
 
-def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -> None:
+def upload_sample_data(
+    data: TableData,
+    entity: Table,
+    sample_storage_config: Optional[DataStorageConfig] = None,
+) -> None:
     """
     Upload Sample data to storage config
     """
     import pandas as pd  # pylint: disable=import-outside-toplevel
 
     try:
-        sample_storage_config: DataStorageConfig = profiler_interface.storage_config
         if not sample_storage_config:
             return
         # Ignore any decoding error for byte data
@@ -92,7 +95,7 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
         pq_buffer = BytesIO()
         df.to_parquet(pq_buffer)
         object_key = _get_object_key(
-            table=profiler_interface.table_entity,
+            table=entity,
             prefix=sample_storage_config.prefix,
             overwrite_data=sample_storage_config.overwriteData,
             file_path_format=sample_storage_config.filePathPattern,