1from __future__ import annotations
 2
 3import functools
 4from collections.abc import Callable
 5from typing import Any
 6
 7from ..db import use_management_connection
 8
 9
10def database_management_command[F: Callable[..., Any]](f: F) -> F:
11    """Run a click command's body through `use_management_connection()`.
12
13    Apply to CLI commands that perform schema changes, migrations, or other
14    database management operations. Inside the command, `get_connection()`
15    returns a connection opened against `POSTGRES_MANAGEMENT_URL` (falling
16    back to `POSTGRES_URL` when unset).
17    """
18
19    @functools.wraps(f)
20    def wrapper(*args: Any, **kwargs: Any) -> Any:
21        with use_management_connection():
22            return f(*args, **kwargs)
23
24    return wrapper  # ty: ignore[invalid-return-type]