openmetadata-ingestion 1.7.4.2__py3-none-any.whl → 1.8.0.0__py3-none-any.whl

metadata/profiler/processor/default.py

@@ -12,55 +12,54 @@
 """
 Default simple profiler to use
 """
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from sqlalchemy.orm import DeclarativeMeta
 
-from metadata.generated.schema.configuration.profilerConfiguration import (
-    ProfilerConfiguration,
-)
 from metadata.generated.schema.entity.data.table import ColumnProfilerConfig
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.settings.settings import Settings
 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
+from metadata.profiler.registry import MetricRegistry
 
 
 def get_default_metrics(
+    metrics_registry: Type[MetricRegistry],
     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),
+        metrics_registry.ROW_COUNT.value,
+        add_props(table=table)(metrics_registry.COLUMN_COUNT.value),
+        add_props(table=table)(metrics_registry.COLUMN_NAMES.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
+            metrics_registry.SYSTEM.value
         ),
         # Column Metrics
-        Metrics.MEDIAN.value,
-        Metrics.FIRST_QUARTILE.value,
-        Metrics.THIRD_QUARTILE.value,
-        Metrics.MEAN.value,
-        Metrics.COUNT.value,
-        Metrics.DISTINCT_COUNT.value,
-        Metrics.DISTINCT_RATIO.value,
-        Metrics.MIN.value,
-        Metrics.MAX.value,
-        Metrics.NULL_COUNT.value,
-        Metrics.NULL_RATIO.value,
-        Metrics.STDDEV.value,
-        Metrics.SUM.value,
-        Metrics.UNIQUE_COUNT.value,
-        Metrics.UNIQUE_RATIO.value,
-        Metrics.IQR.value,
-        Metrics.HISTOGRAM.value,
-        Metrics.NON_PARAMETRIC_SKEW.value,
+        metrics_registry.MEDIAN.value,
+        metrics_registry.FIRST_QUARTILE.value,
+        metrics_registry.THIRD_QUARTILE.value,
+        metrics_registry.MEAN.value,
+        metrics_registry.COUNT.value,
+        metrics_registry.DISTINCT_COUNT.value,
+        metrics_registry.DISTINCT_RATIO.value,
+        metrics_registry.MIN.value,
+        metrics_registry.MAX.value,
+        metrics_registry.NULL_COUNT.value,
+        metrics_registry.NULL_RATIO.value,
+        metrics_registry.STDDEV.value,
+        metrics_registry.SUM.value,
+        metrics_registry.UNIQUE_COUNT.value,
+        metrics_registry.UNIQUE_RATIO.value,
+        metrics_registry.IQR.value,
+        metrics_registry.HISTOGRAM.value,
+        metrics_registry.NON_PARAMETRIC_SKEW.value,
     ]
 
 
