sqlalchemy-foundation-kit 0.0.0__py3-none-any.whl

sqlalchemy_foundation_kit/__init__.py

@@ -0,0 +1,119 @@
+"""Foundation layer for SQLAlchemy-based services with UoW, session management, and observability."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+# Base ORM
+from .base import (
+    DB_NAMING_CONVENTION,
+    Base,
+    BaseTable,
+    DatetimeColumnsMixin,
+    GenericJSONDict,
+    PoolClassStr,
+    PoolRegistry,
+    PydanticJSONB,
+    UnConstrainedEnum,
+    build_engine_kwargs,
+    configure_orjson_serialization,
+    load_orm_metadata,
+    register_pool_class,
+    resolve_pool_class,
+)
+
+# Config
+from .config import (
+    ConnectionSettingsProtocol,
+    PoolSettingsProtocol,
+    PostgresSettingsProtocol,
+    QuerySettingsProtocol,
+)
+
+# Protocols
+from .protocols import (
+    CheckoutRecorder,
+    ErrorRecorder,
+    PoolStatsRecorder,
+    PostgresMetricsProtocol,
+)
+
+# Session Management
+from .session import (
+    DEFAULT_HEALTHCHECK_QUERY,
+    DEFAULT_RETRY_CONFIG,
+    AsyncCConnection,
+    AsyncSessionManager,
+    AsyncSessionManagerBuilder,
+    RetryConfig,
+    create_async_session_manager,
+    retry_async_connection,
+    try_advisory_xact_lock,
+)
+
+# Unit of Work
+from .uow import (
+    AsyncSQLAlchemyUnitOfWork,
+    AsyncSQLAlchemyUowTransaction,
+    AsyncUnitOfWork,
+    AsyncUowTransaction,
+    IsolationLevel,
+    PostgresAdvisoryLockMixin,
+    SupportsAdvisoryLock,
+)
+
+try:
+    __version__ = version("sqlalchemy-foundation-kit")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.1.0"
+
+__all__ = [  # noqa: RUF022
+    # Base ORM
+    "DB_NAMING_CONVENTION",
+    "Base",
+    "BaseTable",
+    "DatetimeColumnsMixin",
+    "GenericJSONDict",
+    "PoolClassStr",
+    "PoolRegistry",
+    "PydanticJSONB",
+    "UnConstrainedEnum",
+    "build_engine_kwargs",
+    "configure_orjson_serialization",
+    "load_orm_metadata",
+    "register_pool_class",
+    "resolve_pool_class",
+    # Config
+    "ConnectionSettingsProtocol",
+    "PoolSettingsProtocol",
+    "PostgresSettingsProtocol",
+    "QuerySettingsProtocol",
+    # Protocols
+    "CheckoutRecorder",
+    "ErrorRecorder",
+    "PoolStatsRecorder",
+    "PostgresMetricsProtocol",
+    # Session Management
+    "DEFAULT_HEALTHCHECK_QUERY",
+    "DEFAULT_RETRY_CONFIG",
+    "AsyncCConnection",
+    "AsyncSessionManager",
+    "AsyncSessionManagerBuilder",
+    "RetryConfig",
+    "create_async_session_manager",
+    "retry_async_connection",
+    "try_advisory_xact_lock",
+    # Unit of Work
+    "AsyncSQLAlchemyUnitOfWork",
+    "AsyncSQLAlchemyUowTransaction",
+    "AsyncUnitOfWork",
+    "AsyncUowTransaction",
+    "IsolationLevel",
+    "PostgresAdvisoryLockMixin",
+    "SupportsAdvisoryLock",
+    # Version
+    "__version__",
+]
+
+# Note: contrib modules are available via:
+# - from sqlalchemy_foundation_kit.contrib.settings import BasePostgresConfig
+# - from sqlalchemy_foundation_kit.contrib.metrics import PostgresMetrics
+# - from sqlalchemy_foundation_kit.contrib.di import AsyncDatabaseProvider, AsyncUnitOfWorkProvider

sqlalchemy_foundation_kit/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.0.0"

sqlalchemy_foundation_kit/_typing.py

@@ -0,0 +1,28 @@
+"""Shared type variables for the library.
+
+Centralized location for all TypeVars to avoid duplication and ensure consistency.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+    from .uow.protocols import AsyncUowTransaction
+
+__all__ = [
+    "SessionT",
+    "T",
+    "T_co",
+]
+
+# Covariant TypeVar for UoW protocols (used in Protocol definitions)
+T_co = TypeVar("T_co", bound="AsyncUowTransaction", covariant=True)
+
+# Invariant TypeVar for UoW implementations
+T = TypeVar("T", bound="AsyncUowTransaction")
+
+# Session TypeVar for AsyncSessionManager and AsyncSessionManagerBuilder
+SessionT = TypeVar("SessionT", bound="AsyncSession")

