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

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

metadata/utils/service_spec/service_spec.py

@@ -2,16 +2,19 @@
 Manifests are used to store class information
 """
 
-from typing import Optional, Type, cast
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Type, cast
 
 from pydantic import model_validator
 
 from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.ingestion.api.steps import Source
+from metadata.ingestion.connections.connection import BaseConnection
 from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.dependency_injector.dependency_injector import Inject, inject
 from metadata.utils.importer import (
     TYPE_SEPARATOR,
     DynamicImportException,
@@ -24,6 +27,14 @@ from metadata.utils.logger import utils_logger
 logger = utils_logger()
 
 
+class SourceLoader(ABC):
+    @abstractmethod
+    def __call__(
+        self, service_type: ServiceType, source_type: str, from_: str
+    ) -> Type[Any]:
+        """Load the service spec for a given service type and source type."""
+
+
 class BaseSpec(BaseModel):
     """
     # The OpenMetadata Ingestion Service Specification (Spec)
@@ -53,6 +64,7 @@ class BaseSpec(BaseModel):
     usage_source_class: Optional[str] = None
     sampler_class: Optional[str] = None
     data_diff: Optional[str] = None
+    connection_class: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -65,8 +77,13 @@ class BaseSpec(BaseModel):
         return values
 
     @classmethod
+    @inject
     def get_for_source(
-        cls, service_type: ServiceType, source_type: str, from_: str = "ingestion"
+        cls,
+        service_type: ServiceType,
+        source_type: str,
+        from_: str = "ingestion",
+        source_loader: Inject[SourceLoader] = None,
     ) -> "BaseSpec":
         """Retrieves the manifest for a given source type. If it does not exist will attempt to retrieve
         a default manifest for the service type.
@@ -79,14 +96,26 @@ class BaseSpec(BaseModel):
         Returns:
             BaseSpec: The manifest for the source type.
         """
-        return cls.model_validate(
-            import_from_module(
-                "metadata.{}.source.{}.{}.{}.ServiceSpec".format(  # pylint: disable=C0209
-                    from_,
-                    service_type.name.lower(),
-                    get_module_dir(source_type),
-                    "service_spec",
-                )
+        if not source_loader:
+            raise ValueError("Source loader is required")
+
+        return cls.model_validate(source_loader(service_type, source_type, from_))
+
+
+class DefaultSourceLoader(SourceLoader):
+    def __call__(
+        self,
+        service_type: ServiceType,
+        source_type: str,
+        from_: str = "ingestion",
+    ) -> Type[Any]:
+        """Default implementation for loading service specifications."""
+        return import_from_module(
+            "metadata.{}.source.{}.{}.{}.ServiceSpec".format(  # pylint: disable=C0209
+                from_,
+                service_type.name.lower(),
+                get_module_dir(source_type),
+                "service_spec",
             )
         )
 
@@ -110,6 +139,10 @@ def import_profiler_class(
     service_type: ServiceType, source_type: str
 ) -> Type[ProfilerInterface]:
     class_path = BaseSpec.get_for_source(service_type, source_type).profiler_class
+    if not class_path:
+        raise ValueError(
+            f"Profiler class not found for service type {service_type} and source type {source_type}"
+        )
     return cast(Type[ProfilerInterface], import_from_module(class_path))
 
 
@@ -127,6 +160,10 @@ def import_test_suite_class(
             ).test_suite_class
         else:
             raise
+    if not class_path:
+        raise ValueError(
+            f"Test suite class not found for service type {service_type} and source type {source_type}"
+        )
     return cast(Type[TestSuiteInterface], import_from_module(class_path))
 
 
@@ -144,4 +181,23 @@ def import_sampler_class(
             ).sampler_class
         else:
             raise
+    if not class_path:
+        raise ValueError(
+            f"Sampler class not found for service type {service_type} and source type {source_type}"
+        )
     return cast(Type[SamplerInterface], import_from_module(class_path))
+
+
+def import_connection_class(
+    service_type: ServiceType,
+    source_type: str,
+) -> Type[BaseConnection]:
+    """
+    Import the connection class for a given service type and source type.
+    """
+    class_path = BaseSpec.get_for_source(service_type, source_type).connection_class
+    if not class_path:
+        raise ValueError(
+            f"Connection class not found for service type {service_type} and source type {source_type}"
+        )
+    return cast(Type[BaseConnection], import_from_module(class_path))

metadata/workflow/profiler.py

@@ -13,10 +13,7 @@ Workflow definition for the profiler
 """
 
 from metadata.ingestion.api.steps import Processor, Sink
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
-from metadata.ingestion.source.connections import get_test_connection_fn
+from metadata.ingestion.source.connections import test_connection_common
 from metadata.profiler.processor.processor import ProfilerProcessor
 from metadata.profiler.source.metadata import OpenMetadataSource
 from metadata.profiler.source.metadata_ext import OpenMetadataSourceExt
@@ -75,9 +72,7 @@ class ProfilerWorkflow(IngestionWorkflow):
             service_config = self.config.source.serviceConnection.root.config
             conn = get_ssl_connection(service_config)
 
-            test_connection_fn = get_test_connection_fn(service_config)
-            result = test_connection_fn(self.metadata, conn, service_config)
-            raise_test_connection_exception(result)
+            test_connection_common(self.metadata, conn, service_config)
 
         return main(self)