stashed
parent
4d84e5c5c9
commit
bb935d4708
|
|
@ -0,0 +1,6 @@
|
||||||
|
# Generated by alembic_setup.py. The URL is injected from the environment at runtime.
|
||||||
|
[alembic]
|
||||||
|
script_location = migrations
|
||||||
|
prepend_sys_path = .
|
||||||
|
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||||
|
timezone = UTC
|
||||||
|
|
@ -0,0 +1,318 @@
|
||||||
|
"""Alembic wiring: scaffolding, autogenerate and applying migrations.
|
||||||
|
|
||||||
|
`backend/alembic.ini` and `backend/migrations/` are generated on first use and are
|
||||||
|
never overwritten, so the plain `alembic` CLI works alongside `db_setup.init_db()`.
|
||||||
|
Models are discovered automatically: every `<package>/models.py` under `backend/`
|
||||||
|
is imported before the metadata is diffed against the live schema.
|
||||||
|
|
||||||
|
python alembic_setup.py [migrate|revision|upgrade|downgrade|current|head] [-m MSG] [-r REV]
|
||||||
|
|
||||||
|
Named `alembic_setup` rather than `alembic` because `backend/` is on `sys.path`,
|
||||||
|
so a module called `alembic.py` would shadow the installed package.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import sys
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
from importlib import import_module
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, AsyncIterator, Callable, Sequence
|
||||||
|
|
||||||
|
from alembic import command
|
||||||
|
from alembic.autogenerate import compare_metadata
|
||||||
|
from alembic.config import Config
|
||||||
|
from alembic.runtime.migration import MigrationContext
|
||||||
|
from alembic.script import ScriptDirectory
|
||||||
|
from sqlalchemy import MetaData, text
|
||||||
|
from sqlalchemy.engine import Connection
|
||||||
|
|
||||||
|
from db_setup import BASE_DIR, Base, close_db, database_url, get_engine, get_settings
|
||||||
|
|
||||||
|
logger = logging.getLogger("db.alembic")
|
||||||
|
|
||||||
|
INI = BASE_DIR / "alembic.ini"
|
||||||
|
MIGRATIONS = BASE_DIR / "migrations"
|
||||||
|
VERSIONS = MIGRATIONS / "versions"
|
||||||
|
SKIP_DIRS = {"migrations", "__pycache__", "tests", "test", ".venv", "venv", "node_modules"}
|
||||||
|
LOCK_ID = 8_412_557_390_112_004 # any 64-bit constant; serialises startup migrations
|
||||||
|
|
||||||
|
INI_TEMPLATE = """\
|
||||||
|
# Generated by alembic_setup.py. The URL is injected from the environment at runtime.
|
||||||
|
[alembic]
|
||||||
|
script_location = migrations
|
||||||
|
prepend_sys_path = .
|
||||||
|
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||||
|
timezone = UTC
|
||||||
|
"""
|
||||||
|
|
||||||
|
ENV_TEMPLATE = '''"""Alembic environment -- generated by alembic_setup.py."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from alembic import context
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
import alembic_setup as setup # noqa: E402
|
||||||
|
import db_setup # noqa: E402
|
||||||
|
|
||||||
|
metadata = setup.target_metadata()
|
||||||
|
options = setup.context_options()
|
||||||
|
|
||||||
|
|
||||||
|
def run(connection) -> None:
|
||||||
|
context.configure(connection=connection, target_metadata=metadata, **options)
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
|
||||||
|
|
||||||
|
if context.is_offline_mode():
|
||||||
|
context.configure(
|
||||||
|
url=db_setup.database_url(async_driver=False),
|
||||||
|
target_metadata=metadata,
|
||||||
|
literal_binds=True,
|
||||||
|
dialect_opts={"paramstyle": "named"},
|
||||||
|
**options,
|
||||||
|
)
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
elif (connection := context.config.attributes.get("connection")) is not None:
|
||||||
|
run(connection) # alembic_setup passed an already-open connection
|
||||||
|
else:
|
||||||
|
asyncio.run(setup.run_standalone(run)) # bare `alembic` CLI
|
||||||
|
'''
|
||||||
|
|
||||||
|
MAKO_TEMPLATE = '''"""${message}
|
||||||
|
|
||||||
|
Revision ID: ${up_revision}
|
||||||
|
Revises: ${down_revision | comma,n}
|
||||||
|
Create Date: ${create_date}
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
${imports if imports else ""}
|
||||||
|
|
||||||
|
revision: str = ${repr(up_revision)}
|
||||||
|
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||||
|
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
${upgrades if upgrades else "pass"}
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
${downgrades if downgrades else "pass"}
|
||||||
|
'''
|
||||||
|
|
||||||
|
|
||||||
|
def scaffold() -> None:
|
||||||
|
"""Create the Alembic layout if missing; existing files are left alone."""
|
||||||
|
VERSIONS.mkdir(parents=True, exist_ok=True)
|
||||||
|
(VERSIONS / ".gitkeep").touch()
|
||||||
|
for path, content in (
|
||||||
|
(INI, INI_TEMPLATE),
|
||||||
|
(MIGRATIONS / "env.py", ENV_TEMPLATE),
|
||||||
|
(MIGRATIONS / "script.py.mako", MAKO_TEMPLATE),
|
||||||
|
):
|
||||||
|
if not path.exists():
|
||||||
|
path.write_text(content, encoding="utf-8")
|
||||||
|
logger.info("created %s", path)
|
||||||
|
|
||||||
|
|
||||||
|
_imported = False
|
||||||
|
|
||||||
|
|
||||||
|
def target_metadata() -> MetaData:
|
||||||
|
"""`Base.metadata` with every discovered model module imported onto it."""
|
||||||
|
global _imported
|
||||||
|
if _imported:
|
||||||
|
return Base.metadata
|
||||||
|
if str(BASE_DIR) not in sys.path:
|
||||||
|
sys.path.insert(0, str(BASE_DIR))
|
||||||
|
names = get_settings().db_model_modules or [
|
||||||
|
".".join(p.relative_to(BASE_DIR).with_suffix("").parts).removesuffix(".__init__")
|
||||||
|
for p in [*BASE_DIR.rglob("models.py"), *BASE_DIR.rglob("models/__init__.py")]
|
||||||
|
if not SKIP_DIRS & set(p.relative_to(BASE_DIR).parts)
|
||||||
|
]
|
||||||
|
for name in names:
|
||||||
|
try:
|
||||||
|
import_module(name)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning("skipping model module %s: %s", name, exc)
|
||||||
|
_imported = True
|
||||||
|
return Base.metadata
|
||||||
|
|
||||||
|
|
||||||
|
def config(connection: Connection | None = None) -> Config:
|
||||||
|
scaffold()
|
||||||
|
cfg = Config(str(INI))
|
||||||
|
cfg.set_main_option("script_location", str(MIGRATIONS))
|
||||||
|
# ConfigParser interpolates '%', which is legal inside a password.
|
||||||
|
cfg.set_main_option("sqlalchemy.url", database_url(async_driver=False).replace("%", "%%"))
|
||||||
|
if connection is not None:
|
||||||
|
cfg.attributes["connection"] = connection
|
||||||
|
return cfg
|
||||||
|
|
||||||
|
|
||||||
|
VERSION_TABLE = "alembic_version"
|
||||||
|
|
||||||
|
|
||||||
|
def _include_object(obj: Any, name: str, type_: str, reflected: bool, compare_to: Any) -> bool:
|
||||||
|
"""Keep autogenerate inside the schemas this application owns."""
|
||||||
|
s = get_settings()
|
||||||
|
if type_ != "table":
|
||||||
|
return True
|
||||||
|
if name == VERSION_TABLE: # Alembic's own bookkeeping; never ours to alter
|
||||||
|
return False
|
||||||
|
return not s.db_schemas or (obj.schema or s.db_default_schema) in s.db_schemas
|
||||||
|
|
||||||
|
|
||||||
|
def _skip_empty(context_: Any, revision: Any, directives: list[Any]) -> None:
|
||||||
|
"""Write no revision file at all when the models and the database already agree."""
|
||||||
|
ops = directives[0].upgrade_ops if directives else None
|
||||||
|
if ops is not None and ops.is_empty():
|
||||||
|
directives[:] = []
|
||||||
|
logger.info("no model changes detected")
|
||||||
|
|
||||||
|
|
||||||
|
def context_options() -> dict[str, Any]:
|
||||||
|
"""Shared by `migrations/env.py` and the in-process drift check."""
|
||||||
|
s = get_settings()
|
||||||
|
return {
|
||||||
|
"compare_type": True,
|
||||||
|
"compare_server_default": True,
|
||||||
|
"include_schemas": bool(s.db_schemas),
|
||||||
|
"version_table_schema": s.db_default_schema or None,
|
||||||
|
"include_object": _include_object,
|
||||||
|
"process_revision_directives": _skip_empty,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
async def _run(fn: Callable[[Connection], Any]) -> Any:
|
||||||
|
"""Run a synchronous Alembic call on the async engine's connection."""
|
||||||
|
async with get_engine().connect() as conn:
|
||||||
|
result = await conn.run_sync(fn)
|
||||||
|
await conn.commit()
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
async def run_standalone(fn: Callable[[Connection], Any]) -> None:
|
||||||
|
"""Entry point for the bare `alembic` CLI, which brings no connection of its own."""
|
||||||
|
try:
|
||||||
|
await _run(fn)
|
||||||
|
finally:
|
||||||
|
await close_db()
|
||||||
|
|
||||||
|
|
||||||
|
def head() -> str | None:
|
||||||
|
"""The latest revision on disk."""
|
||||||
|
heads = ScriptDirectory.from_config(config()).get_heads()
|
||||||
|
return heads[0] if heads else None
|
||||||
|
|
||||||
|
|
||||||
|
async def current() -> str | None:
|
||||||
|
"""The revision the database is stamped with."""
|
||||||
|
opts = {"version_table_schema": get_settings().db_default_schema or None}
|
||||||
|
return await _run(lambda c: MigrationContext.configure(c, opts=opts).get_current_revision())
|
||||||
|
|
||||||
|
|
||||||
|
async def upgrade(revision: str = "head") -> None:
|
||||||
|
await _run(lambda c: command.upgrade(config(c), revision))
|
||||||
|
logger.info("upgraded to %s", revision)
|
||||||
|
|
||||||
|
|
||||||
|
async def downgrade(revision: str = "-1") -> None:
|
||||||
|
await _run(lambda c: command.downgrade(config(c), revision))
|
||||||
|
logger.info("downgraded to %s", revision)
|
||||||
|
|
||||||
|
|
||||||
|
async def autogenerate(message: str = "auto") -> str | None:
|
||||||
|
"""Write a revision if the models have drifted; return its id, or None."""
|
||||||
|
opts = {k: v for k, v in context_options().items() if k != "process_revision_directives"}
|
||||||
|
diffs = await _run(
|
||||||
|
lambda c: compare_metadata(MigrationContext.configure(c, opts=opts), target_metadata())
|
||||||
|
)
|
||||||
|
if not diffs:
|
||||||
|
logger.info("schema matches the models")
|
||||||
|
return None
|
||||||
|
logger.info("%s schema difference(s) detected", len(diffs))
|
||||||
|
before = head()
|
||||||
|
await _run(lambda c: command.revision(config(c), message=message, autogenerate=True))
|
||||||
|
after = head()
|
||||||
|
return after if after != before else None
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def _lock() -> AsyncIterator[None]:
|
||||||
|
"""Advisory lock, so only one worker migrates when several boot at once."""
|
||||||
|
async with get_engine().connect() as conn:
|
||||||
|
await conn.execution_options(isolation_level="AUTOCOMMIT")
|
||||||
|
await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": LOCK_ID})
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
await conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": LOCK_ID})
|
||||||
|
|
||||||
|
|
||||||
|
async def migrate(*, autogen: bool | None = None, message: str = "auto") -> None:
|
||||||
|
"""Apply pending revisions, then any fresh model drift, under the lock."""
|
||||||
|
should_autogen = get_settings().db_autogenerate if autogen is None else autogen
|
||||||
|
async with _lock():
|
||||||
|
await upgrade()
|
||||||
|
if should_autogen and await autogenerate(message):
|
||||||
|
await upgrade()
|
||||||
|
logger.info("database at revision %s", await current())
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv: Sequence[str] | None = None) -> None:
|
||||||
|
parser = argparse.ArgumentParser(prog="alembic_setup.py", description=__doc__)
|
||||||
|
parser.add_argument(
|
||||||
|
"command",
|
||||||
|
nargs="?",
|
||||||
|
default="migrate",
|
||||||
|
choices=["migrate", "revision", "upgrade", "downgrade", "current", "head"],
|
||||||
|
)
|
||||||
|
parser.add_argument("-m", "--message", default="auto", help="revision message")
|
||||||
|
parser.add_argument("-r", "--revision", help="target revision")
|
||||||
|
args = parser.parse_args(argv)
|
||||||
|
logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")
|
||||||
|
|
||||||
|
async def run() -> None:
|
||||||
|
from db_setup import init_db # local import: db_setup imports this module lazily
|
||||||
|
|
||||||
|
try:
|
||||||
|
if args.command == "migrate":
|
||||||
|
await init_db()
|
||||||
|
elif args.command == "revision":
|
||||||
|
print(await autogenerate(args.message) or "no changes")
|
||||||
|
elif args.command == "upgrade":
|
||||||
|
await upgrade(args.revision or "head")
|
||||||
|
elif args.command == "downgrade":
|
||||||
|
await downgrade(args.revision or "-1")
|
||||||
|
elif args.command == "current":
|
||||||
|
print(await current())
|
||||||
|
elif args.command == "head":
|
||||||
|
print(head())
|
||||||
|
finally:
|
||||||
|
await close_db()
|
||||||
|
|
||||||
|
asyncio.run(run())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
@ -0,0 +1,248 @@
|
||||||
|
"""PostgreSQL connection, async SQLAlchemy ORM and session management.
|
||||||
|
|
||||||
|
Configuration comes from the environment, with `.env` read from the repo root or
|
||||||
|
from `backend/` (`Db_USERNAME`, `Db_PASSWORD`, `Db_HOST`, `Db_PORT`, `Db_NAME`, and
|
||||||
|
the `DB_*` tuning fields below). Alembic lives in `alembic_setup.py`; `init_db()`
|
||||||
|
calls into it.
|
||||||
|
|
||||||
|
app = FastAPI(lifespan=lifespan) # migrate on startup
|
||||||
|
async def endpoint(db: AsyncSession = Depends(get_session)): ...
|
||||||
|
async with session_scope() as db: ... # workers and scripts
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import logging
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
from functools import lru_cache
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Annotated, Any, AsyncIterator, Sequence
|
||||||
|
|
||||||
|
from pydantic import field_validator
|
||||||
|
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
|
||||||
|
from sqlalchemy import MetaData, text
|
||||||
|
from sqlalchemy.engine import URL, make_url
|
||||||
|
from sqlalchemy.ext.asyncio import (
|
||||||
|
AsyncEngine,
|
||||||
|
AsyncSession,
|
||||||
|
async_sessionmaker,
|
||||||
|
create_async_engine,
|
||||||
|
)
|
||||||
|
from sqlalchemy.orm import DeclarativeBase
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
logger = logging.getLogger("db")
|
||||||
|
BASE_DIR = Path(__file__).resolve().parent
|
||||||
|
|
||||||
|
|
||||||
|
class Settings(BaseSettings):
|
||||||
|
"""Every field is overridden by an environment variable of the same name."""
|
||||||
|
|
||||||
|
model_config = SettingsConfigDict(
|
||||||
|
env_file=(BASE_DIR.parent / ".env", BASE_DIR / ".env"), extra="ignore"
|
||||||
|
)
|
||||||
|
|
||||||
|
database_url: str = "" # full DSN; wins over the Db_* parts below
|
||||||
|
db_username: str = os.getenv("DB_USERNAME")
|
||||||
|
db_password: str = os.getenv("DB_PASSWORD")
|
||||||
|
db_host: str = os.getenv("DB_HOST")
|
||||||
|
db_port: int = int(os.getenv("DB_PORT"))
|
||||||
|
db_name: str = os.getenv("DB_NAME")
|
||||||
|
db_sslmode: str = "" # e.g. "require" on Azure
|
||||||
|
|
||||||
|
|
||||||
|
db_schemas: Annotated[list[str], NoDecode] = ["app", "ref", "audit", "ai", "staging"]
|
||||||
|
db_default_schema: str = "app" # schema for models that declare none
|
||||||
|
db_echo: bool = False
|
||||||
|
db_pool_size: int = 5
|
||||||
|
db_max_overflow: int = 10
|
||||||
|
db_pool_recycle: int = 1800
|
||||||
|
db_connect_retries: int = 10
|
||||||
|
db_auto_migrate: bool = True # run `upgrade head` on startup
|
||||||
|
db_autogenerate: bool = True # write a revision when models drift from the schema
|
||||||
|
db_model_modules: Annotated[list[str], NoDecode] = [] # empty means auto-discover
|
||||||
|
app_name: str = "hr-ats-portal"
|
||||||
|
|
||||||
|
@field_validator("db_schemas", "db_model_modules", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def _csv(cls, value: Any) -> Any:
|
||||||
|
if isinstance(value, str):
|
||||||
|
return [item.strip() for item in value.split(",") if item.strip()]
|
||||||
|
return value
|
||||||
|
|
||||||
|
def url(self, *, async_driver: bool = True) -> URL:
|
||||||
|
"""DSN with the driver forced; `sslmode` is translated to asyncpg's `ssl`."""
|
||||||
|
url = (
|
||||||
|
make_url(self.database_url)
|
||||||
|
if self.database_url
|
||||||
|
else URL.create(
|
||||||
|
"postgresql",
|
||||||
|
self.db_username,
|
||||||
|
self.db_password,
|
||||||
|
self.db_host,
|
||||||
|
self.db_port,
|
||||||
|
self.db_name,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
query = dict(url.query)
|
||||||
|
if self.db_sslmode:
|
||||||
|
query.setdefault("sslmode", self.db_sslmode)
|
||||||
|
if async_driver and query.pop("sslmode", None) not in (None, "disable", "allow", "prefer"):
|
||||||
|
query["ssl"] = "true"
|
||||||
|
driver = "asyncpg" if async_driver else "psycopg2"
|
||||||
|
return url.set(drivername=f"postgresql+{driver}", query=query)
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=1)
|
||||||
|
def get_settings() -> Settings:
|
||||||
|
return Settings()
|
||||||
|
|
||||||
|
|
||||||
|
def database_url(*, async_driver: bool = True, hide_password: bool = False) -> str:
|
||||||
|
return get_settings().url(async_driver=async_driver).render_as_string(hide_password=hide_password)
|
||||||
|
|
||||||
|
|
||||||
|
NAMING_CONVENTION = {
|
||||||
|
"ix": "ix_%(table_name)s_%(column_0_N_name)s",
|
||||||
|
"uq": "uq_%(table_name)s_%(column_0_N_name)s",
|
||||||
|
"ck": "ck_%(table_name)s_%(constraint_name)s",
|
||||||
|
"fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s",
|
||||||
|
"pk": "pk_%(table_name)s",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class Base(DeclarativeBase):
|
||||||
|
"""Declarative base every model must inherit from."""
|
||||||
|
|
||||||
|
metadata = MetaData(
|
||||||
|
naming_convention=NAMING_CONVENTION,
|
||||||
|
schema=get_settings().db_default_schema or None,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
_engine: AsyncEngine | None = None
|
||||||
|
_sessionmaker: async_sessionmaker[AsyncSession] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_engine() -> AsyncEngine:
|
||||||
|
"""The process-wide AsyncEngine, created on first use."""
|
||||||
|
global _engine
|
||||||
|
if _engine is None:
|
||||||
|
s = get_settings()
|
||||||
|
_engine = create_async_engine(
|
||||||
|
s.url(),
|
||||||
|
echo=s.db_echo,
|
||||||
|
pool_pre_ping=True,
|
||||||
|
pool_size=s.db_pool_size,
|
||||||
|
max_overflow=s.db_max_overflow,
|
||||||
|
pool_recycle=s.db_pool_recycle,
|
||||||
|
connect_args={
|
||||||
|
"server_settings": {"timezone": "UTC", "application_name": s.app_name}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
return _engine
|
||||||
|
|
||||||
|
|
||||||
|
def get_sessionmaker() -> async_sessionmaker[AsyncSession]:
|
||||||
|
global _sessionmaker
|
||||||
|
if _sessionmaker is None:
|
||||||
|
_sessionmaker = async_sessionmaker(
|
||||||
|
bind=get_engine(), class_=AsyncSession, expire_on_commit=False, autoflush=False
|
||||||
|
)
|
||||||
|
return _sessionmaker
|
||||||
|
|
||||||
|
|
||||||
|
async def get_session() -> AsyncIterator[AsyncSession]:
|
||||||
|
"""FastAPI dependency. Rolls back on error; committing is the caller's job."""
|
||||||
|
async with get_sessionmaker()() as session:
|
||||||
|
try:
|
||||||
|
yield session
|
||||||
|
except Exception:
|
||||||
|
await session.rollback()
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def session_scope() -> AsyncIterator[AsyncSession]:
|
||||||
|
"""Transactional session for scripts and workers: commits, or rolls back on error."""
|
||||||
|
async with get_sessionmaker()() as session:
|
||||||
|
try:
|
||||||
|
yield session
|
||||||
|
await session.commit()
|
||||||
|
except Exception:
|
||||||
|
await session.rollback()
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
async def close_db() -> None:
|
||||||
|
"""Dispose of the pool and reset the cached factories."""
|
||||||
|
global _engine, _sessionmaker
|
||||||
|
if _engine is not None:
|
||||||
|
await _engine.dispose()
|
||||||
|
logger.info("connection pool closed")
|
||||||
|
_engine, _sessionmaker = None, None
|
||||||
|
|
||||||
|
|
||||||
|
async def check_connection(retries: int | None = None, delay: float = 1.0) -> None:
|
||||||
|
"""Wait for Postgres to answer `SELECT 1`, retrying with a capped backoff."""
|
||||||
|
attempts = get_settings().db_connect_retries if retries is None else retries
|
||||||
|
for attempt in range(1, max(attempts, 1) + 1):
|
||||||
|
try:
|
||||||
|
async with get_engine().connect() as conn:
|
||||||
|
await conn.execute(text("SELECT 1"))
|
||||||
|
logger.info("connected to %s", database_url(hide_password=True))
|
||||||
|
return
|
||||||
|
except Exception as exc:
|
||||||
|
if attempt >= attempts:
|
||||||
|
raise RuntimeError(f"cannot reach {database_url(hide_password=True)}") from exc
|
||||||
|
logger.warning("database not ready (%s/%s): %s", attempt, attempts, exc)
|
||||||
|
await asyncio.sleep(delay)
|
||||||
|
delay = min(delay * 2, 10.0)
|
||||||
|
|
||||||
|
|
||||||
|
async def create_schemas(schemas: Sequence[str] | None = None) -> None:
|
||||||
|
"""`CREATE SCHEMA IF NOT EXISTS` for every configured schema."""
|
||||||
|
names = list(get_settings().db_schemas if schemas is None else schemas)
|
||||||
|
if not names:
|
||||||
|
return
|
||||||
|
async with get_engine().begin() as conn:
|
||||||
|
for name in names:
|
||||||
|
await conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{name}"'))
|
||||||
|
logger.info("schemas ensured: %s", ", ".join(names))
|
||||||
|
|
||||||
|
|
||||||
|
async def init_db(*, migrate: bool | None = None, autogen: bool | None = None) -> None:
|
||||||
|
"""Connect, create the schemas, then bring migrations up to head."""
|
||||||
|
should_migrate = get_settings().db_auto_migrate if migrate is None else migrate
|
||||||
|
await check_connection()
|
||||||
|
await create_schemas()
|
||||||
|
if should_migrate:
|
||||||
|
from alembic_setup import migrate as run_migrations # local import: avoids a cycle
|
||||||
|
|
||||||
|
await run_migrations(autogen=autogen)
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def lifespan(app: Any = None) -> AsyncIterator[None]:
|
||||||
|
"""FastAPI lifespan: `app = FastAPI(lifespan=lifespan)`."""
|
||||||
|
await init_db()
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
await close_db()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")
|
||||||
|
|
||||||
|
async def _main() -> None:
|
||||||
|
try:
|
||||||
|
await init_db()
|
||||||
|
finally:
|
||||||
|
await close_db()
|
||||||
|
|
||||||
|
asyncio.run(_main())
|
||||||
|
|
@ -0,0 +1,8 @@
|
||||||
|
from fastapi import APIRouter
|
||||||
|
import httpx
|
||||||
|
from db_setup import get_session
|
||||||
|
router = APIRouter()
|
||||||
|
|
||||||
|
@router.get("/inbox")
|
||||||
|
async def get_inbox():
|
||||||
|
return {"message": "Hello, World!"}
|
||||||
|
|
@ -0,0 +1,24 @@
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import fastapi
|
||||||
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from fastapi import FastAPI,APIRouter
|
||||||
|
from db_setup import lifespan
|
||||||
|
from inbox.app import router as inbox_router
|
||||||
|
|
||||||
|
# Without this the db/migration logs have no handler and are swallowed under uvicorn.
|
||||||
|
logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")
|
||||||
|
|
||||||
|
# lifespan connects to Postgres and brings migrations up to head on startup,
|
||||||
|
# and disposes of the connection pool on shutdown.
|
||||||
|
app = FastAPI(lifespan=lifespan)
|
||||||
|
app.add_middleware(
|
||||||
|
CORSMiddleware,
|
||||||
|
allow_origins=["*"],
|
||||||
|
allow_credentials=True,
|
||||||
|
allow_methods=["*"],
|
||||||
|
allow_headers=["*"],
|
||||||
|
)
|
||||||
|
|
||||||
|
app.include_router(inbox_router)
|
||||||
|
|
@ -0,0 +1,39 @@
|
||||||
|
"""Alembic environment -- generated by alembic_setup.py."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from alembic import context
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
import alembic_setup as setup # noqa: E402
|
||||||
|
import db_setup # noqa: E402
|
||||||
|
|
||||||
|
metadata = setup.target_metadata()
|
||||||
|
options = setup.context_options()
|
||||||
|
|
||||||
|
|
||||||
|
def run(connection) -> None:
|
||||||
|
context.configure(connection=connection, target_metadata=metadata, **options)
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
|
||||||
|
|
||||||
|
if context.is_offline_mode():
|
||||||
|
context.configure(
|
||||||
|
url=db_setup.database_url(async_driver=False),
|
||||||
|
target_metadata=metadata,
|
||||||
|
literal_binds=True,
|
||||||
|
dialect_opts={"paramstyle": "named"},
|
||||||
|
**options,
|
||||||
|
)
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
elif (connection := context.config.attributes.get("connection")) is not None:
|
||||||
|
run(connection) # alembic_setup passed an already-open connection
|
||||||
|
else:
|
||||||
|
asyncio.run(setup.run_standalone(run)) # bare `alembic` CLI
|
||||||
|
|
@ -0,0 +1,27 @@
|
||||||
|
"""${message}
|
||||||
|
|
||||||
|
Revision ID: ${up_revision}
|
||||||
|
Revises: ${down_revision | comma,n}
|
||||||
|
Create Date: ${create_date}
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
${imports if imports else ""}
|
||||||
|
|
||||||
|
revision: str = ${repr(up_revision)}
|
||||||
|
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||||
|
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
${upgrades if upgrades else "pass"}
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
${downgrades if downgrades else "pass"}
|
||||||
Loading…
Reference in New Issue