API Reference¶

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__,
        )

class PydanticJSONB(TypeDecorator):
    """SQLAlchemy TypeDecorator for Pydantic models stored as JSONB.

    Validates and serializes values against the supplied Pydantic-compatible type
    on both write and read paths. A ``model_type`` is **required** — if you need
    raw dict storage without validation, use SQLAlchemy's built-in ``JSONB`` directly
    (or the :data:`GenericJSONDict` alias).
    """

    impl = JSONB
    cache_ok = True

    def __init__(self, model_type: type[T], *args: Any, **kwargs: Any) -> None:
        """Initialize the type decorator.

        Args:
            model_type: Pydantic model class (or any type compatible with
                ``pydantic.TypeAdapter``) used to validate and serialize values.
        """
        self.model_type = model_type
        self.adapter: TypeAdapter[T] = TypeAdapter(model_type)
        super().__init__(*args, **kwargs)

    def process_bind_param(self, value: Any, dialect: Any) -> Any:
        if value is None:
            return None

        # Validate before dump to avoid Pydantic serialization warnings when value
        # is a dict (e.g. from model_dump()). This ensures value matches the expected
        # schema and converts it to a model instance if needed.
        validated = self.adapter.validate_python(value)
        return self.adapter.dump_python(validated, mode="json")

    def process_result_value(self, value: Any, dialect: Any) -> Any:
        if value is None:
            return None

        try:
            return self.adapter.validate_python(value)
        except ValidationError:
            logger.warning(
                "Validation error while loading %s from JSONB. Using raw data. "
                "This may indicate legacy data that doesn't match current schema.",
                self.model_type,
            )
            return value

class AsyncSessionManager(Generic[SessionT]):
    """Manages async database sessions with configurable connection pooling.

    Supports two initialization approaches:

    1. **Direct constructor** — all parameters in constructor with defaults.
    2. **Builder pattern** — use :class:`AsyncSessionManagerBuilder` for more
       readable complex configurations.
    """

    def __init__(
        self,
        url: str,
        echo: bool = False,
        poolclass: str | type = "null",
        session_class: type[SessionT] | None = None,
        expire_on_commit: bool = False,
        connect_args: dict[str, object] | None = None,
        isolation_level: str | None = None,
        pool_settings: PoolSettingsProtocol | None = None,
        use_orjson: bool = False,
        metrics: PostgresMetricsProtocol | None = None,
        on_engine_created: Callable[[AsyncEngine], None] | None = None,
        dispose_timeout: float = DEFAULT_DISPOSE_TIMEOUT_SECONDS,
        **kwargs: object,
    ) -> None:
        """Initialize session manager with direct configuration.

        Args:
            url: Database connection URL (required).
            echo: If True, SQLAlchemy will log all SQL statements (default: False).
            poolclass: SQLAlchemy pool class or name (default: "null").
                Use "queue" for production, "null" for testing.
            session_class: Custom ``AsyncSession`` subclass (default: ``AsyncSession``).
            expire_on_commit: If True, objects expire after commit (default: False).
            connect_args: Arguments passed to the database driver (default: None).
            isolation_level: Default transaction isolation level (default: None).
            pool_settings: Pool configuration settings (default: None).
            use_orjson: If True, use orjson for JSON serialization (default: False).
            metrics: Optional metrics collector (default: None).
            on_engine_created: Optional callback invoked with ``AsyncEngine`` after creation.
                Use for OpenTelemetry instrumentation, custom event listeners, etc.
            dispose_timeout: Maximum seconds to wait for engine disposal in :meth:`aclose`
                (default: 30.0). Lower this in tests or short-lived environments; raise it
                if you have long-running transactions that need more time to settle.
            **kwargs: Additional keyword arguments for ``create_async_engine``.
        """
        self._closed = False
        self._close_lock = asyncio.Lock()
        self._dispose_timeout = dispose_timeout
        resolved_poolclass = resolve_pool_class(poolclass)
        engine_kwargs = build_engine_kwargs(
            echo=echo,
            poolclass=resolved_poolclass,
            isolation_level=isolation_level,
            pool_settings=pool_settings,
            connect_args=connect_args,
            extra_kwargs=kwargs,
            use_orjson=use_orjson,
        )

        self._engine: AsyncEngine = create_async_engine(url, **engine_kwargs)
        self._session_maker = cast(
            async_sessionmaker[SessionT],
            async_sessionmaker(
                self._engine,
                class_=session_class or AsyncSession,
                expire_on_commit=expire_on_commit,
            ),
        )

        if metrics:
            attach_metrics(self._engine, metrics)

        if on_engine_created is not None:
            on_engine_created(self._engine)

    async def aclose(self) -> None:
        """Close the engine and all connections."""
        async with self._close_lock:
            if self._closed:
                return

            try:
                # Use shield so disposal runs even if the task is cancelled; timeout avoids indefinite hang.
                await asyncio.wait_for(
                    asyncio.shield(self._engine.dispose()),
                    timeout=self._dispose_timeout,
                )
            except TimeoutError:
                logger.warning(
                    "Engine disposal timed out after %.1f seconds. Some connections may not have closed cleanly.",
                    self._dispose_timeout,
                )
            finally:
                # Always mark as closed to prevent retry loops
                self._closed = True

    def _ensure_not_closed(self) -> None:
        """Ensure that the manager is not closed."""
        if self._closed:
            raise RuntimeError("AsyncSessionManager is closed")

    async def __aenter__(self) -> AsyncSessionManager[SessionT]:
        """Support for async context manager."""
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Close the engine on exit."""
        await self.aclose()

    @property
    def engine(self) -> AsyncEngine:
        """Get the underlying engine."""
        return self._engine

    @property
    def session_maker(self) -> async_sessionmaker[SessionT]:
        """Get the session maker."""
        return self._session_maker

    @asynccontextmanager
    async def get_session(self) -> AsyncIterator[SessionT]:
        """Get a new database session."""
        self._ensure_not_closed()
        async with self._session_maker() as session:
            yield session

    @asynccontextmanager
    async def get_transaction(self, isolation_level: str | None = None) -> AsyncIterator[SessionT]:
        """Get a new database session with automatic transaction management.

        Args:
            isolation_level: Optional isolation level for the transaction.

        Yields:
            Managed async session with active transaction.
        """
        self._ensure_not_closed()
        options = {"isolation_level": isolation_level} if isolation_level else {}
        async with (
            self._session_maker(execution_options=options) as session,
            session.begin(),
        ):
            yield session

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]