diff --git a/docs/module/db.md b/docs/module/db.md index 75f757f..3ca8ab4 100644 --- a/docs/module/db.md +++ b/docs/module/db.md @@ -31,6 +31,25 @@ async def list_users(session: AsyncSession = Depends(db)): The `Database` instance **is** the dependency: use it directly as `Depends(db)`. The whole request runs as a single transaction (CRUD writes use savepoints under it). +The **URL** may be a plain string or a Pydantic [`PostgresDsn`](https://docs.pydantic.dev/latest/api/networks/#pydantic.networks.PostgresDsn). In URL mode you can tune the engine: pass `connect_args` for DBAPI-level options and any other keyword for `create_async_engine` (e.g. `pool_size`, `echo`, `pool_pre_ping`): + +```python +from pydantic_settings import BaseSettings +from pydantic import PostgresDsn + +class Settings(BaseSettings): + database_url: PostgresDsn + +settings = Settings() + +db = Database( + settings.database_url, + pool_size=20, + pool_pre_ping=True, + connect_args={"server_settings": {"application_name": "myapp"}}, +) +``` + ## Committing before the response [`db.install(app)`](../reference/db.md#fastapi_toolsets.db.Database) adds a middleware that commits the request's session when the response starts, after the endpoint returns and before the body is sent. With the middleware installed, the dependency does not commit again. diff --git a/src/fastapi_toolsets/db/core.py b/src/fastapi_toolsets/db/core.py index 6feb413..69775ac 100644 --- a/src/fastapi_toolsets/db/core.py +++ b/src/fastapi_toolsets/db/core.py @@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator from contextlib import AbstractAsyncContextManager, asynccontextmanager from typing import Any +from pydantic import PostgresDsn from sqlalchemy import exc as sa_exc from sqlalchemy.ext.asyncio import ( AsyncEngine, @@ -84,18 +85,20 @@ class Database: untouched). Args: - url: Database connection URL (e.g. ``"postgresql+asyncpg://..."``). + url: Database connection URL. Accepts a plain string or a Pydantic + :class:`~pydantic.PostgresDsn`. engine: An existing :class:`AsyncEngine` to reuse instead of *url*. session_class: Session class for the sessionmaker (e.g. ``EventSession``). expire_on_commit: Expire attributes after commit. Defaults to ``False``. autoflush: Autoflush the session before queries. Defaults to ``True``. + connect_args: DBAPI-level connection arguments forwarded to + :func:`create_async_engine` (URL mode only). **engine_options: Extra keyword arguments forwarded to - :func:`create_async_engine` (URL mode only, e.g. ``pool_size``, - ``echo``, ``connect_args``). + :func:`create_async_engine` (URL mode only). Raises: TypeError: If neither or both of *url* and *engine* are given, or if - *engine_options* are passed together with *engine*. + *connect_args*/*engine_options* are passed together with *engine*. Example: ```python @@ -115,12 +118,13 @@ class Database: def __init__( self, - url: str | None = None, + url: str | PostgresDsn | None = None, *, engine: AsyncEngine | None = None, session_class: type[AsyncSession] = AsyncSession, expire_on_commit: bool = False, autoflush: bool = True, + connect_args: dict[str, Any] | None = None, **engine_options: Any, ) -> None: if (url is None) == (engine is None): @@ -128,10 +132,10 @@ class Database: "Database requires exactly one of 'url' or 'engine' " "(got both or neither)." ) - if engine is not None and engine_options: + if engine is not None and (engine_options or connect_args is not None): raise TypeError( - "engine_options are only valid in URL mode; configure the " - "engine you pass via 'engine=' yourself." + "connect_args/engine_options are only valid in URL mode; " + "configure the engine you pass via 'engine=' yourself." ) if engine is not None: @@ -140,7 +144,11 @@ class Database: else: assert url is not None # guaranteed by the XOR check above self._owns_engine = True - self.engine = create_async_engine(url, **engine_options) + if connect_args is not None: + engine_options["connect_args"] = connect_args + # ``PostgresDsn`` (and other URL objects) are not str subclasses, so + # coerce to the string form SQLAlchemy expects. + self.engine = create_async_engine(str(url), **engine_options) self._sessionmaker: async_sessionmaker[AsyncSession] = async_sessionmaker( self.engine, class_=session_class, diff --git a/tests/test_db.py b/tests/test_db.py index 2b393a9..609c8a1 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -3,12 +3,13 @@ import asyncio import uuid from contextlib import asynccontextmanager -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Depends, FastAPI from fastapi.responses import StreamingResponse from httpx import ASGITransport, AsyncClient +from pydantic import PostgresDsn from sqlalchemy import ( Column, ForeignKey, @@ -94,6 +95,33 @@ class TestDatabaseConstruction: with pytest.raises(TypeError): Database(engine=engine, pool_size=5) + @pytest.mark.anyio + async def test_connect_args_with_engine_raises(self, engine): + """connect_args are rejected in engine= mode.""" + with pytest.raises(TypeError): + Database(engine=engine, connect_args={"server_settings": {}}) + + @pytest.mark.anyio + async def test_accepts_postgres_dsn(self): + """A Pydantic PostgresDsn is coerced to a string URL.""" + dsn = PostgresDsn(DATABASE_URL) + db = Database(dsn) + try: + assert db._owns_engine is True + assert str(db.engine.url) == str(create_async_engine(DATABASE_URL).url) + finally: + await db.engine.dispose() + + @pytest.mark.anyio + async def test_connect_args_forwarded_to_engine(self): + """connect_args are forwarded to create_async_engine in URL mode.""" + connect_args = {"server_settings": {"application_name": "ft_test"}} + with patch("fastapi_toolsets.db.core.create_async_engine") as mocked: + mocked.return_value = MagicMock() + Database(DATABASE_URL, connect_args=connect_args) + _, kwargs = mocked.call_args + assert kwargs["connect_args"] == connect_args + @pytest.mark.anyio async def test_url_mode_owns_engine(self): """URL mode builds and owns the engine."""