@@ -74,12 +73,14 @@ class DefaultProfiler(Profiler):
     def __init__(
         self,
         profiler_interface: ProfilerInterface,
+        metrics_registry: Type[MetricRegistry],
         include_columns: Optional[List[ColumnProfilerConfig]] = None,
         exclude_columns: Optional[List[str]] = None,
-        global_profiler_configuration: Optional[ProfilerConfiguration] = None,
+        global_profiler_configuration: Optional[Settings] = None,
         db_service=None,
     ):
         _metrics = get_default_metrics(
+            metrics_registry=metrics_registry,
             table=profiler_interface.table,
             ometa_client=profiler_interface.ometa_client,
             db_service=db_service,

metadata/profiler/processor/metric_filter.py

@@ -33,22 +33,35 @@ from metadata.profiler.metrics.core import (
     SystemMetric,
     TMetric,
 )
-from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.orm.converter.converter_registry import converter_registry
+from metadata.profiler.registry import MetricRegistry
+from metadata.utils.dependency_injector.dependency_injector import (
+    DependencyNotFoundError,
+    Inject,
+    inject,
+)
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 
 class MetricFilter:
     """Metric filter class for profiler"""
 
+    @inject
     def __init__(
         self,
         metrics: Tuple[Type[TMetric]],
         global_profiler_config: Optional[ProfilerConfiguration] = None,
         table_profiler_config: Optional[TableProfilerConfig] = None,
         column_profiler_config: Optional[List[ColumnProfilerConfig]] = None,
+        metrics_registry: Inject[Type[MetricRegistry]] = None,
     ):
+        if metrics_registry is None:
+            raise DependencyNotFoundError(
+                "MetricRegistry dependency not found. Please ensure the MetricRegistry is properly registered."
+            )
+
         self.metrics = metrics
+        self.metrics_registry = metrics_registry
         self.global_profiler_config = global_profiler_config
         self.table_profiler_config = table_profiler_config
         self.column_profiler_config = column_profiler_config
@@ -196,7 +209,7 @@ class MetricFilter:
 
         metrics = [
             Metric.value
-            for Metric in Metrics
+            for Metric in self.metrics_registry
             if Metric.value.name() in {mtrc.value for mtrc in col_dtype_config.metrics}
             and Metric.value in metrics
         ]
@@ -240,7 +253,7 @@ class MetricFilter:
 
         metrics = [
             Metric.value
-            for Metric in Metrics
+            for Metric in self.metrics_registry
             if Metric.value.name().lower() in {mtrc.lower() for mtrc in metric_names}
             and Metric.value in metrics
         ]

metadata/profiler/processor/models.py

@@ -13,20 +13,30 @@
 Models to map profiler definitions
 JSON workflows to the profiler
 """
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from pydantic import BaseModel, BeforeValidator
 from typing_extensions import Annotated
 
-from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.registry import MetricRegistry
+from metadata.utils.dependency_injector.dependency_injector import (
+    DependencyNotFoundError,
+    Inject,
+    inject,
+)
 
 
-def valid_metric(value: str):
+@inject
+def valid_metric(value: str, metrics: Inject[Type[MetricRegistry]] = None):
     """
     Validate that the input metrics are correctly named
     and can be found in the Registry
     """
-    if not Metrics.get(value.upper()):
+    if metrics is None:
+        raise DependencyNotFoundError(
+            "MetricRegistry dependency not found. Please ensure the MetricRegistry is properly registered."
+        )
+    if not metrics.get(value.upper()):
         raise ValueError(
             f"Metric name {value} is not a proper metric name from the Registry"
         )

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

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+from typing import Tuple, Type
+
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.profiler.interface.profiler_interface import ProfilerInterface
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.service_spec.service_spec import (
+    import_profiler_class,
+    import_sampler_class,
+)
+
+
+class ProfilerResolver(ABC):
+    """Abstract class for the profiler resolver"""
+
+    @staticmethod
+    @abstractmethod
+    def resolve(
+        processing_engine: str, service_type: ServiceType, source_type: str
+    ) -> Tuple[Type[SamplerInterface], Type[ProfilerInterface]]:
+        """Resolve the sampler and profiler based on the processing engine."""
+        raise NotImplementedError
+
+
+class DefaultProfilerResolver(ProfilerResolver):
+    """Default profiler resolver"""
+
+    @staticmethod
+    def resolve(
+        processing_engine: str, service_type: ServiceType, source_type: str
+    ) -> Tuple[Type[SamplerInterface], Type[ProfilerInterface]]:
+        """Resolve the sampler and profiler based on the processing engine."""
+        sampler_class = import_sampler_class(service_type, source_type=source_type)
+        profiler_class = import_profiler_class(service_type, source_type=source_type)
+        return sampler_class, profiler_class

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

@@ -14,7 +14,7 @@ Base source for the profiler used to instantiate a profiler runner with
 its interface
 """
 from copy import deepcopy
-from typing import Optional, cast
+from typing import Optional, Type, cast
 
 from metadata.generated.schema.configuration.profilerConfiguration import (
     ProfilerConfiguration,
@@ -22,10 +22,7 @@ from metadata.generated.schema.configuration.profilerConfiguration import (
 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 Table
-from metadata.generated.schema.entity.services.databaseService import (
-    DatabaseConnection,
-    DatabaseService,
-)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
@@ -36,9 +33,10 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.api.models import ProfilerProcessorConfig, TableConfig
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
-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.registry import MetricRegistry
+from metadata.profiler.source.database.base.profiler_resolver import ProfilerResolver
 from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
 from metadata.sampler.config import (
     get_config_for_table,
@@ -47,12 +45,13 @@ from metadata.sampler.config import (
 )
 from metadata.sampler.models import SampleConfig
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.dependency_injector.dependency_injector import (
+    DependencyNotFoundError,
+    Inject,
+    inject,
+)
 from metadata.utils.logger import profiler_logger
 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()
 
@@ -77,8 +76,7 @@ class ProfilerSource(ProfilerSourceInterface):
         self.ometa_client = ometa_client
         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
 
@@ -122,25 +120,34 @@ class ProfilerSource(ProfilerSourceInterface):
 
         return config_copy
 
+    @inject
     def create_profiler_interface(
         self,
         entity: Table,
         config: Optional[TableConfig],
-        profiler_config: Optional[ProfilerProcessorConfig],
-        schema_entity: Optional[DatabaseSchema],
-        database_entity: Optional[Database],
-        db_service: Optional[DatabaseService],
+        schema_entity: DatabaseSchema,
+        database_entity: Database,
+        profiler_resolver: Inject[Type[ProfilerResolver]] = None,
     ) -> ProfilerInterface:
-        """Create sqlalchemy profiler interface"""
+        """Create the appropriate profiler interface based on processing engine."""
+        if profiler_resolver is None:
+            raise DependencyNotFoundError(
+                "ProfilerResolver dependency not found. Please ensure the ProfilerResolver is properly registered."
+            )
+
+        # NOTE: For some reason I do not understand, if we instantiate this on the __init__ method, we break the
+        # autoclassification workflow. This should be fixed. There should not be an impact on AutoClassification.
+        # We have an issue to track this here: https://github.com/open-metadata/OpenMetadata/issues/21790
         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
+
+        sampler_class, profiler_class = profiler_resolver.resolve(
+            processing_engine=self.get_processing_engine(self.source_config),
+            service_type=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,
@@ -155,6 +162,8 @@ class ProfilerSource(ProfilerSourceInterface):
                 samplingMethodType=self.source_config.samplingMethodType,
                 randomizedSample=self.source_config.randomizedSample,
             ),
+            # TODO: Change this when we have the processing engine configuration implemented. Right now it does nothing.
+            processing_engine=self.get_processing_engine(self.source_config),
         )
 
         profiler_interface: ProfilerInterface = profiler_class.create(
@@ -168,28 +177,33 @@ class ProfilerSource(ProfilerSourceInterface):
         self.interface = profiler_interface
         return self.interface
 
+    @inject
     def get_profiler_runner(
-        self, entity: Table, profiler_config: ProfilerProcessorConfig
+        self,
+        entity: Table,
+        profiler_config: ProfilerProcessorConfig,
+        metrics_registry: Inject[Type[MetricRegistry]] = None,
     ) -> Profiler:
         """
         Returns the runner for the profiler
         """
+        if metrics_registry is None:
+            raise DependencyNotFoundError(
+                "MetricRegistry dependency not found. Please ensure the MetricRegistry is properly registered."
+            )
+
         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,
-            table_config,
-            profiler_config,
-            schema_entity,
-            database_entity,
-            db_service,
+            entity, table_config, schema_entity, database_entity
         )
 
         if not profiler_config.profiler:
             return DefaultProfiler(
                 profiler_interface=profiler_interface,
+                metrics_registry=metrics_registry,
                 include_columns=get_include_columns(entity, table_config),
                 exclude_columns=get_exclude_columns(entity, table_config),
                 global_profiler_configuration=self.global_profiler_configuration,
@@ -197,9 +211,10 @@ class ProfilerSource(ProfilerSourceInterface):
             )
 
         metrics = (
-            [Metrics.get(name) for name in profiler_config.profiler.metrics]
+            [metrics_registry.get(name) for name in profiler_config.profiler.metrics]
             if profiler_config.profiler.metrics
             else get_default_metrics(
+                metrics_registry=metrics_registry,
                 table=profiler_interface.table,
                 ometa_client=self.ometa_client,
                 db_service=db_service,

metadata/profiler/source/profiler_source_interface.py

@@ -16,6 +16,9 @@ Class defining the interface for the profiler source
 from abc import ABC, abstractmethod
 from typing import Optional
 
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    DatabaseServiceProfilerPipeline,
+)
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 
 
@@ -34,20 +37,12 @@ class ProfilerSourceInterface(ABC):
         """Set the interface"""
         raise NotImplementedError
 
-    @abstractmethod
-    def create_profiler_interface(
-        self,
-        entity,
-        config,
-        profiler_config,
-        schema_entity,
-        database_entity,
-        db_service,
-    ) -> ProfilerInterface:
-        """Create the profiler interface"""
-        raise NotImplementedError
-
     @abstractmethod
     def get_profiler_runner(self, entity, profiler_config):
         """Get the profiler runner"""
         raise NotImplementedError
+
+    @staticmethod
+    def get_processing_engine(config: DatabaseServiceProfilerPipeline):
+        """Get the processing engine based on the configuration."""
+        return "Native"

metadata/sampler/sampler_interface.py

@@ -13,12 +13,13 @@ Interface for sampler
 """
 import traceback
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Set, Union
+from typing import List, Optional, Set, Union
 
 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,
+    PartitionProfilerConfig,
     Table,
     TableData,
 )
@@ -65,20 +66,20 @@ class SamplerInterface(ABC):
         include_columns: Optional[List[ColumnProfilerConfig]] = None,
         exclude_columns: Optional[List[str]] = None,
         sample_config: SampleConfig = SampleConfig(),
-        partition_details: Optional[Dict] = None,
+        partition_details: Optional[PartitionProfilerConfig] = None,
         sample_query: Optional[str] = None,
-        storage_config: DataStorageConfig = None,
+        storage_config: Optional[DataStorageConfig] = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
         **__,
     ):
         self.ometa_client = ometa_client
         self._sample = None
-        self._columns: Optional[List[SQALikeColumn]] = None
+        self._columns: List[SQALikeColumn] = []
         self.sample_config = sample_config
 
         self.entity = entity
-        self.include_columns = include_columns
-        self.exclude_columns = exclude_columns
+        self.include_columns = include_columns or []
+        self.exclude_columns = exclude_columns or []
         self.sample_query = sample_query
         self.sample_limit = sample_data_count
         self.partition_details = partition_details
@@ -99,7 +100,7 @@ class SamplerInterface(ABC):
         table_config: Optional[TableConfig] = None,
         storage_config: Optional[DataStorageConfig] = None,
         default_sample_config: Optional[SampleConfig] = None,
-        default_sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        default_sample_data_count: int = SAMPLE_DATA_DEFAULT_COUNT,
         **kwargs,
     ) -> "SamplerInterface":
         """Create sampler"""
@@ -165,16 +166,20 @@ class SamplerInterface(ABC):
 
         return self._columns
 
-    def _get_excluded_columns(self) -> Optional[Set[str]]:
+    def _get_excluded_columns(self) -> Set[str]:
         """Get excluded  columns for table being profiled"""
         if self.exclude_columns:
             return set(self.exclude_columns)
         return set()
 
-    def _get_included_columns(self) -> Optional[Set[str]]:
+    def _get_included_columns(self) -> Set[str]:
         """Get include columns for table being profiled"""
         if self.include_columns:
-            return {include_col.columnName for include_col in self.include_columns}
+            return {
+                include_col.columnName
+                for include_col in self.include_columns
+                if include_col.columnName
+            }
         return set()
 
     @property

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -73,6 +73,13 @@ class BigQuerySampler(SQASampler):
             **kwargs,
         )
         self.raw_dataset_type: Optional[TableType] = entity.tableType
+        if self._table.__table__.schema and self.entity.database.name:
+            self._table.__table__.schema = (
+                f"{self.entity.database.name}.{self._table.__table__.schema}"
+            )
+            self._table.__table_args__[
+                "schema"
+            ] = f"{self.entity.database.name}.{self._table.__table_args__['schema']}"
 
     def set_tablesample(self, selectable: SqaTable):
         """Set the TABLESAMPLE clause for BigQuery