sqlalchemy_foundation_kit/base/__init__.py

@@ -0,0 +1,46 @@
+"""Base ORM models and utilities.
+
+Public API for base functionality. Import directly from submodules for clarity:
+    - base.engine - Engine configuration (build_engine_kwargs, resolve_pool_class)
+    - base.serialization - JSON serialization (configure_orjson_serialization)
+    - base.metadata - Metadata loading (load_orm_metadata)
+    - base.models - ORM base classes (Base, BaseTable, mixins)
+    - base.types - Custom SQLAlchemy types (PydanticJSONB)
+"""
+
+from __future__ import annotations
+
+from .engine import (
+    PoolClassStr,
+    PoolRegistry,
+    build_engine_kwargs,
+    register_pool_class,
+    resolve_pool_class,
+)
+from .metadata import load_orm_metadata
+from .models import (
+    DB_NAMING_CONVENTION,
+    Base,
+    BaseTable,
+    DatetimeColumnsMixin,
+    UnConstrainedEnum,
+)
+from .serialization import configure_orjson_serialization
+from .types import GenericJSONDict, PydanticJSONB
+
+__all__ = [
+    "DB_NAMING_CONVENTION",
+    "Base",
+    "BaseTable",
+    "DatetimeColumnsMixin",
+    "GenericJSONDict",
+    "PoolClassStr",
+    "PoolRegistry",
+    "PydanticJSONB",
+    "UnConstrainedEnum",
+    "build_engine_kwargs",
+    "configure_orjson_serialization",
+    "load_orm_metadata",
+    "register_pool_class",
+    "resolve_pool_class",
+]

sqlalchemy_foundation_kit/base/_optional.py

@@ -0,0 +1,37 @@
+"""Utilities for handling optional dependencies."""
+
+from __future__ import annotations
+
+import importlib
+import types
+
+
+def require_optional(module_name: str, extra_name: str) -> types.ModuleType:
+    """Import an optional dependency or raise a helpful error.
+
+    Args:
+        module_name: Name of the module to import (e.g., "orjson", "opentelemetry").
+        extra_name: Name of the pip extra that provides this dependency (e.g., "json", "telemetry").
+
+    Returns:
+        The imported module.
+
+    Raises:
+        ImportError: If the module is not installed, with installation instructions.
+
+    Examples:
+        >>> orjson = require_optional("orjson", "json")
+        >>> from opentelemetry import trace
+        # or
+        >>> otel = require_optional("opentelemetry", "telemetry")
+    """
+    try:
+        return importlib.import_module(module_name)
+    except ImportError as e:
+        raise ImportError(
+            f"{module_name} is required for this functionality. "
+            f"Install it with: pip install 'sqlalchemy-foundation-kit[{extra_name}]'"
+        ) from e
+
+
+__all__ = ["require_optional"]

sqlalchemy_foundation_kit/base/engine.py

