84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
 
|
"""PostgreSQL access: one pool for the app role, plain connections for other roles."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import threading
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Literal
|
||
|
|
|
||
|
|
import psycopg
|
||
|
|
from psycopg_pool import ConnectionPool
|
||
|
|
|
||
|
|
from .settings import Settings, get_settings
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
Role = Literal["app", "worker", "migrate"]
|
||
|
|
MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations"
|
||
|
|
APP_NAMES: dict[str, str] = {"app": "ps-app", "worker": "ps-worker", "migrate": "ps-migrate"}
|
||
|
|
|
||
|
|
|
||
|
|
class DatabaseUnavailable(Exception):
|
||
|
|
"""No connection could be made (bad secret, server down, pool timeout)."""
|
||
|
|
|
||
|
|
|
||
|
|
def connect_kwargs(role: Role, settings: Settings | None = None) -> dict[str, Any]:
|
||
|
|
"""Connection keyword arguments for a role. The worker logs in as the app user."""
|
||
|
|
s = settings or get_settings()
|
||
|
|
db_role = "migrate" if role == "migrate" else "app"
|
||
|
|
return {
|
||
|
|
"host": s.pg_host,
|
||
|
|
"port": s.pg_port,
|
||
|
|
"dbname": s.pg_database,
|
||
|
|
"user": s.pg_user(db_role),
|
||
|
|
"password": s.pg_password(db_role),
|
||
|
|
"application_name": APP_NAMES[role],
|
||
|
|
"connect_timeout": 5,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def connect_as(role: Role, settings: Settings | None = None, **extra: Any) -> psycopg.Connection[Any]:
|
||
|
|
"""A new connection as the given role; the caller closes it."""
|
||
|
|
kwargs = connect_kwargs(role, settings)
|
||
|
|
kwargs.update(extra)
|
||
|
|
return psycopg.connect(**kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
_pool: ConnectionPool[Any] | None = None
|
||
|
|
_pool_lock = threading.Lock()
|
||
|
|
|
||
|
|
|
||
|
|
def get_pool(settings: Settings | None = None) -> ConnectionPool[Any]:
|
||
|
|
"""The app role's pool, opened on first use without waiting for the server."""
|
||
|
|
global _pool
|
||
|
|
with _pool_lock:
|
||
|
|
if _pool is None:
|
||
|
|
try:
|
||
|
|
kwargs = connect_kwargs("app", settings)
|
||
|
|
except Exception as e:
|
||
|
|
raise DatabaseUnavailable(str(e)) from None
|
||
|
|
_pool = ConnectionPool(kwargs=kwargs, min_size=1, max_size=10, open=False, name="ps-app")
|
||
|
|
_pool.open(wait=False)
|
||
|
|
return _pool
|
||
|
|
|
||
|
|
|
||
|
|
def close_pool() -> None:
|
||
|
|
global _pool
|
||
|
|
with _pool_lock:
|
||
|
|
if _pool is not None:
|
||
|
|
_pool.close()
|
||
|
|
_pool = None
|
||
|
|
|
||
|
|
|
||
|
|
def ping(settings: Settings | None = None, timeout: float = 3.0) -> None:
|
||
|
|
"""SELECT 1 as the app role; raises DatabaseUnavailable."""
|
||
|
|
try:
|
||
|
|
pool = get_pool(settings)
|
||
|
|
with pool.connection(timeout=timeout) as conn:
|
||
|
|
conn.execute("SELECT 1").fetchone()
|
||
|
|
except DatabaseUnavailable:
|
||
|
|
raise
|
||
|
|
except Exception as e:
|
||
|
|
raise DatabaseUnavailable(type(e).__name__) from None
|