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

sqlalchemy_foundation_kit/contrib/settings/postgres.py

@@ -0,0 +1,183 @@
+"""Base PostgreSQL configuration using pydantic-settings."""
+
+from __future__ import annotations
+
+from typing import Literal
+from urllib.parse import quote_plus
+
+from ...base.engine import PoolClassStr
+
+try:
+    from pydantic import BaseModel, Field, SecretStr, model_validator
+    from pydantic_settings import BaseSettings, SettingsConfigDict
+except ImportError as _e:
+    raise ImportError(
+        "pydantic and pydantic-settings are required for this functionality. "
+        "Install with: pip install 'sqlalchemy-foundation-kit[settings]'"
+    ) from _e
+
+PostgresIsolationLevel = Literal["READ UNCOMMITTED", "READ COMMITTED", "REPEATABLE READ", "SERIALIZABLE"]
+PostgresJit = Literal["off", "on"]
+
+
+class ConnectionSettings(BaseModel):
+    """PostgreSQL connection configuration.
+
+    Groups all connection-related parameters: host, port, credentials, and database name.
+
+    Examples:
+        >>> connection = ConnectionSettings(
+        ...     host="localhost",
+        ...     port=5432,
+        ...     user="postgres",
+        ...     password=SecretStr("secret"),
+        ...     database="mydb",
+        ... )
+    """
+
+    host: str = Field(default="localhost", description="PostgreSQL host")
+    port: int = Field(default=5432, ge=1, le=65535, description="PostgreSQL port")
+    user: str = Field(default="postgres", description="PostgreSQL user")
+    password: SecretStr = Field(description="PostgreSQL password")
+    database: str = Field(description="Database name")
+
+
+class PoolSettings(BaseModel):
+    """PostgreSQL connection pool configuration.
+
+    Groups all connection pool-related parameters for SQLAlchemy engine.
+
+    Examples:
+        >>> pool = PoolSettings(size=10, max_overflow=20)
+    """
+
+    kind: PoolClassStr = Field(
+        default="async_adapted_queue",
+        description="PostgreSQL pool implementation kind",
+    )
+    size: int = Field(default=10, ge=1, description="Connection pool size")
+    max_overflow: int = Field(default=20, ge=0, description="Additional connections when pool is exhausted")
+    pre_ping: bool = Field(default=True, description="Check connection health before use (pre-ping)")
+    recycle: int = Field(default=3600, ge=-1, description="Recycle connections after N seconds")
+    timeout: float = Field(
+        default=30.0,
+        ge=0,
+        description="Seconds to wait before giving up on getting connection",
+    )
+
+    @model_validator(mode="after")
+    def _validate_pool_settings(self) -> PoolSettings:
+        """Validate pool configuration constraints."""
+        if self.kind == "static" and self.max_overflow > 0:
+            raise ValueError(f"max_overflow must be 0 for static pool, got {self.max_overflow}")
+        return self
+
+
+class QuerySettings(BaseModel):
+    """PostgreSQL query and performance configuration.
+
+    Groups query execution, caching, and transaction isolation settings.
+
+    Examples:
+        >>> query = QuerySettings(echo=False, isolation_level="READ COMMITTED")
+    """
+
+    echo: bool = Field(default=False, description="Echo SQL queries")
+    statement_cache_size: int = Field(default=0, ge=0, description="Statement cache size")
+    prepared_statement_cache_size: int = Field(default=0, ge=0, description="Prepared statement cache size")
+    isolation_level: PostgresIsolationLevel | None = Field(default=None, description="Transaction isolation level")
+
+
+class BasePostgresConfig(BaseSettings):
+    """Base PostgreSQL configuration.
+
+    Organized configuration for PostgreSQL database connections with grouped settings.
+
+    Attributes:
+        connection: Connection parameters (host, port, credentials, database).
+        pool: Connection pool configuration (size, overflow, timeouts).
+        query: Query execution settings (echo, caching, isolation level).
+        application_name: Application name for connection identification.
+        db_schema: Optional PostgreSQL schema name.
+        use_orjson_serialization: Use orjson for JSON serialization (requires orjson).
+        jit: JIT compilation setting (off/on) for PgBouncer compatibility.
+        metrics_enabled: Enable connection pool metrics collection.
+
+    Examples:
+        >>> config = BasePostgresConfig(
+        ...     connection=ConnectionSettings(
+        ...         host="localhost",
+        ...         user="postgres",
+        ...         password=SecretStr("secret"),
+        ...         database="mydb",
+        ...     ),
+        ...     application_name="my-service",
+        ... )
+        >>> dsn = config.to_dsn()
+        >>> host = config.connection.host
+        >>> pool_size = config.pool.size
+    """
+
+    # Grouped configuration
+    connection: ConnectionSettings = Field(description="PostgreSQL connection settings")
+    pool: PoolSettings = Field(default_factory=PoolSettings, description="Connection pool settings")
+    query: QuerySettings = Field(default_factory=QuerySettings, description="Query execution settings")
+
+    # Top-level settings
+    application_name: str = Field(description="Application name for PostgreSQL")
+    db_schema: str | None = Field(default=None, description="PostgreSQL schema name")
+    use_orjson_serialization: bool = Field(
+        default=True,
+        description="Use orjson for JSON serialization (requires orjson installed)",
+    )
+    jit: PostgresJit | None = Field(default="off", description="JIT setting (off/on)")
+    metrics_enabled: bool = Field(default=False, description="Enable PostgreSQL metrics")
+
+    def __repr__(self) -> str:
+        """Return representation with masked password."""
+        return f"{self.__class__.__name__}(dsn={self.to_dsn(mask_password=True)!r})"
+
+    def to_dsn(self, driver: str | None = "asyncpg", mask_password: bool = False) -> str:
+        """Build PostgreSQL DSN for async connections.
+
+        This library is async-only, so DSN always includes asyncpg driver by default.
+
+        Args:
+            driver: Driver name (default: 'asyncpg' for async connections).
+                Pass None to omit driver suffix.
+            mask_password: If True, the password will be masked (default: False).
+
+        Returns:
+            PostgreSQL connection string (e.g., 'postgresql+asyncpg://user:pass@host:5432/db').
+
+        Examples:
+            >>> config.to_dsn()
+            'postgresql+asyncpg://user:secret@localhost:5432/mydb'
+            >>> config.to_dsn(mask_password=True)
+            'postgresql+asyncpg://user:**********@localhost:5432/mydb'
+        """
+        user = quote_plus(self.connection.user)
+        password = "**********" if mask_password else quote_plus(self.connection.password.get_secret_value())
+        scheme = f"postgresql+{driver}" if driver else "postgresql"
+        return f"{scheme}://{user}:{password}@{self.connection.host}:{self.connection.port}/{self.connection.database}"
+
+
+class BasePostgresMigrationsConfig(BaseSettings):
+    """Base configuration for PostgreSQL database migrations."""
+
+    model_config = SettingsConfigDict(
+        extra="ignore",
+        env_nested_delimiter="__",
+        case_sensitive=False,
+    )
+
+    postgres: BasePostgresConfig
+
+
+__all__ = [
+    "BasePostgresConfig",
+    "BasePostgresMigrationsConfig",
+    "ConnectionSettings",
+    "PoolSettings",
+    "QuerySettings",
+]