@@ -0,0 +1,256 @@
+"""SQLAlchemy engine configuration utilities."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar, Literal
+
+from sqlalchemy.pool import (
+    AsyncAdaptedQueuePool,
+    FallbackAsyncAdaptedQueuePool,
+    NullPool,
+    QueuePool,
+    SingletonThreadPool,
+    StaticPool,
+)
+
+if TYPE_CHECKING:
+    from ..config import PoolSettingsProtocol
+
+PoolClassStr = Literal[
+    "null",
+    "queue",
+    "singleton_thread",
+    "async_adapted_queue",
+    "fallback_async_adapted_queue",
+    "static",
+]
+
+
+class PoolRegistry:
+    """Registry for SQLAlchemy pool classes.
+
+    Provides a centralized registry for pool classes that follows the Open/Closed principle:
+    - Open for extension: custom pools can be registered via :meth:`register`
+    - Closed for modification: built-in pools are immutable
+
+    This design allows library users to register custom pool implementations without
+    modifying library code.
+
+    Examples:
+        Register a custom pool class:
+            >>> class MyCustomPool(QueuePool):
+            ...     pass
+            >>> PoolRegistry.register("custom", MyCustomPool)
+            >>> pool = PoolRegistry.resolve("custom")
+
+        Override built-in pool (not recommended, but possible):
+            >>> PoolRegistry.register("queue", MyCustomQueuePool, override=True)
+    """
+
+    _pools: ClassVar[dict[str, type]] = {
+        "null": NullPool,
+        "queue": QueuePool,
+        "singleton_thread": SingletonThreadPool,
+        "async_adapted_queue": AsyncAdaptedQueuePool,
+        "fallback_async_adapted_queue": FallbackAsyncAdaptedQueuePool,
+        "static": StaticPool,
+    }
+
+    @classmethod
+    def register(cls, name: str, pool_class: type, *, override: bool = False) -> None:
+        """Register a custom pool class.
+
+        Args:
+            name: Pool class identifier (lowercase recommended).
+            pool_class: Pool class type to register.
+            override: If True, allows overriding built-in pools (use with caution).
+
+        Raises:
+            ValueError: If name already exists and override=False.
+
+        Examples:
+            >>> PoolRegistry.register("my_pool", MyCustomPool)
+        """
+        if name in cls._pools and not override:
+            raise ValueError(
+                f"Pool class '{name}' is already registered. "
+                f"Use override=True to replace it (not recommended for built-ins)."
+            )
+        cls._pools[name] = pool_class
+
+    @classmethod
+    def resolve(cls, name: str) -> type:
+        """Resolve pool class by name.
+
+        Args:
+            name: Pool class identifier.
+
+        Returns:
+            Pool class type.
+
+        Raises:
+            ValueError: If pool class name is not registered.
+
+        Examples:
+            >>> pool = PoolRegistry.resolve("queue")
+            >>> pool
+            <class 'sqlalchemy.pool.QueuePool'>
+        """
+        try:
+            return cls._pools[name.lower()]
+        except KeyError as e:
+            available = ", ".join(sorted(cls._pools.keys()))
+            raise ValueError(f"Unknown pool class: {name}. Available: {available}") from e
+
+    @classmethod
+    def list_available(cls) -> list[str]:
+        """List all registered pool class names.
+
+        Returns:
+            Sorted list of registered pool names.
+
+        Examples:
+            >>> PoolRegistry.list_available()
+            ['async_adapted_queue', 'fallback_async_adapted_queue', 'null', 'queue', ...]
+        """
+        return sorted(cls._pools.keys())
+
+
+def resolve_pool_class(poolclass: PoolClassStr | str | type) -> type:
+    """Resolve pool class from string name or return the class directly.
+
+    Args:
+        poolclass: Pool class name (e.g., "null", "queue") or actual class type.
+
+    Returns:
+        Pool class type.
+
+    Raises:
+        ValueError: If pool class name is not recognized.
+    """
+    if isinstance(poolclass, str):
+        return PoolRegistry.resolve(poolclass)
+
+    return poolclass
+
+
+def register_pool_class(name: str, pool_class: type, *, override: bool = False) -> None:
+    """Register a custom pool class.
+
+    Convenience wrapper around :meth:`PoolRegistry.register` for users who prefer
+    a functional API over the class-based one.
+
+    Args:
+        name: Pool class identifier (lowercase recommended).
+        pool_class: Pool class type to register.
+        override: If True, allows overriding built-in pools (use with caution).
+
+    Raises:
+        ValueError: If name already exists and override=False.
+    """
+    PoolRegistry.register(name, pool_class, override=override)
+
+
+def build_engine_kwargs(
+    echo: bool,
+    poolclass: type,
+    isolation_level: str | None,
+    pool_settings: PoolSettingsProtocol | None,
+    connect_args: dict[str, object] | None,
+    extra_kwargs: dict[str, object],
+    use_orjson: bool = False,
+) -> dict[str, object]:
+    """Build SQLAlchemy engine keyword arguments.
+
+    Args:
+        echo: If True, SQLAlchemy will log all SQL statements.
+        poolclass: SQLAlchemy pool class.
+        isolation_level: Default transaction isolation level.
+        pool_settings: Pool configuration settings (validated by caller, e.g., Pydantic).
+        connect_args: Arguments passed to the database driver.
+        extra_kwargs: Additional keyword arguments for create_async_engine.
+        use_orjson: If True, use orjson for JSON serialization.
+
+    Returns:
+        Dictionary of engine keyword arguments ready for create_async_engine().
+
+    Raises:
+        ImportError: If use_orjson is True but orjson is not installed.
+
+    Examples:
+        >>> kwargs = build_engine_kwargs(
+        ...     echo=False,
+        ...     poolclass=NullPool,
+        ...     isolation_level=None,
+        ...     pool_settings=None,
+        ...     connect_args=None,
+        ...     extra_kwargs={},
+        ...     use_orjson=False,
+        ... )
+        >>> kwargs["echo"]
+        False
+    """
+    engine_kwargs: dict[str, object] = {
+        "echo": echo,
+        "poolclass": poolclass,
+        "isolation_level": isolation_level,
+        "pool_pre_ping": pool_settings.pre_ping if pool_settings else True,
+    }
+
+    if use_orjson:
+        from .serialization import configure_orjson_serialization  # noqa: PLC0415
+
+        engine_kwargs.update(configure_orjson_serialization())
+
+    if pool_settings:
+        _apply_pool_settings(engine_kwargs, poolclass, pool_settings)
+
+    if connect_args:
+        engine_kwargs["connect_args"] = {k: v for k, v in connect_args.items() if v is not None}
+
+    if extra_kwargs:
+        engine_kwargs.update(extra_kwargs)
+
+    return engine_kwargs
+
+
+def _apply_pool_settings(
+    engine_kwargs: dict[str, object],
+    poolclass: type,
+    pool_settings: PoolSettingsProtocol,
+) -> None:
+    """Apply pool settings to engine kwargs.
+
+    Maps pool settings to SQLAlchemy's expected ``pool_*`` keyword arguments.
+    Note: pool_pre_ping is already set in build_engine_kwargs, so not applied here.
+
+    Only applies pool size/overflow/recycle/timeout for pool classes that support them.
+    Checks if the pool class accepts these parameters via hasattr to avoid passing
+    unsupported kwargs to pools like NullPool or StaticPool.
+
+    Args:
+        engine_kwargs: Dictionary to update with pool configuration.
+        poolclass: Pool class being configured.
+        pool_settings: Pool configuration settings.
+    """
+    # Build pool parameters dict
+    params = {
+        "pool_size": pool_settings.size,
+        "max_overflow": pool_settings.max_overflow,
+        "pool_recycle": pool_settings.recycle,
+        "pool_timeout": pool_settings.timeout,
+    }
+
+    # Apply only non-None parameters
+    # SQLAlchemy pool classes that don't support these params will ignore them
+    # or raise TypeError if passed, so we rely on the pool class itself to validate
+    engine_kwargs.update({k: v for k, v in params.items() if v is not None})
+
+
+__all__ = [
+    "PoolClassStr",
+    "PoolRegistry",
+    "build_engine_kwargs",
+    "register_pool_class",
+    "resolve_pool_class",
+]

sqlalchemy_foundation_kit/base/metadata.py

@@ -0,0 +1,57 @@
+"""ORM metadata loading utilities."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from importlib import import_module
+
+from sqlalchemy import MetaData
+
+from .models import Base
+
+
+def load_orm_metadata(models_modules: Iterable[str], metadata: MetaData | None = None) -> MetaData:
+    """Load all ORM models metadata synchronously.
+
+    Imports specified modules to ensure that all SQLAlchemy models are
+    registered in the metadata. This is useful for migrations and schema
+    introspection with tools like Alembic.
+
+    Args:
+        models_modules: Iterable of module paths to import (e.g., ["myapp.models", "myapp.core.models"]).
+        metadata: Optional specific MetaData object to use. If None, uses Base.metadata.
+
+    Returns:
+        MetaData object containing all registered models from the imported modules.
+
+    Examples:
+        Load models from multiple modules:
+            >>> from sqlalchemy_foundation_kit.base import load_orm_metadata
+            >>> metadata = load_orm_metadata([
+            ...     "myapp.users.models",
+            ...     "myapp.orders.models",
+            ...     "myapp.products.models",
+            ... ])
+            >>> len(metadata.tables)
+            15
+
+        Use with custom metadata:
+            >>> from sqlalchemy import MetaData
+            >>> custom_meta = MetaData(schema="public")
+            >>> metadata = load_orm_metadata(["myapp.models"], metadata=custom_meta)
+
+        Typical usage in Alembic env.py:
+            >>> from sqlalchemy_foundation_kit.base import Base, load_orm_metadata
+            >>> target_metadata = Base.metadata
+            >>> load_orm_metadata(["myapp.models"])  # Register all models
+            >>> # Now target_metadata.tables contains all tables
+    """
+    for module in models_modules:
+        import_module(module)
+
+    return metadata if metadata is not None else Base.metadata
+
+
+__all__ = [
+    "load_orm_metadata",
+]

