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

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,71 @@ 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
+
+    @property
+    def session(self):
+        """Table name attribute access"""
+        return self._session
 
     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 +115,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 +134,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 +153,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 +173,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 +208,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,

metadata/profiler/source/database/base/profiler_source.py

@@ -14,19 +14,14 @@ Base source for the profiler used to instantiate a profiler runner with
 its interface
 """
 from copy import deepcopy
-from typing import List, Optional, Tuple, Type, cast
-
-from sqlalchemy import MetaData
+from typing import Optional, cast
 
 from metadata.generated.schema.configuration.profilerConfiguration import (
     ProfilerConfiguration,
 )
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
-from metadata.generated.schema.entity.data.table import ColumnProfilerConfig, Table
-from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
-    DatalakeConnection,
-)
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseService,
@@ -45,11 +40,19 @@ from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.core import Profiler
 from metadata.profiler.processor.default import DefaultProfiler, get_default_metrics
 from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
-from metadata.utils.importer import import_from_module
+from metadata.sampler.config import (
+    get_config_for_table,
+    get_exclude_columns,
+    get_include_columns,
+)
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.logger import profiler_logger
-from metadata.utils.service_spec.service_spec import BaseSpec
-
-NON_SQA_DATABASE_CONNECTIONS = (DatalakeConnection,)
+from metadata.utils.profiler_utils import get_context_entities
+from metadata.utils.service_spec.service_spec import (
+    import_profiler_class,
+    import_sampler_class,
+)
 
 logger = profiler_logger()
 
@@ -62,21 +65,21 @@ class ProfilerSource(ProfilerSourceInterface):
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
-        database: DatabaseService,
+        database: Database,
         ometa_client: OpenMetadata,
         global_profiler_configuration: ProfilerConfiguration,
     ):
+        self.config = config
         self.service_conn_config = self._copy_service_config(config, database)
-        self.source_config = DatabaseServiceProfilerPipeline.model_validate(
-            config.source.sourceConfig.config
-        )
         self.profiler_config = ProfilerProcessorConfig.model_validate(
             config.processor.model_dump().get("config")
         )
         self.ometa_client = ometa_client
-        self.profiler_interface_type: str = config.source.type.lower()
-        self.sqa_metadata = self._set_sqa_metadata()
+        self._interface_type: str = config.source.type.lower()
         self._interface = None
+        # We define this in create_profiler_interface to help us reuse
+        # this method for the sampler, which does not have a DatabaseServiceProfilerPipeline
+        self.source_config = None
         self.global_profiler_configuration = global_profiler_configuration
 
     @property
@@ -91,68 +94,8 @@ class ProfilerSource(ProfilerSourceInterface):
         """Set the interface"""
         self._interface = interface
 
-    def _set_sqa_metadata(self):
-        """Set sqlalchemy metadata"""
-        if not isinstance(self.service_conn_config, NON_SQA_DATABASE_CONNECTIONS):
-            return MetaData()
-        return None
-
-    @staticmethod
-    def get_config_for_table(entity: Table, profiler_config) -> Optional[TableConfig]:
-        """Get config for a specific entity
-
-        Args:
-            entity: table entity
-        """
-        for table_config in profiler_config.tableConfig or []:
-            if table_config.fullyQualifiedName.root == entity.fullyQualifiedName.root:
-                return table_config
-
-        for schema_config in profiler_config.schemaConfig or []:
-            if (
-                schema_config.fullyQualifiedName.root
-                == entity.databaseSchema.fullyQualifiedName
-            ):
-                return TableConfig.from_database_and_schema_config(
-                    schema_config, entity.fullyQualifiedName.root
-                )
-        for database_config in profiler_config.databaseConfig or []:
-            if (
-                database_config.fullyQualifiedName.root
-                == entity.database.fullyQualifiedName
-            ):
-                return TableConfig.from_database_and_schema_config(
-                    database_config, entity.fullyQualifiedName.root
-                )
-
-        return None
-
-    def _get_include_columns(
-        self, entity, entity_config: Optional[TableConfig]
-    ) -> Optional[List[ColumnProfilerConfig]]:
-        """get included columns"""
-        if entity_config and entity_config.columnConfig:
-            return entity_config.columnConfig.includeColumns
-
-        if entity.tableProfilerConfig:
-            return entity.tableProfilerConfig.includeColumns
-
-        return None
-
-    def _get_exclude_columns(
-        self, entity, entity_config: Optional[TableConfig]
-    ) -> Optional[List[str]]:
-        """get included columns"""
-        if entity_config and entity_config.columnConfig:
-            return entity_config.columnConfig.excludeColumns
-
-        if entity.tableProfilerConfig:
-            return entity.tableProfilerConfig.excludeColumns
-
-        return None
-
     def _copy_service_config(
-        self, config: OpenMetadataWorkflowConfig, database: DatabaseService
+        self, config: OpenMetadataWorkflowConfig, database: Database
     ) -> DatabaseConnection:
         """Make a copy of the service config and update the database name
 
@@ -189,75 +132,50 @@ class ProfilerSource(ProfilerSourceInterface):
         db_service: Optional[DatabaseService],
     ) -> ProfilerInterface:
         """Create sqlalchemy profiler interface"""
