chore: Database enhancement (#329)

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.

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,

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."""