sqlalchemy_foundation_kit/contrib/telemetry/__init__.py

@@ -0,0 +1,17 @@
+"""OpenTelemetry tracing integration."""
+
+from __future__ import annotations
+
+from .instrumentations import (
+    instrument_asyncpg,
+    instrument_engine,
+    instrument_sqlalchemy,
+)
+from .uow import TracedAsyncUnitOfWork
+
+__all__ = [
+    "TracedAsyncUnitOfWork",
+    "instrument_asyncpg",
+    "instrument_engine",
+    "instrument_sqlalchemy",
+]

sqlalchemy_foundation_kit/contrib/telemetry/instrumentations.py

@@ -0,0 +1,101 @@
+"""OpenTelemetry instrumentation functions for SQLAlchemy and asyncpg."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncEngine
+
+
+def instrument_sqlalchemy(engine: Any | None = None, **kwargs: Any) -> None:
+    """Instrument SQLAlchemy engine for OpenTelemetry tracing.
+
+    Automatically traces all SQLAlchemy operations including queries,
+    commits, and rollbacks.
+
+    Args:
+        engine: Optional SQLAlchemy engine to instrument. If None, all engines.
+        **kwargs: Additional keyword arguments passed to SQLAlchemyInstrumentor.
+
+    Raises:
+        ImportError: If opentelemetry-instrumentation-sqlalchemy is not installed.
+
+    Examples:
+        >>> from sqlalchemy import create_engine
+        >>> from sqlalchemy_foundation_kit.contrib.telemetry import instrument_sqlalchemy
+        >>> engine = create_engine("postgresql://...")
+        >>> instrument_sqlalchemy(engine=engine)
+    """
+    try:
+        from opentelemetry.instrumentation.sqlalchemy import (  # noqa: PLC0415
+            SQLAlchemyInstrumentor,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "opentelemetry-instrumentation-sqlalchemy not installed. "
+            "Install with: pip install 'sqlalchemy-foundation-kit[telemetry]'"
+        ) from e
+
+    call_kwargs: dict[str, Any] = dict(kwargs)
+    if engine is not None:
+        call_kwargs["engine"] = engine
+
+    SQLAlchemyInstrumentor().instrument(**call_kwargs)
+
+
+def instrument_engine(engine: AsyncEngine, **kwargs: Any) -> None:
+    """Attach OpenTelemetry tracing to a specific async SQLAlchemy engine.
+
+    Designed to be passed as the ``on_engine_created`` hook of
+    :class:`~sqlalchemy_foundation_kit.session.AsyncSessionManager` or
+    :func:`~sqlalchemy_foundation_kit.session.create_async_session_manager`.
+
+    Args:
+        engine: The AsyncEngine to instrument.
+        **kwargs: Additional keyword arguments passed to SQLAlchemyInstrumentor.
+
+    Raises:
+        ImportError: If opentelemetry-instrumentation-sqlalchemy is not installed.
+
+    Examples:
+        >>> from sqlalchemy_foundation_kit.session import create_async_session_manager
+        >>> from sqlalchemy_foundation_kit.contrib.telemetry import instrument_engine
+        >>> manager = create_async_session_manager(config, on_engine_created=instrument_engine)
+    """
+    instrument_sqlalchemy(engine=engine.sync_engine, **kwargs)
+
+
+def instrument_asyncpg(**kwargs: Any) -> None:
+    """Instrument asyncpg connections for OpenTelemetry tracing.
+
+    Automatically traces all asyncpg database operations at the connection level.
+
+    Args:
+        **kwargs: Additional keyword arguments passed to AsyncPGInstrumentor.
+
+    Raises:
+        ImportError: If opentelemetry-instrumentation-asyncpg is not installed.
+
+    Examples:
+        >>> from sqlalchemy_foundation_kit.contrib.telemetry import instrument_asyncpg
+        >>> instrument_asyncpg()
+    """
+    try:
+        from opentelemetry.instrumentation.asyncpg import (  # noqa: PLC0415
+            AsyncPGInstrumentor,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "opentelemetry-instrumentation-asyncpg not installed. "
+            "Install with: pip install 'sqlalchemy-foundation-kit[telemetry]'"
+        ) from e
+
+    AsyncPGInstrumentor().instrument(**kwargs)
+
+
+__all__ = [
+    "instrument_asyncpg",
+    "instrument_engine",
+    "instrument_sqlalchemy",
+]

sqlalchemy_foundation_kit/contrib/telemetry/uow.py

@@ -0,0 +1,227 @@
+"""Unit of Work with OpenTelemetry tracing support."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncIterator, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from typing import TYPE_CHECKING, Generic
+
+from ..._typing import T
+from ...base._optional import require_optional
+from ...uow import AsyncSQLAlchemyUnitOfWork, IsolationLevel
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span, Tracer
+    from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+logger = logging.getLogger(__name__)
+
+try:
+    trace = require_optional("opentelemetry.trace", "telemetry")
+    from opentelemetry.trace import Status, StatusCode
+
+    HAS_OTEL = True
+except ImportError:
+    HAS_OTEL = False
+    trace = None  # type: ignore[assignment]
+
+
+class TracedAsyncUnitOfWork(AsyncSQLAlchemyUnitOfWork[T], Generic[T]):
+    """Unit of Work with automatic OpenTelemetry tracing.
+
+    Automatically creates spans for transaction() and query() operations,
+    including transaction attributes (isolation level, duration, outcome).
+
+    Example:
+        from opentelemetry import trace
+        from opentelemetry.sdk.trace import TracerProvider
+
+        trace.set_tracer_provider(TracerProvider())
+
+        uow = TracedAsyncUnitOfWork(
+            session_maker=session_maker,
+            transaction_factory=MyTransaction,
+            service_name="my-service",
+        )
+
+        async with uow.transaction() as tx:
+            # This operation is automatically traced
+            user = await tx.users.create(...)
+    """
+
+    def __init__(
+        self,
+        session_maker: async_sessionmaker[AsyncSession],
+        transaction_factory: Callable[[AsyncSession], T],
+        service_name: str = "sqlalchemy-foundation-kit",
+        *,
+        flush_before_commit: bool = True,
+    ) -> None:
+        """Initialize traced unit of work.
+
+        Args:
+            session_maker: SQLAlchemy async session maker.
+            transaction_factory: Factory function to create transaction objects.
+            service_name: Service name for OpenTelemetry tracer.
+            flush_before_commit: Default ``flush_before_commit`` policy applied when
+                :meth:`transaction` is called without an explicit override.
+        """
+        super().__init__(session_maker, transaction_factory, flush_before_commit=flush_before_commit)
+        self._tracer: Tracer | None = trace.get_tracer(service_name) if HAS_OTEL else None
+
+    @asynccontextmanager
+    async def _traced(
+        self,
+        operation: str,
+        context_manager: AbstractAsyncContextManager[T],
+        attributes: dict[str, str | bool] | None = None,
+    ) -> AsyncIterator[T]:
+        """Generic tracing wrapper for UoW operations.
+
+        Args:
+            operation: Operation name (e.g., "transaction", "query").
+            context_manager: Async context manager to wrap with tracing.
+            attributes: Optional span attributes to set.
+
+        Yields:
+            Transaction object from the wrapped context manager.
+        """
+        if not self._tracer:
+            # No tracing available, delegate to wrapped context manager
+            async with context_manager as result:
+                yield result
+            return
+
+        span: Span = self._tracer.start_span(f"uow.{operation}")
+        span.set_attribute("db.operation", operation)
+
+        if attributes:
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+        try:
+            async with context_manager as result:
+                yield result
+
+            span.set_status(Status(StatusCode.OK))
+
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            raise
+
+        finally:
+            span.end()
+
+    @asynccontextmanager
+    async def transaction(
+        self,
+        isolation_level: IsolationLevel | str | None = None,
+        flush_before_commit: bool | None = None,
+    ) -> AsyncIterator[T]:
+        """Create a new transaction context with tracing.
+
+        Automatically creates a span named "uow.transaction" with attributes:
+        - db.operation: "transaction"
+        - db.isolation_level: The isolation level (if specified)
+
+        Args:
+            isolation_level: Optional transaction isolation level.
+            flush_before_commit: If True, flush before commit.
+
+        Yields:
+            Transaction object with repositories.
+        """
+        attributes: dict[str, str | bool] = {}
+        if isolation_level:
+            attributes["db.isolation_level"] = str(isolation_level)
+
+        async with self._traced(
+            operation="transaction",
+            context_manager=super().transaction(
+                isolation_level=isolation_level,
+                flush_before_commit=flush_before_commit,
+            ),
+            attributes=attributes,
+        ) as tx:
+            yield tx
+
+    @asynccontextmanager
+    async def managed_session(
+        self,
+        isolation_level: IsolationLevel | str | None = None,
+    ) -> AsyncIterator[tuple[T, AsyncSession]]:
+        """Create a session with manual transaction control and tracing.
+
+        Automatically creates a span named "uow.managed_session" with attributes:
+        - db.operation: "managed_session"
+        - db.isolation_level: The isolation level (if specified)
+
+        Args:
+            isolation_level: Optional transaction isolation level.
+
+        Yields:
+            Tuple of (transaction object, session) for manual control.
+        """
+        attributes: dict[str, str | bool] = {}
+        if isolation_level:
+            attributes["db.isolation_level"] = str(isolation_level)
+
+        if not self._tracer:
+            # No tracing available, delegate directly
+            async with super().managed_session(isolation_level=isolation_level) as result:
+                yield result
+            return
+
+        span: Span = self._tracer.start_span("uow.managed_session")
+        span.set_attribute("db.operation", "managed_session")
+
+        if attributes:
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+        try:
+            async with super().managed_session(isolation_level=isolation_level) as result:
+                yield result
+
+            span.set_status(Status(StatusCode.OK))
+
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            raise
+
+        finally:
+            span.end()
+
+    @asynccontextmanager
+    async def query(
+        self,
+        isolation_level: IsolationLevel | str | None = None,
+    ) -> AsyncIterator[T]:
+        """Create a read-only query context with tracing.
+
+        Automatically creates a span named "uow.query" with attributes:
+        - db.operation: "query"
+        - db.isolation_level: The isolation level (if specified)
+
+        Args:
+            isolation_level: Optional transaction isolation level.
+
+        Yields:
+            Query object with repositories.
+        """
+        attributes: dict[str, str | bool] = {}
+        if isolation_level:
+            attributes["db.isolation_level"] = str(isolation_level)
+
+        async with self._traced(
+            operation="query",
+            context_manager=super().query(isolation_level=isolation_level),
+            attributes=attributes,
+        ) as qx:
+            yield qx
+
+
+__all__ = ["TracedAsyncUnitOfWork"]

