openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

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
@@ -317,7 +308,6 @@ class Profiler(Generic[TMetric]):
                 col,
                 metric,
                 current_col_results,
-                table=self.table,
             )
 
     def _prepare_table_metrics(self) -> List:
@@ -495,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)
@@ -512,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
@@ -603,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"),

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/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,