src/common/migrate.py - master - azathd/mutiagent

src/common/migrate.py
266 строк8 KB
Your Name
init
15 май 2026, 19:11
15 май 2026, 19:11f0c1163
"""Database schema: Alembic (sync psycopg) or inline SQL via asyncpg."""

from __future__ import annotations

import asyncio
from pathlib import Path

import asyncpg
from alembic import command
from alembic.config import Config
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

from src.common.db_url import database_url_async, database_url_sync
from src.common.logging import get_logger

log = get_logger("migrate")

# Embedded DDL — works without deploy/postgres_schema.sql on disk (dev volume mounts).
SCHEMA_STATEMENTS: tuple[str, ...] = (
    """
    CREATE TABLE IF NOT EXISTS users (
        id UUID PRIMARY KEY,
        username VARCHAR(255) NOT NULL UNIQUE,
        email VARCHAR(255),
        marzban_uuid VARCHAR(64),
        created_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS user_profiles (
        id UUID PRIMARY KEY,
        user_id UUID REFERENCES users (id),
        baseline_bytes_hour BIGINT DEFAULT 0,
        last_seen_at TIMESTAMPTZ,
        meta JSONB DEFAULT '{}'
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS user_limits (
        id UUID PRIMARY KEY,
        user_id UUID UNIQUE REFERENCES users (id),
        daily_bytes BIGINT,
        monthly_bytes BIGINT,
        allowed_countries JSONB DEFAULT '[]'
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS incidents (
        id UUID PRIMARY KEY,
        title VARCHAR(512) NOT NULL,
        status VARCHAR(64) DEFAULT 'new',
        severity VARCHAR(32) NOT NULL,
        summary TEXT,
        created_at TIMESTAMPTZ DEFAULT now(),
        updated_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS actions_queue (
        id UUID PRIMARY KEY,
        incident_id UUID REFERENCES incidents (id),
        action_type VARCHAR(64) NOT NULL,
        target VARCHAR(255) NOT NULL,
        payload JSONB DEFAULT '{}',
        status VARCHAR(32) DEFAULT 'proposed',
        severity VARCHAR(32) DEFAULT 'low',
        created_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS hitl_requests (
        id UUID PRIMARY KEY,
        action_id UUID REFERENCES actions_queue (id),
        status VARCHAR(32) DEFAULT 'pending',
        requested_by VARCHAR(128) DEFAULT 'orchestrator',
        resolved_by VARCHAR(128),
        created_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS audit_log (
        id UUID PRIMARY KEY,
        action_id VARCHAR(64) NOT NULL UNIQUE,
        action_type VARCHAR(64) NOT NULL,
        actor VARCHAR(128) NOT NULL,
        target VARCHAR(255) NOT NULL,
        payload JSONB DEFAULT '{}',
        result JSONB DEFAULT '{}',
        prev_hash VARCHAR(64) DEFAULT '',
        entry_hash VARCHAR(64) NOT NULL,
        created_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS reports (
        id UUID PRIMARY KEY,
        period VARCHAR(32) NOT NULL,
        object_key VARCHAR(512) NOT NULL,
        formats JSONB DEFAULT '[]',
        created_at TIMESTAMPTZ DEFAULT now()
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS alembic_version (
        version_num VARCHAR(32) NOT NULL,
        CONSTRAINT alembic_version_pkc PRIMARY KEY (version_num)
    )
    """,
    """
    INSERT INTO alembic_version (version_num) VALUES ('20240516_full')
    ON CONFLICT (version_num) DO NOTHING
    """,
)

AUDIT_TRIGGER_STATEMENTS: tuple[str, ...] = (
    """
    CREATE OR REPLACE FUNCTION audit_log_deny_mutation() RETURNS trigger AS $$
    BEGIN
        RAISE EXCEPTION 'audit_log is append-only';
    END;
    $$ LANGUAGE plpgsql
    """,
    "DROP TRIGGER IF EXISTS audit_log_no_update ON audit_log",
    """
    CREATE TRIGGER audit_log_no_update
        BEFORE UPDATE OR DELETE ON audit_log
        FOR EACH ROW EXECUTE PROCEDURE audit_log_deny_mutation()
    """,
)


def _asyncpg_dsn() -> str:
    return database_url_async().replace("postgresql+asyncpg://", "postgresql://")


