Compare commits

..
6 Commits
24 changed files with 855 additions and 279 deletions
+1
View File
@@ -0,0 +1 @@
# Authentication
+31
View File
@@ -87,6 +87,37 @@ await wait_for_row_change(
)
```
## Creating a database
!!! info "Added in `v2.1`"
[`create_database`](../reference/db.md#fastapi_toolsets.db.create_database) creates a database at a given URL. It connects to *server_url* and issues a `CREATE DATABASE` statement:
```python
from fastapi_toolsets.db import create_database
SERVER_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres"
await create_database(db_name="myapp_test", server_url=SERVER_URL)
```
For test isolation with automatic cleanup, use [`create_worker_database`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_worker_database) from the `pytest` module instead — it handles drop-before, create, and drop-after automatically.
## Cleaning up tables
!!! info "Added in `v2.1`"
[`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.cleanup_tables) truncates all tables:
```python
from fastapi_toolsets.db import cleanup_tables
@pytest.fixture(autouse=True)
async def clean(db_session):
yield
await cleanup_tables(session=db_session, base=Base)
```
---
[:material-api: API Reference](../reference/db.md)
+30 -7
View File
@@ -36,7 +36,13 @@ This mounts the `/metrics` endpoint that Prometheus can scrape.
### Providers
Providers are called once at startup and register metrics that are updated externally (e.g. counters, histograms):
Providers are called once at startup by `init_metrics`. The return value (the Prometheus metric object) is stored in the registry and can be retrieved later with [`registry.get(name)`](../reference/metrics.md#fastapi_toolsets.metrics.registry.MetricsRegistry.get).
Use providers when you want **deferred initialization**: the Prometheus metric is not registered with the global `CollectorRegistry` until `init_metrics` runs, not at import time. This is particularly useful for testing — importing the module in a test suite without calling `init_metrics` leaves no metrics registered, avoiding cross-test pollution.
It is also useful when metrics are defined across multiple modules and merged with `include_registry`: any code that needs a metric can call `metrics.get()` on the shared registry instead of importing the metric directly from its origin module.
If neither of these applies to you, declaring metrics at module level (e.g. `HTTP_REQUESTS = Counter(...)`) is simpler and equally valid.
```python
from prometheus_client import Counter, Histogram
@@ -50,15 +56,32 @@ def request_duration():
return Histogram("request_duration_seconds", "Request duration")
```
### Collectors
Collectors are called on every scrape. Use them for metrics that reflect current state (e.g. gauges):
To use a provider's metric elsewhere (e.g. in a middleware), call `metrics.get()` inside the handler — **not** at module level, as providers are only initialized when `init_metrics` runs:
```python
async def metrics_middleware(request: Request, call_next):
response = await call_next(request)
metrics.get("http_requests").labels(
method=request.method, status=response.status_code
).inc()
return response
```
### Collectors
Collectors are called on every scrape. Use them for metrics that reflect current state (e.g. gauges).
!!! warning "Declare the metric at module level"
Do **not** instantiate the Prometheus metric inside the collector function. Doing so recreates it on every scrape, raising `ValueError: Duplicated timeseries in CollectorRegistry`. Declare it once at module level instead:
```python
from prometheus_client import Gauge
_queue_depth = Gauge("queue_depth", "Current queue depth")
@metrics.register(collect=True)
def queue_depth():
gauge = Gauge("queue_depth", "Current queue depth")
gauge.set(get_current_queue_depth())
def collect_queue_depth():
_queue_depth.set(get_current_queue_depth())
```
## Merging registries
+16 -4
View File
@@ -40,10 +40,10 @@ async def http_client(db_session):
## Database sessions in tests
Use [`create_db_session`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_db_session) to create an isolated `AsyncSession` for a test:
Use [`create_db_session`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_db_session) to create an isolated `AsyncSession` for a test, combined with [`create_worker_database`](../reference/pytest.md#fastapi_toolsets.pytest.utils.create_worker_database) to set up a per-worker database:
```python
from fastapi_toolsets.pytest import create_db_session, create_worker_database
from fastapi_toolsets.pytest import create_worker_database, create_db_session
@pytest.fixture(scope="session")
async def worker_db_url():
@@ -64,16 +64,28 @@ async def db_session(worker_db_url):
!!! info
In this example, the database is reset between each test using the argument `cleanup=True`.
Use [`worker_database_url`](../reference/pytest.md#fastapi_toolsets.pytest.utils.worker_database_url) to derive the per-worker URL manually if needed:
```python
from fastapi_toolsets.pytest import worker_database_url
url = worker_database_url("postgresql+asyncpg://user:pass@localhost/test_db", default_test_db="test")
# e.g. "postgresql+asyncpg://user:pass@localhost/test_db_gw0" under xdist
```
## Parallel testing with pytest-xdist
The examples above are already compatible with parallel test execution with `pytest-xdist`.
## Cleaning up tables
If you want to manually clean up a database you can use [`cleanup_tables`](../reference/pytest.md#fastapi_toolsets.pytest.utils.cleanup_tables), this will truncates all tables between tests for fast isolation:
!!! warning
Since `V2.1.0` `cleanup_tables` now live in `fastapi_toolsets.db`. For backward compatibility the function is still available in `fastapi_toolsets.pytest`, but this will be remove in `V3.0.0`.
If you want to manually clean up a database you can use [`cleanup_tables`](../reference/db.md#fastapi_toolsets.db.cleanup_tables), this will truncate all tables between tests for fast isolation:
```python
from fastapi_toolsets.pytest import cleanup_tables
from fastapi_toolsets.db import cleanup_tables
@pytest.fixture(autouse=True)
async def clean(db_session):
+6
View File
@@ -7,6 +7,8 @@ You can import them directly from `fastapi_toolsets.db`:
```python
from fastapi_toolsets.db import (
LockMode,
cleanup_tables,
create_database,
create_db_dependency,
create_db_context,
get_transaction,
@@ -26,3 +28,7 @@ from fastapi_toolsets.db import (
## ::: fastapi_toolsets.db.lock_tables
## ::: fastapi_toolsets.db.wait_for_row_change
## ::: fastapi_toolsets.db.create_database
## ::: fastapi_toolsets.db.cleanup_tables
-2
View File
@@ -24,5 +24,3 @@ from fastapi_toolsets.pytest import (
## ::: fastapi_toolsets.pytest.utils.worker_database_url
## ::: fastapi_toolsets.pytest.utils.create_worker_database
## ::: fastapi_toolsets.pytest.utils.cleanup_tables
+9
View File
@@ -0,0 +1,9 @@
from fastapi import FastAPI
from fastapi_toolsets.exceptions import init_exceptions_handlers
from .routes import router
app = FastAPI()
init_exceptions_handlers(app=app)
app.include_router(router=router)
+9
View File
@@ -0,0 +1,9 @@
from fastapi_toolsets.crud import CrudFactory
from .models import OAuthAccount, OAuthProvider, Team, User, UserToken
TeamCrud = CrudFactory(model=Team)
UserCrud = CrudFactory(model=User)
UserTokenCrud = CrudFactory(model=UserToken)
OAuthProviderCrud = CrudFactory(model=OAuthProvider)
OAuthAccountCrud = CrudFactory(model=OAuthAccount)
+15
View File
@@ -0,0 +1,15 @@
from fastapi import Depends
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from fastapi_toolsets.db import create_db_context, create_db_dependency
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost:5432/postgres"
engine = create_async_engine(url=DATABASE_URL, future=True)
async_session_maker = async_sessionmaker(bind=engine, expire_on_commit=False)
get_db = create_db_dependency(session_maker=async_session_maker)
get_db_context = create_db_context(session_maker=async_session_maker)
SessionDep = Depends(get_db)
+105
View File
@@ -0,0 +1,105 @@
import enum
from datetime import datetime
from uuid import UUID
from sqlalchemy import (
Boolean,
DateTime,
Enum,
ForeignKey,
Integer,
String,
UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from fastapi_toolsets.models import TimestampMixin, UUIDMixin
class Base(DeclarativeBase, UUIDMixin):
type_annotation_map = {
str: String(),
int: Integer(),
UUID: PG_UUID(as_uuid=True),
datetime: DateTime(timezone=True),
}
class UserRole(enum.Enum):
admin = "admin"
moderator = "moderator"
user = "user"
class Team(Base, TimestampMixin):
__tablename__ = "teams"
name: Mapped[str] = mapped_column(String, unique=True, index=True)
users: Mapped[list["User"]] = relationship(back_populates="team")
class User(Base, TimestampMixin):
__tablename__ = "users"
username: Mapped[str] = mapped_column(String, unique=True, index=True)
email: Mapped[str | None] = mapped_column(
String, unique=True, index=True, nullable=True
)
hashed_password: Mapped[str | None] = mapped_column(String, nullable=True)
is_active: Mapped[bool] = mapped_column(Boolean, default=True)
role: Mapped[UserRole] = mapped_column(Enum(UserRole), default=UserRole.user)
team_id: Mapped[UUID | None] = mapped_column(ForeignKey("teams.id"), nullable=True)
team: Mapped["Team | None"] = relationship(back_populates="users")
oauth_accounts: Mapped[list["OAuthAccount"]] = relationship(back_populates="user")
tokens: Mapped[list["UserToken"]] = relationship(back_populates="user")
class UserToken(Base, TimestampMixin):
"""API tokens for a user (multiple allowed)."""
__tablename__ = "user_tokens"
user_id: Mapped[UUID] = mapped_column(ForeignKey("users.id"))
# Store hashed token value
token_hash: Mapped[str] = mapped_column(String, unique=True, index=True)
name: Mapped[str | None] = mapped_column(String, nullable=True)
expires_at: Mapped[datetime | None] = mapped_column(
DateTime(timezone=True), nullable=True
)
user: Mapped["User"] = relationship(back_populates="tokens")
class OAuthProvider(Base, TimestampMixin):
"""Configurable OAuth2 / OpenID Connect provider."""
__tablename__ = "oauth_providers"
slug: Mapped[str] = mapped_column(String, unique=True, index=True)
name: Mapped[str] = mapped_column(String)
client_id: Mapped[str] = mapped_column(String)
client_secret: Mapped[str] = mapped_column(String)
discovery_url: Mapped[str] = mapped_column(String, nullable=False)
scopes: Mapped[str] = mapped_column(String, default="openid email profile")
is_active: Mapped[bool] = mapped_column(Boolean, default=True)
accounts: Mapped[list["OAuthAccount"]] = relationship(back_populates="provider")
class OAuthAccount(Base, TimestampMixin):
"""OAuth2 / OpenID Connect account linked to a user."""
__tablename__ = "oauth_accounts"
__table_args__ = (
UniqueConstraint("provider_id", "subject", name="uq_oauth_provider_subject"),
)
user_id: Mapped[UUID] = mapped_column(ForeignKey("users.id"))
provider_id: Mapped[UUID] = mapped_column(ForeignKey("oauth_providers.id"))
# OAuth `sub` / OpenID subject identifier
subject: Mapped[str] = mapped_column(String)
user: Mapped["User"] = relationship(back_populates="oauth_accounts")
provider: Mapped["OAuthProvider"] = relationship(back_populates="accounts")
+122
View File
@@ -0,0 +1,122 @@
from typing import Annotated
from uuid import UUID
import bcrypt
from fastapi import APIRouter, Form, HTTPException, Response, Security
from fastapi_toolsets.dependencies import PathDependency
from .crud import UserCrud, UserTokenCrud
from .db import SessionDep
from .models import OAuthProvider, User, UserToken
from .schemas import (
ApiTokenCreateRequest,
ApiTokenResponse,
RegisterRequest,
UserCreate,
UserResponse,
)
from .security import auth, cookie_auth, create_api_token
ProviderDep = PathDependency(
model=OAuthProvider,
field=OAuthProvider.slug,
session_dep=SessionDep,
param_name="slug",
)
def hash_password(password: str) -> str:
return bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
def verify_password(plain: str, hashed: str) -> bool:
return bcrypt.checkpw(plain.encode(), hashed.encode())
router = APIRouter(prefix="/auth")
@router.post("/register", response_model=UserResponse, status_code=201)
async def register(body: RegisterRequest, session: SessionDep):
existing = await UserCrud.first(
session=session, filters=[User.username == body.username]
)
if existing:
raise HTTPException(status_code=409, detail="Username already taken")
user = await UserCrud.create(
session=session,
obj=UserCreate(
username=body.username,
email=body.email,
hashed_password=hash_password(body.password),
),
)
return user
@router.post("/token", status_code=204)
async def login(
session: SessionDep,
response: Response,
username: Annotated[str, Form()],
password: Annotated[str, Form()],
):
user = await UserCrud.first(session=session, filters=[User.username == username])
if (
not user
or not user.hashed_password
or not verify_password(password, user.hashed_password)
):
raise HTTPException(status_code=401, detail="Invalid credentials")
if not user.is_active:
raise HTTPException(status_code=403, detail="Account disabled")
cookie_auth.set_cookie(response, str(user.id))
@router.post("/logout", status_code=204)
async def logout(response: Response):
cookie_auth.delete_cookie(response)
@router.get("/me", response_model=UserResponse)
async def me(user: User = Security(auth)):
return user
@router.post("/tokens", response_model=ApiTokenResponse, status_code=201)
async def create_token(
body: ApiTokenCreateRequest,
user: User = Security(auth),
):
raw, token_row = await create_api_token(
user.id, name=body.name, expires_at=body.expires_at
)
return ApiTokenResponse(
id=token_row.id,
name=token_row.name,
expires_at=token_row.expires_at,
created_at=token_row.created_at,
token=raw,
)
@router.delete("/tokens/{token_id}", status_code=204)
async def revoke_token(
session: SessionDep,
token_id: UUID,
user: User = Security(auth),
):
if not await UserTokenCrud.first(
session=session,
filters=[UserToken.id == token_id, UserToken.user_id == user.id],
):
raise HTTPException(status_code=404, detail="Token not found")
await UserTokenCrud.delete(
session=session,
filters=[UserToken.id == token_id, UserToken.user_id == user.id],
)
@@ -0,0 +1,64 @@
from datetime import datetime
from uuid import UUID
from pydantic import EmailStr
from fastapi_toolsets.schemas import PydanticBase
class RegisterRequest(PydanticBase):
username: str
password: str
email: EmailStr | None = None
class UserResponse(PydanticBase):
id: UUID
username: str
email: str | None
role: str
is_active: bool
model_config = {"from_attributes": True}
class ApiTokenCreateRequest(PydanticBase):
name: str | None = None
expires_at: datetime | None = None
class ApiTokenResponse(PydanticBase):
id: UUID
name: str | None
expires_at: datetime | None
created_at: datetime
# Only populated on creation
token: str | None = None
model_config = {"from_attributes": True}
class OAuthProviderResponse(PydanticBase):
slug: str
name: str
model_config = {"from_attributes": True}
class UserCreate(PydanticBase):
username: str
email: str | None = None
hashed_password: str | None = None
class UserTokenCreate(PydanticBase):
user_id: UUID
token_hash: str
name: str | None = None
expires_at: datetime | None = None
class OAuthAccountCreate(PydanticBase):
user_id: UUID
provider_id: UUID
subject: str
@@ -0,0 +1,100 @@
import hashlib
from datetime import datetime, timezone
from uuid import UUID
from fastapi import HTTPException
from sqlalchemy.orm import selectinload
from fastapi_toolsets.exceptions import UnauthorizedError
from fastapi_toolsets.security import (
APIKeyHeaderAuth,
BearerTokenAuth,
CookieAuth,
MultiAuth,
)
from .crud import UserCrud, UserTokenCrud
from .db import get_db_context
from .models import User, UserRole, UserToken
from .schemas import UserTokenCreate
SESSION_COOKIE = "session"
SECRET_KEY = "123456789"
def _hash_token(token: str) -> str:
return hashlib.sha256(token.encode()).hexdigest()
async def _verify_token(token: str, role: UserRole | None = None) -> User:
async with get_db_context() as db:
user_token = await UserTokenCrud.first(
session=db,
filters=[UserToken.token_hash == _hash_token(token)],
load_options=[selectinload(UserToken.user)],
)
if user_token is None or not user_token.user.is_active:
raise UnauthorizedError()
if user_token.expires_at and user_token.expires_at < datetime.now(timezone.utc):
raise UnauthorizedError()
user = user_token.user
if role is not None and user.role != role:
raise HTTPException(status_code=403, detail="Insufficient permissions")
return user
async def _verify_cookie(user_id: str, role: UserRole | None = None) -> User:
async with get_db_context() as db:
user = await UserCrud.first(
session=db,
filters=[User.id == UUID(user_id)],
)
if not user or not user.is_active:
raise UnauthorizedError()
if role is not None and user.role != role:
raise HTTPException(status_code=403, detail="Insufficient permissions")
return user
bearer_auth = BearerTokenAuth(
validator=_verify_token,
prefix="ctf_",
)
header_auth = APIKeyHeaderAuth(
name="X-API-Key",
validator=_verify_token,
)
cookie_auth = CookieAuth(
name=SESSION_COOKIE,
validator=_verify_cookie,
secret_key=SECRET_KEY,
)
auth = MultiAuth(bearer_auth, header_auth, cookie_auth)
async def create_api_token(
user_id: UUID,
*,
name: str | None = None,
expires_at: datetime | None = None,
) -> tuple[str, UserToken]:
raw = bearer_auth.generate_token()
async with get_db_context() as db:
token_row = await UserTokenCrud.create(
session=db,
obj=UserTokenCreate(
user_id=user_id,
token_hash=_hash_token(raw),
name=name,
expires_at=expires_at,
),
)
return raw, token_row
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "fastapi-toolsets"
version = "2.0.0"
version = "2.1.0"
description = "Production-ready utilities for FastAPI applications"
readme = "README.md"
license = "MIT"
+1 -1
View File
@@ -21,4 +21,4 @@ Example usage:
return Response(data={"user": user.username}, message="Success")
"""
__version__ = "2.0.0"
__version__ = "2.1.0"
+69 -2
View File
@@ -7,17 +7,19 @@ from enum import Enum
from typing import Any, TypeVar
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase
from .exceptions import NotFoundError
__all__ = [
"LockMode",
"cleanup_tables",
"create_database",
"create_db_context",
"create_db_dependency",
"lock_tables",
"get_transaction",
"lock_tables",
"wait_for_row_change",
]
@@ -188,6 +190,71 @@ async def lock_tables(
yield session
async def create_database(
db_name: str,
*,
server_url: str,
) -> None:
"""Create a database.
Connects to *server_url* using ``AUTOCOMMIT`` isolation and issues a
``CREATE DATABASE`` statement for *db_name*.
Args:
db_name: Name of the database to create.
server_url: URL used for server-level DDL (must point to an existing
database on the same server).
Example:
```python
from fastapi_toolsets.db import create_database
SERVER_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres"
await create_database("myapp_test", server_url=SERVER_URL)
```
"""
engine = create_async_engine(server_url, isolation_level="AUTOCOMMIT")
try:
async with engine.connect() as conn:
await conn.execute(text(f"CREATE DATABASE {db_name}"))
finally:
await engine.dispose()
async def cleanup_tables(
session: AsyncSession,
base: type[DeclarativeBase],
) -> None:
"""Truncate all tables for fast between-test cleanup.
Executes a single ``TRUNCATE … RESTART IDENTITY CASCADE`` statement
across every table in *base*'s metadata, which is significantly faster
than dropping and re-creating tables between tests.
This is a no-op when the metadata contains no tables.
Args:
session: An active async database session.
base: SQLAlchemy DeclarativeBase class containing model metadata.
Example:
```python
@pytest.fixture
async def db_session(worker_db_url):
async with create_db_session(worker_db_url, Base) as session:
yield session
await cleanup_tables(session, Base)
```
"""
tables = base.metadata.sorted_tables
if not tables:
return
table_names = ", ".join(f'"{t.name}"' for t in tables)
await session.execute(text(f"TRUNCATE {table_names} RESTART IDENTITY CASCADE"))
await session.commit()
_M = TypeVar("_M", bound=DeclarativeBase)
+1 -1
View File
@@ -51,7 +51,7 @@ def init_metrics(
"""
for provider in registry.get_providers():
logger.debug("Initialising metric provider '%s'", provider.name)
provider.func()
registry._instances[provider.name] = provider.func()
# Partition collectors and cache env check at startup — both are stable for the app lifetime.
async_collectors = [
+21 -45
View File
@@ -19,31 +19,11 @@ class Metric:
class MetricsRegistry:
"""Registry for managing Prometheus metric providers and collectors.
Example:
```python
from prometheus_client import Counter, Gauge
from fastapi_toolsets.metrics import MetricsRegistry
metrics = MetricsRegistry()
@metrics.register
def http_requests():
return Counter("http_requests_total", "Total HTTP requests", ["method", "status"])
@metrics.register(name="db_pool")
def database_pool_size():
return Gauge("db_pool_size", "Database connection pool size")
@metrics.register(collect=True)
def collect_queue_depth(gauge=Gauge("queue_depth", "Current queue depth")):
gauge.set(get_current_queue_depth())
```
"""
"""Registry for managing Prometheus metric providers and collectors."""
def __init__(self) -> None:
self._metrics: dict[str, Metric] = {}
self._instances: dict[str, Any] = {}
def register(
self,
@@ -61,17 +41,6 @@ class MetricsRegistry:
name: Metric name (defaults to function name).
collect: If ``True``, the function is called on every scrape.
If ``False`` (default), called once at init time.
Example:
```python
@metrics.register
def my_counter():
return Counter("my_counter", "A counter")
@metrics.register(collect=True, name="queue")
def collect_queue_depth():
gauge.set(compute_depth())
```
"""
def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
@@ -87,6 +56,25 @@ class MetricsRegistry:
return decorator(func)
return decorator
def get(self, name: str) -> Any:
"""Return the metric instance created by a provider.
Args:
name: The metric name (defaults to the provider function name).
Raises:
KeyError: If the metric name is unknown or ``init_metrics`` has not
been called yet.
"""
if name not in self._instances:
if name in self._metrics:
raise KeyError(
f"Metric '{name}' exists but has not been initialized yet. "
"Ensure init_metrics() has been called before accessing metric instances."
)
raise KeyError(f"Unknown metric '{name}'.")
return self._instances[name]
def include_registry(self, registry: "MetricsRegistry") -> None:
"""Include another :class:`MetricsRegistry` into this one.
@@ -95,18 +83,6 @@ class MetricsRegistry:
Raises:
ValueError: If a metric name already exists in the current registry.
Example:
```python
main = MetricsRegistry()
sub = MetricsRegistry()
@sub.register
def sub_metric():
return Counter("sub_total", "Sub counter")
main.include_registry(sub)
```
"""
for metric_name, definition in registry._metrics.items():
if metric_name in self._metrics:
+129 -159
View File
@@ -1,12 +1,12 @@
"""Pytest helper utilities for FastAPI testing."""
import os
import warnings
from collections.abc import AsyncGenerator, Callable
from contextlib import asynccontextmanager
from typing import Any
from httpx import ASGITransport, AsyncClient
from sqlalchemy import text
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import (
AsyncSession,
@@ -15,7 +15,134 @@ from sqlalchemy.ext.asyncio import (
)
from sqlalchemy.orm import DeclarativeBase
from ..db import create_db_context
from sqlalchemy import text
from ..db import (
cleanup_tables as _cleanup_tables,
create_database,
create_db_context,
)
async def cleanup_tables(
session: AsyncSession,
base: type[DeclarativeBase],
) -> None:
"""Truncate all tables for fast between-test cleanup.
.. deprecated::
Import ``cleanup_tables`` from ``fastapi_toolsets.db`` instead.
This re-export will be removed in v3.0.0.
"""
warnings.warn(
"Importing cleanup_tables from fastapi_toolsets.pytest is deprecated "
"and will be removed in v3.0.0. "
"Use 'from fastapi_toolsets.db import cleanup_tables' instead.",
DeprecationWarning,
stacklevel=2,
)
await _cleanup_tables(session=session, base=base)
def _get_xdist_worker(default_test_db: str) -> str:
"""Return the pytest-xdist worker name, or *default_test_db* when not running under xdist.
Reads the ``PYTEST_XDIST_WORKER`` environment variable that xdist sets
automatically in each worker process (e.g. ``"gw0"``, ``"gw1"``).
When xdist is not installed or not active, the variable is absent and
*default_test_db* is returned instead.
Args:
default_test_db: Fallback value returned when ``PYTEST_XDIST_WORKER``
is not set.
"""
return os.environ.get("PYTEST_XDIST_WORKER", default_test_db)
def worker_database_url(database_url: str, default_test_db: str) -> str:
"""Derive a per-worker database URL for pytest-xdist parallel runs.
Appends ``_{worker_name}`` to the database name so each xdist worker
operates on its own database. When not running under xdist,
``_{default_test_db}`` is appended instead.
The worker name is read from the ``PYTEST_XDIST_WORKER`` environment
variable (set automatically by xdist in each worker process).
Args:
database_url: Original database connection URL.
default_test_db: Suffix appended to the database name when
``PYTEST_XDIST_WORKER`` is not set.
Returns:
A database URL with a worker- or default-specific database name.
"""
worker = _get_xdist_worker(default_test_db=default_test_db)
url = make_url(database_url)
url = url.set(database=f"{url.database}_{worker}")
return url.render_as_string(hide_password=False)
@asynccontextmanager
async def create_worker_database(
database_url: str,
default_test_db: str = "test_db",
) -> AsyncGenerator[str, None]:
"""Create and drop a per-worker database for pytest-xdist isolation.
Derives a worker-specific database URL using :func:`worker_database_url`,
then delegates to :func:`~fastapi_toolsets.db.create_database` to create
and drop it. Intended for use as a **session-scoped** fixture.
When running under xdist the database name is suffixed with the worker
name (e.g. ``_gw0``). Otherwise it is suffixed with *default_test_db*.
Args:
database_url: Original database connection URL (used as the server
connection and as the base for the worker database name).
default_test_db: Suffix appended to the database name when
``PYTEST_XDIST_WORKER`` is not set. Defaults to ``"test_db"``.
Yields:
The worker-specific database URL.
Example:
```python
from fastapi_toolsets.pytest import create_worker_database, create_db_session
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/test_db"
@pytest.fixture(scope="session")
async def worker_db_url():
async with create_worker_database(DATABASE_URL) as url:
yield url
@pytest.fixture
async def db_session(worker_db_url):
async with create_db_session(
worker_db_url, Base, cleanup=True
) as session:
yield session
```
"""
worker_url = worker_database_url(
database_url=database_url, default_test_db=default_test_db
)
worker_db_name: str = make_url(worker_url).database # type: ignore[assignment]
engine = create_async_engine(database_url, isolation_level="AUTOCOMMIT")
try:
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
await create_database(db_name=worker_db_name, server_url=database_url)
yield worker_url
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
finally:
await engine.dispose()
@asynccontextmanager
@@ -156,160 +283,3 @@ async def create_db_session(
await conn.run_sync(base.metadata.drop_all)
finally:
await engine.dispose()
def _get_xdist_worker(default_test_db: str) -> str:
"""Return the pytest-xdist worker name, or *default_test_db* when not running under xdist.
Reads the ``PYTEST_XDIST_WORKER`` environment variable that xdist sets
automatically in each worker process (e.g. ``"gw0"``, ``"gw1"``).
When xdist is not installed or not active, the variable is absent and
*default_test_db* is returned instead.
Args:
default_test_db: Fallback value returned when ``PYTEST_XDIST_WORKER``
is not set.
"""
return os.environ.get("PYTEST_XDIST_WORKER", default_test_db)
def worker_database_url(database_url: str, default_test_db: str) -> str:
"""Derive a per-worker database URL for pytest-xdist parallel runs.
Appends ``_{worker_name}`` to the database name so each xdist worker
operates on its own database. When not running under xdist,
``_{default_test_db}`` is appended instead.
The worker name is read from the ``PYTEST_XDIST_WORKER`` environment
variable (set automatically by xdist in each worker process).
Args:
database_url: Original database connection URL.
default_test_db: Suffix appended to the database name when
``PYTEST_XDIST_WORKER`` is not set.
Returns:
A database URL with a worker- or default-specific database name.
Example:
```python
# With PYTEST_XDIST_WORKER="gw0":
url = worker_database_url(
"postgresql+asyncpg://user:pass@localhost/test_db",
default_test_db="test",
)
# "postgresql+asyncpg://user:pass@localhost/test_db_gw0"
# Without PYTEST_XDIST_WORKER:
url = worker_database_url(
"postgresql+asyncpg://user:pass@localhost/test_db",
default_test_db="test",
)
# "postgresql+asyncpg://user:pass@localhost/test_db_test"
```
"""
worker = _get_xdist_worker(default_test_db=default_test_db)
url = make_url(database_url)
url = url.set(database=f"{url.database}_{worker}")
return url.render_as_string(hide_password=False)
@asynccontextmanager
async def create_worker_database(
database_url: str,
default_test_db: str = "test_db",
) -> AsyncGenerator[str, None]:
"""Create and drop a per-worker database for pytest-xdist isolation.
Intended for use as a **session-scoped** fixture. Connects to the server
using the original *database_url* (with ``AUTOCOMMIT`` isolation for DDL),
creates a dedicated database for the worker, and yields the worker-specific
URL. On cleanup the worker database is dropped.
When running under xdist the database name is suffixed with the worker
name (e.g. ``_gw0``). Otherwise it is suffixed with *default_test_db*.
Args:
database_url: Original database connection URL.
default_test_db: Suffix appended to the database name when
``PYTEST_XDIST_WORKER`` is not set. Defaults to ``"test_db"``.
Yields:
The worker-specific database URL.
Example:
```python
from fastapi_toolsets.pytest import (
create_worker_database, create_db_session,
)
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/test_db"
@pytest.fixture(scope="session")
async def worker_db_url():
async with create_worker_database(DATABASE_URL) as url:
yield url
@pytest.fixture
async def db_session(worker_db_url):
async with create_db_session(
worker_db_url, Base, cleanup=True
) as session:
yield session
```
"""
worker_url = worker_database_url(
database_url=database_url, default_test_db=default_test_db
)
worker_db_name = make_url(worker_url).database
engine = create_async_engine(
database_url,
isolation_level="AUTOCOMMIT",
)
try:
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
await conn.execute(text(f"CREATE DATABASE {worker_db_name}"))
yield worker_url
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
finally:
await engine.dispose()
async def cleanup_tables(
session: AsyncSession,
base: type[DeclarativeBase],
) -> None:
"""Truncate all tables for fast between-test cleanup.
Executes a single ``TRUNCATE … RESTART IDENTITY CASCADE`` statement
across every table in *base*'s metadata, which is significantly faster
than dropping and re-creating tables between tests.
This is a no-op when the metadata contains no tables.
Args:
session: An active async database session.
base: SQLAlchemy DeclarativeBase class containing model metadata.
Example:
```python
@pytest.fixture
async def db_session(worker_db_url):
async with create_db_session(worker_db_url, Base) as session:
yield session
await cleanup_tables(session, Base)
```
"""
tables = base.metadata.sorted_tables
if not tables:
return
table_names = ", ".join(f'"{t.name}"' for t in tables)
await session.execute(text(f"TRUNCATE {table_names} RESTART IDENTITY CASCADE"))
await session.commit()
+87 -1
View File
@@ -4,10 +4,15 @@ import asyncio
import uuid
import pytest
from sqlalchemy import text
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase
from fastapi_toolsets.db import (
LockMode,
cleanup_tables,
create_database,
create_db_context,
create_db_dependency,
get_transaction,
@@ -15,8 +20,9 @@ from fastapi_toolsets.db import (
wait_for_row_change,
)
from fastapi_toolsets.exceptions import NotFoundError
from fastapi_toolsets.pytest import create_db_session
from .conftest import DATABASE_URL, Base, Role, RoleCrud, User
from .conftest import DATABASE_URL, Base, Role, RoleCrud, User, UserCrud
class TestCreateDbDependency:
@@ -344,3 +350,83 @@ class TestWaitForRowChange:
with pytest.raises(NotFoundError):
await wait_for_row_change(db_session, Role, role.id, interval=0.05)
await delete_task
class TestCreateDatabase:
"""Tests for create_database."""
@pytest.mark.anyio
async def test_creates_database(self):
"""Database is created by create_database."""
target_url = (
make_url(DATABASE_URL)
.set(database="test_create_db_general")
.render_as_string(hide_password=False)
)
expected_db: str = make_url(target_url).database # type: ignore[assignment]
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
try:
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {expected_db}"))
await create_database(db_name=expected_db, server_url=DATABASE_URL)
async with engine.connect() as conn:
result = await conn.execute(
text("SELECT 1 FROM pg_database WHERE datname = :name"),
{"name": expected_db},
)
assert result.scalar() == 1
# Cleanup
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {expected_db}"))
finally:
await engine.dispose()
class TestCleanupTables:
"""Tests for cleanup_tables helper."""
@pytest.mark.anyio
async def test_truncates_all_tables(self):
"""All table rows are removed after cleanup_tables."""
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
role = Role(id=uuid.uuid4(), name="cleanup_role")
session.add(role)
await session.flush()
user = User(
id=uuid.uuid4(),
username="cleanup_user",
email="cleanup@test.com",
role_id=role.id,
)
session.add(user)
await session.commit()
# Verify rows exist
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 1
assert users_count == 1
await cleanup_tables(session, Base)
# Verify tables are empty
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 0
assert users_count == 0
@pytest.mark.anyio
async def test_noop_for_empty_metadata(self):
"""cleanup_tables does not raise when metadata has no tables."""
class EmptyBase(DeclarativeBase):
pass
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
# Should not raise
await cleanup_tables(session, EmptyBase)
+36
View File
@@ -159,6 +159,42 @@ class TestMetricsRegistry:
assert registry.get_all()[0].func is second
class TestGet:
"""Tests for MetricsRegistry.get method."""
def test_get_returns_instance_after_init(self):
"""get() returns the metric instance stored by init_metrics."""
app = FastAPI()
registry = MetricsRegistry()
@registry.register
def my_gauge():
return Gauge("get_test_gauge", "A test gauge")
init_metrics(app, registry)
instance = registry.get("my_gauge")
assert isinstance(instance, Gauge)
def test_get_raises_for_registered_but_not_initialized(self):
"""get() raises KeyError with an informative message when init_metrics was not called."""
registry = MetricsRegistry()
@registry.register
def my_counter():
return Counter("get_uninit_counter", "A counter")
with pytest.raises(KeyError, match="not been initialized yet"):
registry.get("my_counter")
def test_get_raises_for_unknown_name(self):
"""get() raises KeyError when the metric name is not registered at all."""
registry = MetricsRegistry()
with pytest.raises(KeyError, match="Unknown metric"):
registry.get("nonexistent")
class TestIncludeRegistry:
"""Tests for MetricsRegistry.include_registry method."""
+1 -55
View File
@@ -8,11 +8,10 @@ from httpx import AsyncClient
from sqlalchemy import select, text
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import DeclarativeBase, selectinload
from sqlalchemy.orm import selectinload
from fastapi_toolsets.fixtures import Context, FixtureRegistry
from fastapi_toolsets.pytest import (
cleanup_tables,
create_async_client,
create_db_session,
create_worker_database,
@@ -406,7 +405,6 @@ class TestCreateWorkerDatabase:
) as url:
assert make_url(url).database == expected_db
# Verify the database exists while inside the context
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
result = await conn.execute(
@@ -416,7 +414,6 @@ class TestCreateWorkerDatabase:
assert result.scalar() == 1
await engine.dispose()
# After context exit the database should be dropped
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
result = await conn.execute(
@@ -439,7 +436,6 @@ class TestCreateWorkerDatabase:
async with create_worker_database(DATABASE_URL) as url:
assert make_url(url).database == expected_db
# Verify the database exists while inside the context
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
result = await conn.execute(
@@ -449,7 +445,6 @@ class TestCreateWorkerDatabase:
assert result.scalar() == 1
await engine.dispose()
# After context exit the database should be dropped
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
result = await conn.execute(
@@ -467,18 +462,15 @@ class TestCreateWorkerDatabase:
worker_database_url(DATABASE_URL, default_test_db="unused")
).database
# Pre-create the database to simulate a stale leftover
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {expected_db}"))
await conn.execute(text(f"CREATE DATABASE {expected_db}"))
await engine.dispose()
# Should succeed despite the database already existing
async with create_worker_database(DATABASE_URL) as url:
assert make_url(url).database == expected_db
# Verify cleanup after context exit
engine = create_async_engine(DATABASE_URL, isolation_level="AUTOCOMMIT")
async with engine.connect() as conn:
result = await conn.execute(
@@ -487,49 +479,3 @@ class TestCreateWorkerDatabase:
)
assert result.scalar() is None
await engine.dispose()
class TestCleanupTables:
"""Tests for cleanup_tables helper."""
@pytest.mark.anyio
async def test_truncates_all_tables(self):
"""All table rows are removed after cleanup_tables."""
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
role = Role(id=uuid.uuid4(), name="cleanup_role")
session.add(role)
await session.flush()
user = User(
id=uuid.uuid4(),
username="cleanup_user",
email="cleanup@test.com",
role_id=role.id,
)
session.add(user)
await session.commit()
# Verify rows exist
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 1
assert users_count == 1
await cleanup_tables(session, Base)
# Verify tables are empty
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 0
assert users_count == 0
@pytest.mark.anyio
async def test_noop_for_empty_metadata(self):
"""cleanup_tables does not raise when metadata has no tables."""
class EmptyBase(DeclarativeBase):
pass
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
# Should not raise
await cleanup_tables(session, EmptyBase)
Generated
+1 -1
View File
@@ -251,7 +251,7 @@ wheels = [
[[package]]
name = "fastapi-toolsets"
version = "2.0.0"
version = "2.1.0"
source = { editable = "." }
dependencies = [
{ name = "asyncpg" },