sqlalchemy_foundation_kit/base/models.py

@@ -0,0 +1,101 @@
+"""Database base models and mixins."""
+
+from __future__ import annotations
+
+import datetime
+import enum
+import uuid
+from functools import partial
+from typing import Any, ClassVar
+
+from sqlalchemy import TIMESTAMP, Enum, MetaData, func
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.orm import DeclarativeBase, Mapped, declared_attr, mapped_column
+from sqlalchemy.types import TypeEngine
+
+DB_NAMING_CONVENTION: dict[str, str] = {
+    "ix": "%(column_0_label)s_idx",
+    "uq": "%(table_name)s_%(column_0_name)s_key",
+    "ck": "%(table_name)s_%(constraint_name)s_check",
+    "fk": "%(table_name)s_%(column_0_name)s_fkey",
+    "pk": "%(table_name)s_pkey",
+}
+
+
+class Base(DeclarativeBase):
+    """Base class for all ORM models."""
+
+    type_annotation_map: ClassVar[dict[type, TypeEngine[Any]]] = {
+        uuid.UUID: postgresql.UUID(),
+        datetime.datetime: TIMESTAMP(timezone=True),
+    }
+    metadata = MetaData(naming_convention=DB_NAMING_CONVENTION)
+
+
+class BaseTable(Base):
+    """Base table class with __repr__."""
+
+    __abstract__ = True
+
+    def __repr__(self) -> str:
+        columns = {column.name: getattr(self, column.name) for column in self.__table__.columns}
+        return f"<{self.__tablename__}: {', '.join(f'{k}={v}' for k, v in columns.items())}>"
+
+
+class DatetimeColumnsMixin:
+    """Mixin for tables that need created_at and updated_at timestamps.
+
+    Control indexing via __created_at_index__ and __updated_at_index__ class variables in the model.
+    Defaults to False.
+    """
+
+    __created_at_index__: ClassVar[bool] = False
+    __updated_at_index__: ClassVar[bool] = False
+
+    @declared_attr
+    def created_at(self) -> Mapped[datetime.datetime]:
+        return mapped_column(
+            server_default=func.timezone("UTC", func.now()),
+            index=self.__created_at_index__,
+        )
+
+    @declared_attr
+    def updated_at(self) -> Mapped[datetime.datetime]:
+        return mapped_column(
+            server_default=func.timezone("UTC", func.now()),
+            onupdate=func.timezone("UTC", func.now()),
+            index=self.__updated_at_index__,
+        )
+
+
+def _extract_enum_values(enum_obj: type[enum.Enum] | list[object]) -> list[object]:
+    """Extract values from Python enum or return the list as-is.
+
+    Args:
+        enum_obj: Either a Python Enum class with __members__ or a list of values.
+
+    Returns:
+        List of enum values (extracting .value attribute if available) or the input list.
+
+    Examples:
+        >>> from enum import Enum
+        >>> class Color(Enum):
+        ...     RED = "red"
+        ...     BLUE = "blue"
+        >>> _extract_enum_values(Color)
+        ["red", "blue"]
+        >>> _extract_enum_values(["red", "blue"])
+        ["red", "blue"]
+    """
+    if hasattr(enum_obj, "__members__"):
+        return [getattr(item, "value", item) for item in enum_obj]
+    return list(enum_obj)  # type: ignore[arg-type]
+
+
+UnConstrainedEnum = partial(
+    Enum,
+    native_enum=False,
+    create_constraint=False,
+    validate_strings=True,
+    values_callable=_extract_enum_values,
+)