openmetadata-ingestion 1.7.5.3__py3-none-any.whl → 1.8.1.0__py3-none-any.whl

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

metadata/utils/dependency_injector/dependency_injector.py

@@ -0,0 +1,403 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Dependency injection utilities for OpenMetadata.
+
+This module provides a type-safe dependency injection system that uses Python's type hints
+to automatically inject dependencies into functions and methods.
+
+Example:
+    ```python
+    from typing import Annotated
+    from metadata.utils.dependency_injector import Inject, inject, DependencyContainer
+
+    class Database:
+        def __init__(self, connection_string: str):
+            self.connection_string = connection_string
+
+    # Register a dependency
+    container = DependencyContainer[Callable]()
+    container.register(Database, lambda: Database("postgresql://localhost:5432"))
+
+    # Use dependency injection
+    @inject
+    def get_user(user_id: int, db: Inject[Database]) -> dict:
+        return db.query(f"SELECT * FROM users WHERE id = {user_id}")
+    ```
+"""
+from functools import wraps
+from threading import RLock
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
+
+T = TypeVar("T")
+
+
+class DependencyInjectionError(Exception):
+    """Base exception for dependency injection errors."""
+
+    pass
+
+
+class DependencyNotFoundError(DependencyInjectionError):
+    """Raised when a required dependency is not found in the container."""
+
+    pass
+
+
+class InvalidInjectionTypeError(DependencyInjectionError):
+    """Raised when an invalid injection type is used."""
+
+    pass
+
+
+if TYPE_CHECKING:
+    Inject = Annotated[Union[T, None], "Inject Marker"]
+else:
+
+    class Inject(Generic[T]):
+        """
+        Type for dependency injection that uses types as keys.
+
+        This type is used to mark parameters that should be automatically injected
+        by the dependency container. Injection is always treated as Optioonal. It can be overriden.
+
+        Type Parameters:
+            T: The type of the dependency to inject
+
+        Example:
+            ```python
+            @inject
+            def my_function(db: Inject[Database]):
+                # db will be automatically injected
+                pass
+            ```
+        """
+
+
+class DependencyContainer:
+    """
+    Thread-safe singleton container for managing dependencies.
+
+    This container uses RLock to support reentrant locking, allowing the same thread
+    to acquire the lock multiple times. It maintains two dictionaries:
+    - _dependencies: The original registered dependencies
+    - _overrides: Temporary overrides that take precedence over original dependencies
+
+    Type Parameters:
+        T: The base type for all dependencies in this container
+
+    Example:
+        ```python
+        container = DependencyContainer[Callable]()
+        container.register(Database, lambda: Database("postgresql://localhost:5432"))
+        container.override(Database, lambda: Database("postgresql://test:5432"))
+        ```
+    """
+
+    _instance: Optional["DependencyContainer"] = None
+    _lock = RLock()
+    _dependencies: Dict[str, Callable[[], Any]] = {}
+    _overrides: Dict[str, Callable[[], Any]] = {}
+
+    def __new__(cls) -> "DependencyContainer":
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def get_key(self, dependency_type: Type[Any]) -> str:
+        """
+        Get the key for a dependency.
+        """
+        if get_origin(dependency_type) is type:
+            inner_type = get_args(dependency_type)[0]
+            return f"Type[{inner_type.__name__}]"
+        return dependency_type.__name__
+
+    def register(
+        self, dependency_type: Type[Any], dependency: Callable[[], Any]
+    ) -> None:
+        """
+        Register a dependency with the container.
+
+        Args:
+            dependency_type: The type of the dependency to register
+            dependency: The dependency to register (usually a factory function)
+
+        Example:
+            ```python
+            container.register(Database, lambda: Database("postgresql://localhost:5432"))
+            container.register(Type[Metrics], lambda: Metrics)  # For registering types themselves
+            ```
+        """
+        with self._lock:
+            self._dependencies[self.get_key(dependency_type)] = dependency
+
+    def override(
+        self, dependency_type: Type[Any], dependency: Callable[[], Any]
+    ) -> None:
+        """
+        Override a dependency with a new implementation.
+
+        The override takes precedence over the original dependency and is useful
+        for testing or temporary changes.
+
+        Args:
+            dependency_type: The type of the dependency to override
+            dependency: The new dependency implementation
+
+        Example:
+            ```python
+            container.override(Database, lambda: Database("postgresql://test:5432"))
+            ```
+        """
+        with self._lock:
+            self._overrides[self.get_key(dependency_type)] = dependency
+
+    def remove_override(self, dependency_type: Type[T]) -> None:
+        """
+        Remove an override for a dependency.
+
+        Args:
+            dependency_type: The type of the dependency override to remove
+
+        Example:
+            ```python
+            container.remove_override(Database)
+            ```
+        """
+        with self._lock:
+            self._overrides.pop(self.get_key(dependency_type), None)
+
+    def get(self, dependency_type: Type[Any]) -> Optional[Any]:
+        """
+        Get a dependency from the container.
+
+        Checks overrides first, then falls back to original dependencies.
+
+        Args:
+            dependency_type: The type of the dependency to retrieve
+
+        Returns:
+            The dependency if found, None otherwise
+
+        Example:
+            ```python
+            db_factory = container.get(Database)
+            if db_factory:
+                db = db_factory()
+            ```
+        """
+        with self._lock:
+            factory = self._overrides.get(
+                self.get_key(dependency_type)
+            ) or self._dependencies.get(self.get_key(dependency_type))
+            if factory is None:
+                return None
+            return factory()
+
+    def clear(self) -> None:
+        """
+        Clear all dependencies and overrides.
+
+        Example:
+            ```python
+            container.clear()  # Remove all registered dependencies and overrides
+            ```
+        """
+        with self._lock:
+            self._dependencies.clear()
+            self._overrides.clear()
+
+    def has(self, dependency_type: Type[T]) -> bool:
+        """
+        Check if a dependency exists in the container.
+
+        Args:
+            dependency_type: The type to check
+
+        Returns:
+            True if the dependency exists, False otherwise
+
+        Example:
+            ```python
+            if container.has(Database):
+                print("Database dependency is registered")
+            ```"""
+        with self._lock:
+            return (
+                self.get_key(dependency_type) in self._overrides
+                or self.get_key(dependency_type) in self._dependencies
+            )
+
+
+def inject(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Decorator to inject dependencies based on type hints.
+
+    This decorator automatically injects dependencies into function parameters
+    based on their type hints. It uses types as keys for dependency lookup and
+    allows explicit injection by passing dependencies as keyword arguments.
+
+    Args:
+        func: The function to inject dependencies into
+
+    Returns:
+        A function with dependencies injected
+
+    Example:
+        ```python
+        @inject
+        def get_user(user_id: int, db: Inject[Database]) -> dict:
+            return db.query(f"SELECT * FROM users WHERE id = {user_id}")
+
+        # Dependencies can also be passed explicitly
+        get_user(user_id=1, db=Database("postgresql://localhost:5432"))
+        ```
+    """
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        container = DependencyContainer()
+        type_hints = get_type_hints(func, include_extras=True)
+
+        for param_name, param_type in type_hints.items():
+            # Skip if parameter is already provided explicitly
+            if param_name in kwargs:
+                continue
+
+            # Check if it's an Inject type
+            if is_inject_type(param_type):
+                dependency_type = extract_inject_arg(param_type)
+                dependency = container.get(dependency_type)
+                if dependency is None:
+                    raise DependencyNotFoundError(
+                        f"Dependency of type {dependency_type} not found in container. "
+                        f"Make sure to register it using container.register({dependency_type.__name__}, ...)"
+                    )
+                kwargs[param_name] = dependency
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def is_inject_type(tp: Any) -> bool:
+    """
+    Check if a type is an Inject type or Optional[Inject].
+
+    Args:
+        tp: The type to check
+
+    Returns:
+        True if the type is Inject or Optional[Inject], False otherwise
+    """
+    origin = get_origin(tp)
+    if origin is Inject:
+        return True
+    if origin is Union:
+        args = get_args(tp)
+        return any(get_origin(arg) is Inject for arg in args)
+    return False
+
+
+def extract_inject_arg(tp: Any) -> Any:
+    """
+    Extract the type argument from an Inject type.
+
+    Args:
+        tp: The type to extract from
+
+    Returns:
+        The type argument from the Inject type
+
+    Raises:
+        InvalidInjectionTypeError: If the type is not Inject or Optional[Inject]
+    """
+    origin = get_origin(tp)
+    if origin is Inject:
+        return get_args(tp)[0]
+    if origin is Union:
+        for arg in get_args(tp):
+            if get_origin(arg) is Inject:
+                return get_args(arg)[0]
+    raise InvalidInjectionTypeError(
+        f"Type {tp} is not Inject or Optional[Inject]. "
+        f"Use Annotated[YourType, 'Inject'] to mark a parameter for injection."
+    )
+
+
+def inject_class_attributes(cls: Type[Any]) -> Type[Any]:
+    """
+    Decorator to inject dependencies into class-level (static) attributes based on type hints.
+
+    This decorator automatically injects dependencies into class attributes
+    based on their type hints. The dependencies are shared across all instances
+    of the class.
+
+    Args:
+        cls: The class to inject dependencies into
+
+    Returns:
+        A class with dependencies injected into its class-level attributes
+
+    Example:
+        ```python
+        @inject_class_attributes
+        class UserService:
+            db: ClassVar[Inject[Database]]
+            cache: ClassVar[Inject[Cache]]
+
+            @classmethod
+            def get_user(cls, user_id: int) -> dict:
+                return cls.db.query(f"SELECT * FROM users WHERE id = {user_id}")
+        ```
+    """
+    container = DependencyContainer()
+    type_hints = get_type_hints(cls, include_extras=True)
+
+    # Inject dependencies into class attributes
+    for attr_name, attr_type in type_hints.items():
+        # Skip if attribute is already set
+        if hasattr(cls, attr_name):
+            continue
+
+        # Check if it's an Inject type
+        if is_inject_type(attr_type):
+            dependency_type = extract_inject_arg(attr_type)
+            dependency = container.get(dependency_type)
+            if dependency is None:
+                raise DependencyNotFoundError(
+                    f"Dependency of type {dependency_type} not found in container. "
+                    f"Make sure to register it using container.register({dependency_type.__name__}, ...)"
+                )
+            setattr(cls, attr_name, dependency)
+
+    return cls