core-api/alembic/env.py

"""
Alembic Environment Configuration

Async migration environment for SQLAlchemy 2.0 with asyncpg.

Alembic refuses to run without an explicit target (core-api T-3): a bare
`alembic current` used to fall back to `get_settings().database_url`, which
is composed from `.env`'s production `POSTGRES_*` values — the most natural
command in the repo was quietly talking to production. `-x target=test`
names the database `TEST_DATABASE_URL` points at; `-x target=production`
names today's settings-derived database, deliberately. Nothing else
defaults, and the check runs — and can exit — before any engine is built.
"""
import asyncio
import os
import re
import sys
from logging.config import fileConfig
from pathlib import Path

from sqlalchemy import pool
from sqlalchemy.engine import Connection, make_url
from sqlalchemy.ext.asyncio import async_engine_from_config

from alembic import context

# Import our models (config is imported lazily, inside _production_url, so
# that resolving a bad or missing target never has to construct Settings).
from src.shared.database import Base

# Import all models to ensure they're registered with Base.metadata
from src.domains.auth.models import User, Role, UserRole, UserPreferences, ApiKey  # noqa: F401
from src.domains.dashboard.models import QuickLink, DashboardWidget  # noqa: F401
from src.domains.book.models import (  # noqa: F401
    BookRecord, BookTag, BookRecordLink, BookAttachment, BookEvent, BookTenantCounter,
)
from src.domains.scribe.models import (  # noqa: F401
    ScribeAddress, ScribeMessage, ScribeMailboxEntry,
)
from src.domains.tenants.models import Tenant, TenantConsumer  # noqa: F401

REPO_ROOT = Path(__file__).resolve().parent.parent

# Alembic Config object
config = context.config

# Interpret the config file for Python logging
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# Model metadata for autogenerate support
target_metadata = Base.metadata


def _test_url() -> str:
    """TEST_DATABASE_URL from the environment, else the repo's own `.env` —
    read that one variable only; the file is never loaded or sourced whole."""
    value = os.environ.get("TEST_DATABASE_URL")
    if value:
        return value

    env_path = REPO_ROOT / ".env"
    if env_path.is_file():
        with open(env_path) as f:
            for line in f:
                match = re.match(r"^TEST_DATABASE_URL=(.*)$", line.rstrip("\n"))
                if match:
                    return match.group(1).strip().strip("'\"")

    raise SystemExit(
        "alembic: -x target=test needs TEST_DATABASE_URL (environment, or "
        "the repo's .env) — none found."
    )


def _production_url() -> str:
    """Today's settings-derived database — deliberately, only when asked
    for by name."""
    from src.shared.config import get_settings

    return get_settings().database_url


def _resolve_target_url() -> str:
    """-x target=test|production names the database to migrate. No target,
    or anything else, exits non-zero here — before any engine is built."""
    x_args = context.get_x_argument(as_dictionary=True)
    target = x_args.get("target")

    if target == "test":
        url = _test_url()
    elif target == "production":
        url = _production_url()
    else:
        raise SystemExit(
            "alembic: no migration target named. Pass one of:\n"
            "  -x target=test        (the database TEST_DATABASE_URL names)\n"
            "  -x target=production  (today's settings-derived database)\n"
            f"Got -x target={target!r}. Nothing defaults."
        )

    if url.startswith("postgresql://"):
        url = url.replace("postgresql://", "postgresql+asyncpg://", 1)
    return url


db_url = _resolve_target_url()

# Name the target before migrating anything — never a host's credentials,
# never the URL itself.
print(f"alembic target: database={make_url(db_url).database!r}", file=sys.stderr)


def run_migrations_offline() -> None:
    """
    Run migrations in 'offline' mode.

    Generates SQL script without connecting to the database.
    """
    context.configure(
        url=db_url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()


def do_run_migrations(connection: Connection) -> None:
    """
    Run migrations with the given connection.
    """
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
    )

    with context.begin_transaction():
        context.run_migrations()


async def run_async_migrations() -> None:
    """
    Run migrations in 'online' mode with async engine.
    """
    configuration = config.get_section(config.config_ini_section) or {}
    configuration["sqlalchemy.url"] = db_url

    connectable = async_engine_from_config(
        configuration,
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)

    await connectable.dispose()


def run_migrations_online() -> None:
    """
    Run migrations in 'online' mode.
    """
    asyncio.run(run_async_migrations())


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()