Api-Gateway/app/main.py

import asyncio
from contextlib import asynccontextmanager, suppress

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import RedirectResponse
from sqlalchemy import text

from app import config, portal, proxy, retention, security
from app.admin import routes as admin_routes
from app.admin import stats as admin_stats
from app.admin.deps import LoginRequired, login_redirect_handler
from app.database import Base, SessionLocal, engine
from app.models import User


def enable_incremental_vacuum() -> None:
    """Deleted log rows should hand their pages back to the filesystem;
    switching auto_vacuum requires a one-time VACUUM (cannot run in a
    transaction, hence the raw connection)."""
    raw = engine.raw_connection()
    try:
        cursor = raw.cursor()
        mode = cursor.execute("PRAGMA auto_vacuum").fetchone()[0]
        if mode != 2:  # 2 = INCREMENTAL
            cursor.execute("PRAGMA auto_vacuum=INCREMENTAL")
            raw.commit()
            cursor.execute("VACUUM")
            raw.commit()
    finally:
        raw.close()


def _table_exists(db, name: str) -> bool:
    return db.execute(text(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=:n"
    ), {"n": name}).first() is not None


def migrate_schema() -> None:
    """Bring older databases up to date (SQLite)."""
    with engine.begin() as conn:
        cols = [row[1] for row in conn.exec_driver_sql("PRAGMA table_info(request_logs)")]
        if cols and "endpoint_id" not in cols:
            conn.exec_driver_sql(
                "ALTER TABLE request_logs ADD COLUMN endpoint_id INTEGER "
                "REFERENCES endpoints(id) ON DELETE SET NULL"
            )
        for column, ddl in (
            ("query_string", "VARCHAR(2048) DEFAULT ''"),
            ("request_body", "TEXT DEFAULT ''"),
            ("response_body", "TEXT DEFAULT ''"),
        ):
            if cols and column not in cols:
                conn.exec_driver_sql(f"ALTER TABLE request_logs ADD COLUMN {column} {ddl}")
        service_cols = [row[1] for row in conn.exec_driver_sql("PRAGMA table_info(services)")]
        if service_cols and "verify_tls" not in service_cols:
            conn.exec_driver_sql("ALTER TABLE services ADD COLUMN verify_tls BOOLEAN DEFAULT 1")


def seed_and_migrate() -> None:
    db = SessionLocal()
    try:
        # Role-era databases: replace role->service grants with per-key grants
        # on a service-wide catch-all endpoint, then drop the role tables.
        if _table_exists(db, "roles"):
            db.execute(text(
                "INSERT INTO endpoints (service_id, method, path, description) "
                "SELECT id, '*', '/**', 'Migrated: full service access' FROM services"
            ))
            if _table_exists(db, "role_service_access"):
                db.execute(text(
                    "INSERT OR IGNORE INTO key_endpoint_access (api_key_id, endpoint_id) "
                    "SELECT k.id, e.id "
                    "FROM api_keys k "
                    "JOIN users u ON u.id = k.user_id "
                    "JOIN role_service_access rsa ON rsa.role_id = u.role_id "
                    "JOIN endpoints e ON e.service_id = rsa.service_id AND e.path = '/**'"
                ))
                db.execute(text("DROP TABLE role_service_access"))
            if _table_exists(db, "role_permissions"):
                db.execute(text("DROP TABLE role_permissions"))
            db.execute(text("ALTER TABLE users DROP COLUMN role_id"))
            db.execute(text("DROP TABLE roles"))
            db.commit()
            print("Migrated role-based access to per-key endpoint grants.")

        if db.query(User).count() == 0:
            db.add(User(
                username=config.DEFAULT_ADMIN_USERNAME,
                password_hash=security.hash_password(config.DEFAULT_ADMIN_PASSWORD),
            ))
            db.commit()
            print(f"Seeded admin user '{config.DEFAULT_ADMIN_USERNAME}' "
                  f"(password: '{config.DEFAULT_ADMIN_PASSWORD}' — change it after first login).")
    finally:
        db.close()


@asynccontextmanager
async def lifespan(app: FastAPI):
    Base.metadata.create_all(engine)
    migrate_schema()
    seed_and_migrate()
    enable_incremental_vacuum()
    retention_task = asyncio.create_task(retention.retention_loop())
    yield
    retention_task.cancel()
    with suppress(asyncio.CancelledError):
        await retention_task
    await proxy.close_client()


# Built-in docs/openapi are disabled: the gateway serves its own consumer-facing
# /docs and per-key /openapi.json (app/portal.py) at those paths instead.
app = FastAPI(title="API Gateway", version="2.0.0", lifespan=lifespan,
              docs_url=None, redoc_url=None, openapi_url=None)
app.add_exception_handler(LoginRequired, login_redirect_handler)


@app.get("/", include_in_schema=False)
def root():
    return RedirectResponse("/admin")


@app.get("/health", include_in_schema=False)
def health():
    return {"status": "ok"}


app.include_router(admin_routes.router)
app.include_router(admin_stats.router)
app.include_router(portal.router)
app.mount("/static", StaticFiles(directory=str(config.BASE_DIR / "app" / "static")), name="static")
# The proxy catch-all (/{slug}/...) must come last so it never shadows
# /admin, /static, /docs or /health.
app.include_router(proxy.router)