-        profiler_class = self.import_profiler_class(
-            ServiceType.Database, source_type=self.profiler_interface_type
+        self.source_config = DatabaseServiceProfilerPipeline.model_validate(
+            self.config.source.sourceConfig.config
+        )
+        profiler_class = import_profiler_class(
+            ServiceType.Database, source_type=self._interface_type
+        )
+        sampler_class = import_sampler_class(
+            ServiceType.Database, source_type=self._interface_type
         )
+        # This is shared between the sampler and profiler interfaces
+        sampler_interface: SamplerInterface = sampler_class.create(
+            service_connection_config=self.service_conn_config,
+            ometa_client=self.ometa_client,
+            entity=entity,
+            schema_entity=schema_entity,
+            database_entity=database_entity,
+            table_config=config,
+            default_sample_config=SampleConfig(
+                profile_sample=self.source_config.profileSample,
+                profile_sample_type=self.source_config.profileSampleType,
+                sampling_method_type=self.source_config.samplingMethodType,
+            ),
+        )
+
         profiler_interface: ProfilerInterface = profiler_class.create(
-            entity,
-            schema_entity,
-            database_entity,
-            db_service,
-            config,
-            profiler_config,
-            self.source_config,
-            self.service_conn_config,
-            self.ometa_client,
-            sqa_metadata=self.sqa_metadata,
+            entity=entity,
+            source_config=self.source_config,
+            service_connection_config=self.service_conn_config,
+            sampler=sampler_interface,
+            ometa_client=self.ometa_client,
         )  # type: ignore
 
         self.interface = profiler_interface
         return self.interface
 
-    def import_profiler_class(
-        self, service_type: ServiceType, source_type: str
-    ) -> Type[ProfilerInterface]:
-        class_path = BaseSpec.get_for_source(service_type, source_type).profiler_class
-        return cast(Type[ProfilerInterface], import_from_module(class_path))
-
-    def _get_context_entities(
-        self, entity: Table
-    ) -> Tuple[DatabaseSchema, Database, DatabaseService]:
-        schema_entity = None
-        database_entity = None
-        db_service = None
-
-        if entity.databaseSchema:
-            schema_entity_list = self.ometa_client.es_search_from_fqn(
-                entity_type=DatabaseSchema,
-                fqn_search_string=entity.databaseSchema.fullyQualifiedName,
-                fields="databaseSchemaProfilerConfig",
-            )
-            if schema_entity_list:
-                schema_entity = schema_entity_list[0]
-
-        if entity.database:
-            database_entity_list = self.ometa_client.es_search_from_fqn(
-                entity_type=Database,
-                fqn_search_string=entity.database.fullyQualifiedName,
-                fields="databaseProfilerConfig",
-            )
-            if database_entity_list:
-                database_entity = database_entity_list[0]
-
-        if entity.service:
-            db_service_list = self.ometa_client.es_search_from_fqn(
-                entity_type=DatabaseService,
-                fqn_search_string=entity.service.fullyQualifiedName,
-            )
-            if db_service_list:
-                db_service = db_service_list[0]
-
-        return schema_entity, database_entity, db_service
-
     def get_profiler_runner(
         self, entity: Table, profiler_config: ProfilerProcessorConfig
     ) -> Profiler:
         """
         Returns the runner for the profiler
         """
-        table_config = self.get_config_for_table(entity, profiler_config)
-        schema_entity, database_entity, db_service = self._get_context_entities(
-            entity=entity
+        table_config = get_config_for_table(entity, profiler_config)
+        schema_entity, database_entity, db_service = get_context_entities(
+            entity=entity, metadata=self.ometa_client
         )
         profiler_interface = self.create_profiler_interface(
             entity,
@@ -271,8 +189,8 @@ class ProfilerSource(ProfilerSourceInterface):
         if not profiler_config.profiler:
             return DefaultProfiler(
                 profiler_interface=profiler_interface,
-                include_columns=self._get_include_columns(entity, table_config),
-                exclude_columns=self._get_exclude_columns(entity, table_config),
+                include_columns=get_include_columns(entity, table_config),
+                exclude_columns=get_exclude_columns(entity, table_config),
                 global_profiler_configuration=self.global_profiler_configuration,
                 db_service=db_service,
             )
@@ -290,7 +208,7 @@ class ProfilerSource(ProfilerSourceInterface):
         return Profiler(
             *metrics,  # type: ignore
             profiler_interface=profiler_interface,
-            include_columns=self._get_include_columns(entity, table_config),
-            exclude_columns=self._get_exclude_columns(entity, table_config),
+            include_columns=get_include_columns(entity, table_config),
+            exclude_columns=get_exclude_columns(entity, table_config),
             global_profiler_configuration=self.global_profiler_configuration,
         )

metadata/profiler/source/database/bigquery/profiler_source.py

@@ -15,10 +15,10 @@ Bigquery Profiler source
 
 from copy import deepcopy
 
+from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
@@ -34,14 +34,14 @@ class BigQueryProfilerSource(ProfilerSource):
     """override the base profiler source to handle BigQuery specific connection configs"""
 
     def _copy_service_config(
-        self, config: OpenMetadataWorkflowConfig, database: DatabaseService
+        self, config: OpenMetadataWorkflowConfig, database: Database
     ) -> BigQueryConnection:
         """Make a copy of the database connection config. If MultiProjectId is used, replace it
         with SingleProjectId with the database name being profiled. We iterate over all non filtered
         database in workflow.py `def execute`.
 
         Args:
-            database (DatabaseService): a database entity
+            database (Database): a database entity
 
         Returns:
             DatabaseConnection

metadata/profiler/source/database/databricks/profiler_source.py

@@ -1,9 +1,8 @@
 """Extend the ProfilerSource class to add support for Databricks is_disconnect SQA method"""
-
 from metadata.generated.schema.configuration.profilerConfiguration import (
     ProfilerConfiguration,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
@@ -25,7 +24,7 @@ class DataBricksProfilerSource(ProfilerSource):
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
-        database: DatabaseService,
+        database: Database,
         ometa_client: OpenMetadata,
         global_profiler_config: ProfilerConfiguration,
     ):