FastAPI Alembic Migrations

main.py

from multiprocessing import Process, Queue
from pathlib import Path
from typing import Literal

import alembic.command
from alembic.config import Config
from fastapi import FastAPI
from pydantic import BaseModel


app = FastAPI()


alembic_config_path = (Path(__file__).parent / '../../alembic.ini').resolve()


class RunCommandRequestBody(BaseModel):
    cmd: str


class AlembicMigrationRequestBody(BaseModel):
    revision: str = 'head'


class AlembicMigrationResponse(BaseModel):
    success: bool


AlembicCommand = Literal['upgrade', 'downgrade']


def alembic_migrate(queue: Queue, cmd: AlembicCommand, revision: str):
    # TODO: Make sure errors get reported to Sentry
    try:
        cfg = Config(alembic_config_path)
        getattr(alembic.command, cmd)(cfg, revision)
    except Exception as e:
        queue.put(False)
        raise e
    queue.put(True)


def alembic_migrate_process(cmd: AlembicCommand, revision: str) -> bool:
    """
    Must run in a subprocess, because otherwise Alembic will mess up FastAPI's logging.
    See: https://stackoverflow.com/questions/24622170/using-alembic-api-from-inside-application-code
    """
    queue = Queue()
    process = Process(target=alembic_migrate, args=(queue, cmd, revision))
    process.start()
    process.join(timeout=10 * 60)  # 10 min timeout
    return queue.get()


@app.post('/alembic/upgrade/', response_model=AlembicMigrationResponse)
async def alembic_upgrade(body: AlembicMigrationRequestBody):
    success = alembic_migrate_process('upgrade', body.revision)
    return AlembicMigrationResponse(success=success)


@app.post('/alembic/downgrade/', response_model=AlembicMigrationResponse)
async def alembic_downgrade(body: AlembicMigrationRequestBody):
    success = alembic_migrate_process('downgrade', body.revision)
    return AlembicMigrationResponse(success=success)


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app)