sqlalchemy_foundation_kit/protocols/__init__.py

@@ -0,0 +1,21 @@
+"""Protocols for database infrastructure.
+
+This module defines protocols (structural subtyping) for key abstractions,
+allowing dependency inversion without tight coupling to specific implementations.
+"""
+
+from __future__ import annotations
+
+from .metrics import (
+    CheckoutRecorder,
+    ErrorRecorder,
+    PoolStatsRecorder,
+    PostgresMetricsProtocol,
+)
+
+__all__ = [
+    "CheckoutRecorder",
+    "ErrorRecorder",
+    "PoolStatsRecorder",
+    "PostgresMetricsProtocol",
+]

sqlalchemy_foundation_kit/protocols/metrics.py

@@ -0,0 +1,77 @@
+"""Metrics protocols for database monitoring.
+
+Protocols are split by capability (ISP). Implementations may satisfy
+the narrow protocols selectively (e.g., only record errors), and
+``PostgresMetricsProtocol`` aggregates them for convenience.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class PoolStatsRecorder(Protocol):
+    """Capability protocol for recording pool statistics."""
+
+    def record_pool_stats(
+        self,
+        pool_size: int,
+        pool_checked_out: int,
+        pool_overflow: int,
+    ) -> None:
+        """Record database connection pool statistics.
+
+        Args:
+            pool_size: Current total number of connections in the pool.
+            pool_checked_out: Number of connections currently in use.
+            pool_overflow: Number of connections over the configured pool_size.
+        """
+        ...
+
+
+class CheckoutRecorder(Protocol):
+    """Capability protocol for recording connection checkout duration."""
+
+    def record_checkout(self, duration: float) -> None:
+        """Record a database connection checkout from the pool.
+
+        Args:
+            duration: Time taken to acquire the connection from the pool, in seconds.
+        """
+        ...
+
+
+class ErrorRecorder(Protocol):
+    """Capability protocol for recording database errors."""
+
+    def record_error(self, error_type: str, is_timeout: bool = False) -> None:
+        """Record a database connection or execution error.
+
+        Args:
+            error_type: The type of error that occurred (e.g., "OperationalError").
+            is_timeout: True if this error was specifically a connection checkout timeout.
+        """
+        ...
+
+
+class PostgresMetricsProtocol(PoolStatsRecorder, CheckoutRecorder, ErrorRecorder, Protocol):
+    """Composite protocol covering all PostgreSQL metrics capabilities.
+
+    Aggregates the narrow capability protocols for convenience. Implementations
+    that only need a subset can implement the individual protocols directly.
+
+    Examples:
+        >>> class MyMetrics:
+        ...     def record_pool_stats(self, pool_size: int, pool_checked_out: int, pool_overflow: int) -> None: ...
+        ...     def record_checkout(self, duration: float) -> None: ...
+        ...     def record_error(self, error_type: str, is_timeout: bool = False) -> None: ...
+        >>> metrics: PostgresMetricsProtocol = MyMetrics()
+    """
+
+
+__all__ = [
+    "CheckoutRecorder",
+    "ErrorRecorder",
+    "PoolStatsRecorder",
+    "PostgresMetricsProtocol",
+]

sqlalchemy_foundation_kit/session/__init__.py

@@ -0,0 +1,27 @@
+"""Async session management module."""
+
+from __future__ import annotations
+
+from .builder import AsyncSessionManagerBuilder
+from .connection import AsyncCConnection
+from .factories import create_async_session_manager
+from .locks import try_advisory_xact_lock
+from .manager import AsyncSessionManager
+from .retry import (
+    DEFAULT_HEALTHCHECK_QUERY,
+    DEFAULT_RETRY_CONFIG,
+    RetryConfig,
+    retry_async_connection,
+)
+
+__all__ = [
+    "DEFAULT_HEALTHCHECK_QUERY",
+    "DEFAULT_RETRY_CONFIG",
+    "AsyncCConnection",
+    "AsyncSessionManager",
+    "AsyncSessionManagerBuilder",
+    "RetryConfig",
+    "create_async_session_manager",
+    "retry_async_connection",
+    "try_advisory_xact_lock",
+]