def _project_root() -> Path:
    root = Path(__file__).resolve().parents[2]
    return root if (root / "deploy" / "postgres_schema.sql").exists() else Path("/app")


def split_sql_script(sql: str) -> list[str]:
    """Split SQL on semicolons outside $$ ... $$ blocks."""
    statements: list[str] = []
    buf: list[str] = []
    in_dollar = False
    for line in sql.splitlines(keepends=True):
        if "$$" in line:
            in_dollar = not in_dollar
        buf.append(line)
        if not in_dollar and line.rstrip().endswith(";"):
            chunk = "".join(buf).strip()
            buf.clear()
            if chunk and not chunk.startswith("--"):
                statements.append(chunk)
    tail = "".join(buf).strip()
    if tail and not tail.startswith("--"):
        statements.append(tail)
    return statements


async def _run_asyncpg_statements(statements: tuple[str, ...] | list[str], *, label: str) -> None:
    conn = await asyncpg.connect(_asyncpg_dsn())
    try:
        for stmt in statements:
            try:
                await conn.execute(stmt)
            except Exception as exc:
                log.error("schema_stmt_failed", label=label, error=str(exc), stmt=stmt[:120])
                raise
    finally:
        await conn.close()


def run_alembic_upgrade() -> None:
    """Apply all pending migrations to head (requires psycopg in the image)."""
    root = _project_root()
    ini = root / "alembic.ini"
    if not ini.exists():
        ini = Path("/app/alembic.ini")

    cfg = Config(str(ini))
    cfg.set_main_option("script_location", str(root / "migrations"))
    cfg.set_main_option("sqlalchemy.url", database_url_sync())

    log.info("alembic_upgrade_start")
    command.upgrade(cfg, "head")
    log.info("alembic_upgrade_done")


async def incidents_table_exists() -> bool:
    engine = create_async_engine(database_url_async())
    try:
        async with engine.connect() as conn:
            reg = await conn.scalar(text("SELECT to_regclass('public.incidents')"))
            return reg is not None
    finally:
        await engine.dispose()


async def apply_schema_inline_async() -> None:
    """Create core tables via asyncpg (one statement = one commit)."""
    if await incidents_table_exists():
        log.info("schema_already_present")
        return

    log.info("schema_inline_apply_start", statements=len(SCHEMA_STATEMENTS))
    await _run_asyncpg_statements(SCHEMA_STATEMENTS, label="core")
    try:
        await _run_asyncpg_statements(AUDIT_TRIGGER_STATEMENTS, label="audit_trigger")
    except Exception as exc:
        log.warning("audit_trigger_skipped", error=str(exc))

    if not await incidents_table_exists():
        raise RuntimeError("incidents table still missing after inline schema apply")
    log.info("schema_inline_apply_done")


async def apply_schema_sql_file_async() -> None:
    path = _project_root() / "deploy" / "postgres_schema.sql"
    if not path.is_file():
        raise FileNotFoundError(f"Schema SQL not found: {path}")
    statements = split_sql_script(path.read_text(encoding="utf-8"))
    log.info("schema_file_apply_start", path=str(path), statements=len(statements))
    await _run_asyncpg_statements(statements, label="file")
    log.info("schema_file_apply_done")


async def ensure_database_schema() -> None:
    """Ensure Postgres schema exists (Alembic preferred, inline SQL fallback)."""
    if await incidents_table_exists():
        return

    try:
        await asyncio.to_thread(run_alembic_upgrade)
        if await incidents_table_exists():
            return
    except Exception as exc:
        log.warning("alembic_upgrade_failed", error=str(exc))

    await apply_schema_inline_async()

    if await incidents_table_exists():
        return

    try:
        await apply_schema_sql_file_async()
    except FileNotFoundError:
        pass

    if not await incidents_table_exists():
        raise RuntimeError(
            "Postgres schema not applied: incidents table missing. "
            "Run: make init-db  OR  docker compose run --rm --entrypoint '' orchestrator python -m src.common.migrate"
        )


async def _main() -> None:
    configure = __import__("src.common.logging", fromlist=["configure_logging"]).configure_logging
    configure()
    await ensure_database_schema()
    log.info("schema_ready")


if __name__ == "__main__":
    asyncio.run(_main())