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

sqlalchemy_foundation_kit/session/builder.py

@@ -0,0 +1,293 @@
+"""Builder pattern for AsyncSessionManager configuration.
+
+Provides a fluent interface for constructing AsyncSessionManager instances
+with many optional parameters, following the Builder pattern to simplify
+complex object creation and improve code readability.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Generic
+
+from .._typing import SessionT
+from .manager import AsyncSessionManager
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlalchemy.ext.asyncio import AsyncEngine
+
+    from ..config import PoolSettingsProtocol
+    from ..protocols import PostgresMetricsProtocol
+
+
+class AsyncSessionManagerBuilder(Generic[SessionT]):
+    """Builder for AsyncSessionManager construction.
+
+    Provides a fluent API for configuring AsyncSessionManager with many optional
+    parameters. This pattern improves code readability and follows KISS principle
+    by avoiding constructors with 10+ parameters.
+
+    Examples:
+        Basic usage:
+            >>> manager = (
+            ...     AsyncSessionManagerBuilder("postgresql+asyncpg://...")
+            ...     .with_pool("queue")
+            ...     .with_echo(True)
+            ...     .build()
+            ... )
+
+        Advanced configuration:
+            >>> manager = (
+            ...     AsyncSessionManagerBuilder[CustomSession]("postgresql+asyncpg://...")
+            ...     .with_session_class(CustomSession)
+            ...     .with_pool("queue", pool_settings=settings.pool)
+            ...     .with_metrics(metrics)
+            ...     .with_callbacks(on_engine_created=instrument_engine)
+            ...     .with_json_serialization(orjson=True)
+            ...     .with_isolation_level("READ COMMITTED")
+            ...     .build()
+            ... )
+
+        Reusable configuration:
+            >>> builder = (
+            ...     AsyncSessionManagerBuilder("postgresql+asyncpg://...")
+            ...     .with_pool("queue")
+            ...     .with_metrics(metrics)
+            ... )
+            >>> manager1 = builder.with_echo(True).build()
+            >>> manager2 = builder.with_echo(False).build()
+    """
+
+    def __init__(self, url: str) -> None:
+        """Initialize builder with database URL (required).
+
+        Args:
+            url: Database connection URL (required parameter).
+        """
+        self._url = url
+        self._echo: bool = False
+        self._poolclass: str | type = "null"
+        self._session_class: type[SessionT] | None = None
+        self._expire_on_commit: bool = False
+        self._connect_args: dict[str, object] | None = None
+        self._isolation_level: str | None = None
+        self._pool_settings: PoolSettingsProtocol | None = None
+        self._use_orjson: bool = False
+        self._metrics: PostgresMetricsProtocol | None = None
+        self._on_engine_created: Callable[[AsyncEngine], None] | None = None
+        self._dispose_timeout: float | None = None
+        self._extra_kwargs: dict[str, object] = {}
+
+    def with_echo(self, echo: bool = True) -> AsyncSessionManagerBuilder[SessionT]:
+        """Enable SQL statement logging.
+
+        Args:
+            echo: If True, SQLAlchemy logs all SQL statements.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._echo = echo
+        return self
+
+    def with_pool(
+        self,
+        poolclass: str | type,
+        pool_settings: PoolSettingsProtocol | None = None,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Configure connection pool.
+
+        Args:
+            poolclass: Pool class name or type (e.g., "queue", "null").
+            pool_settings: Optional pool configuration settings.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._poolclass = poolclass
+        self._pool_settings = pool_settings
+        return self
+
+    def with_session_class(
+        self,
+        session_class: type[SessionT],
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Use custom session class.
+
+        Args:
+            session_class: Custom AsyncSession subclass.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._session_class = session_class
+        return self
+
+    def with_expire_on_commit(
+        self,
+        expire: bool = True,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Configure object expiration behavior after commit.
+
+        Args:
+            expire: If True, all objects expire after commit.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._expire_on_commit = expire
+        return self
+
+    def with_connect_args(
+        self,
+        **connect_args: object,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Add database driver connection arguments.
+
+        Args:
+            **connect_args: Arguments passed to the database driver.
+
+        Returns:
+            Self for method chaining.
+
+        Examples:
+            >>> builder.with_connect_args(
+            ...     server_settings={"application_name": "myapp"},
+            ...     command_timeout=60,
+            ... )
+        """
+        if self._connect_args is None:
+            self._connect_args = {}
+        self._connect_args.update(connect_args)
+        return self
+
+    def with_isolation_level(
+        self,
+        isolation_level: str,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Set default transaction isolation level.
+
+        Args:
+            isolation_level: Default isolation level (e.g., "READ COMMITTED").
+
+        Returns:
+            Self for method chaining.
+        """
+        self._isolation_level = isolation_level
+        return self
+
+    def with_metrics(
+        self,
+        metrics: PostgresMetricsProtocol,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Enable connection pool metrics collection.
+
+        Args:
+            metrics: Metrics collector implementing PostgresMetricsProtocol.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._metrics = metrics
+        return self
+
+    def with_callbacks(
+        self,
+        on_engine_created: Callable[[AsyncEngine], None],
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Register engine creation callback.
+
+        Useful for attaching instrumentation (OpenTelemetry), custom event
+        listeners, or debug hooks right after engine creation.
+
+        Args:
+            on_engine_created: Callback invoked with AsyncEngine after creation.
+
+        Returns:
+            Self for method chaining.
+
+        Examples:
+            >>> from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+            >>> def instrument(engine):
+            ...     SQLAlchemyInstrumentor().instrument(engine=engine.sync_engine)
+            >>> builder.with_callbacks(on_engine_created=instrument)
+        """
+        self._on_engine_created = on_engine_created
+        return self
+
+    def with_json_serialization(
+        self,
+        orjson: bool = True,
+    ) -> AsyncSessionManagerBuilder[SessionT]:
+        """Configure JSON serialization backend.
+
+        Args:
+            orjson: If True, use orjson for faster JSON serialization.
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ImportError: If orjson=True but orjson is not installed.
+        """
+        self._use_orjson = orjson
+        return self
+
+    def with_extra_kwargs(self, **kwargs: object) -> AsyncSessionManagerBuilder[SessionT]:
+        """Add additional keyword arguments for create_async_engine.
+
+        Args:
+            **kwargs: Additional engine configuration.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._extra_kwargs.update(kwargs)
+        return self
+
+    def with_dispose_timeout(self, timeout: float) -> AsyncSessionManagerBuilder[SessionT]:
+        """Configure how long :meth:`AsyncSessionManager.aclose` waits for engine disposal.
+
+        Args:
+            timeout: Maximum seconds to wait for the engine to dispose.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._dispose_timeout = timeout
+        return self
+
+    def build(self) -> AsyncSessionManager[SessionT]:
+        """Build AsyncSessionManager instance with configured parameters.
+
+        Returns:
+            Configured AsyncSessionManager instance.
+
+        Raises:
+            ImportError: If use_orjson=True but orjson is not installed.
+
+        Examples:
+            >>> manager = builder.build()
+            >>> async with manager.get_session() as session:
+            ...     await session.execute(...)
+        """
+        kwargs: dict[str, object] = {
+            "url": self._url,
+            "echo": self._echo,
+            "poolclass": self._poolclass,
+            "session_class": self._session_class,
+            "expire_on_commit": self._expire_on_commit,
+            "connect_args": self._connect_args,
+            "isolation_level": self._isolation_level,
+            "pool_settings": self._pool_settings,
+            "use_orjson": self._use_orjson,
+            "metrics": self._metrics,
+            "on_engine_created": self._on_engine_created,
+        }
+        if self._dispose_timeout is not None:
+            kwargs["dispose_timeout"] = self._dispose_timeout
+        return AsyncSessionManager(**kwargs, **self._extra_kwargs)  # type: ignore[arg-type]
+
+
+__all__ = ["AsyncSessionManagerBuilder"]

sqlalchemy_foundation_kit/session/connection.py

@@ -0,0 +1,33 @@
+"""Custom connection class for pgbouncer compatibility (async)."""
+
+import uuid
+
+import asyncpg
+
+
+class AsyncCConnection(asyncpg.Connection):
+    """Custom async connection class for pgbouncer magic.
+
+    This subclass overrides only the private method _get_unique_id so that
+    prepared statement identifiers are unique per connection. That is required
+    when using pgbouncer in transaction mode, where the same server connection
+    may be reused for different logical connections.
+
+    See: https://github.com/sqlalchemy/sqlalchemy/issues/6467
+    """
+
+    def _get_unique_id(self, prefix: str) -> str:
+        """Generate unique ID for prepared statements.
+
+        Args:
+            prefix: Prefix for the unique ID.
+
+        Returns:
+            A unique string ID including the prefix and a UUID.
+        """
+        return f"__asyncpg_{prefix}_{uuid.uuid4()}__"
+
+
+__all__ = [
+    "AsyncCConnection",
+]

sqlalchemy_foundation_kit/session/factories.py

@@ -0,0 +1,104 @@
+"""Async database session manager factory."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ..config import PostgresSettingsProtocol
+from .connection import AsyncCConnection
+from .manager import AsyncSessionManager
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import asyncpg
+    from sqlalchemy.ext.asyncio import AsyncEngine
+
+    from ..protocols import PostgresMetricsProtocol
+
+
+def create_async_session_manager(
+    postgres_config: PostgresSettingsProtocol,
+    application_name: str | None = None,
+    metrics: PostgresMetricsProtocol | None = None,
+    on_engine_created: Callable[[AsyncEngine], None] | None = None,
+    connection_class: type[asyncpg.Connection] | None = None,
+    extra_server_settings: dict[str, str] | None = None,
+    extra_connect_args: dict[str, object] | None = None,
+    **kwargs: Any,
+) -> AsyncSessionManager[AsyncSession]:
+    """Create async session manager with PostgreSQL-specific configuration.
+
+    Args:
+        postgres_config: PostgreSQL configuration implementing PostgresSettingsProtocol.
+        application_name: Optional custom application name. If None, uses postgres_config.application_name.
+        metrics: Optional metrics collector for connection pool monitoring.
+        on_engine_created: Optional callback invoked with the AsyncEngine right after creation.
+            Use it to attach OpenTelemetry instrumentation, custom listeners, etc.
+        connection_class: Custom asyncpg Connection subclass. Defaults to ``AsyncCConnection``
+            which provides pgbouncer transaction-mode compatibility.
+        extra_server_settings: Additional PostgreSQL ``server_settings`` to merge with defaults
+            (e.g., ``{"statement_timeout": "30000", "timezone": "UTC"}``). User-provided keys
+            override library defaults.
+        extra_connect_args: Additional asyncpg ``connect_args`` to merge with defaults
+            (e.g., ``{"command_timeout": 60}``). User-provided keys override library defaults.
+        **kwargs: Additional keyword arguments passed to AsyncSessionManager.
+
+    Returns:
+        Configured AsyncSessionManager instance.
+
+    Examples:
+        Basic usage:
+            >>> manager = create_async_session_manager(postgres_config)
+
+        With metrics:
+            >>> from sqlalchemy_foundation_kit.contrib.metrics import PostgresMetrics
+            >>> manager = create_async_session_manager(postgres_config, metrics=PostgresMetrics())
+
+        With custom server settings and command timeout:
+            >>> manager = create_async_session_manager(
+            ...     postgres_config,
+            ...     extra_server_settings={"statement_timeout": "30000", "timezone": "UTC"},
+            ...     extra_connect_args={"command_timeout": 60},
+            ... )
+
+        With OpenTelemetry tracing bound to this engine:
+            >>> from sqlalchemy_foundation_kit.contrib.telemetry import instrument_engine
+            >>> manager = create_async_session_manager(
+            ...     postgres_config,
+            ...     on_engine_created=instrument_engine,
+            ... )
+    """
+    app_name = application_name or postgres_config.application_name
+
+    # Build server settings with optional overrides
+    server_settings: dict[str, str] = {
+        "application_name": app_name,
+        **({"jit": postgres_config.jit} if postgres_config.jit is not None else {}),
+        **({"search_path": postgres_config.db_schema} if postgres_config.db_schema is not None else {}),
+        **(extra_server_settings or {}),
+    }
+
+    # Build connect args with optional overrides
+    connect_args: dict[str, object] = {
+        "server_settings": server_settings,
+        "statement_cache_size": postgres_config.query.statement_cache_size,
+        "prepared_statement_cache_size": postgres_config.query.prepared_statement_cache_size,
+        "connection_class": connection_class or AsyncCConnection,
+        **(extra_connect_args or {}),
+    }
+
+    return AsyncSessionManager(
+        url=postgres_config.to_dsn(),
+        echo=postgres_config.query.echo,
+        poolclass=postgres_config.pool.kind,
+        connect_args=connect_args,
+        isolation_level=postgres_config.query.isolation_level,
+        pool_settings=postgres_config.pool,
+        use_orjson=postgres_config.use_orjson_serialization,
+        metrics=metrics,
+        on_engine_created=on_engine_created,
+        **kwargs,
+    )

sqlalchemy_foundation_kit/session/locks.py

@@ -0,0 +1,68 @@
+"""PostgreSQL advisory locks (async)."""
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+# PostgreSQL bigint (signed 64-bit) range constants
+_INT64_OFFSET: int = 1 << 63  # 2^63 = 9223372036854775808
+_INT64_MASK: int = (1 << 64) - 1  # 2^64 - 1 = 18446744073709551615
+
+
+async def try_advisory_xact_lock(session: AsyncSession, key: str | int) -> bool:
+    """Acquire a Postgres transaction-scoped advisory lock.
+
+    Uses ``pg_try_advisory_xact_lock``: non-blocking, released automatically
+    at transaction end. String keys are hashed to integers. The key is then
+    truncated to signed 64-bit as Postgres expects.
+
+    Args:
+        session: SQLAlchemy AsyncSession within an active transaction.
+        key: Lock identifier (string or integer). Strings are hashed to integers.
+
+    Returns:
+        True if lock was acquired, False if already held by another session.
+
+    Examples:
+        >>> async with session_maker() as session:
+        ...     async with session.begin():
+        ...         if await try_advisory_xact_lock(session, "my_operation"):
+        ...             # Perform protected operation
+        ...             await session.execute(...)
+        ...             await session.commit()
+    """
+    # Convert string keys to integers via hashing
+    int_key = hash(key) if isinstance(key, str) else key
+
+    result = await session.execute(
+        text("SELECT pg_try_advisory_xact_lock(:k)"),
+        {"k": _to_signed64(int_key)},
+    )
+    return bool(result.scalar())
+
+
+def _to_signed64(key: int) -> int:
+    """Wrap integer to PostgreSQL signed 64-bit bigint range.
+
+    PostgreSQL advisory locks use bigint (signed 64-bit integers).
+    This function wraps arbitrary Python ints into the range [-2^63, 2^63-1].
+
+    Args:
+        key: Integer of any size.
+
+    Returns:
+        Equivalent value in range [-9223372036854775808, 9223372036854775807].
+
+    Examples:
+        >>> _to_signed64(12345)
+        12345
+        >>> _to_signed64(2**64 + 100)
+        100
+        >>> _to_signed64(-1)
+        -1
+    """
+    return ((key + _INT64_OFFSET) & _INT64_MASK) - _INT64_OFFSET
+
+
+__all__ = [
+    "try_advisory_xact_lock",
+]