Spaces:
Sleeping
Sleeping
File size: 2,666 Bytes
e86dfae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
Migration runner used at container start.
Three situations have to be handled, because this project ran
Base.metadata.create_all() at startup for a long time and never ran Alembic:
1. Fresh database, nothing in it
-> upgrade to head normally.
2. Existing database created by create_all(), no alembic_version table
-> the tables already exist, so running the baseline revision would
fail with "relation already exists". Adopt the schema by stamping
it at the baseline, then apply anything newer.
3. Database already under Alembic control
-> upgrade to head normally.
Run directly: python -m scripts.migrate
"""
import logging
import sys
from sqlalchemy import inspect
from alembic import command
from alembic.config import Config
from app.db.session import engine, ensure_indexes
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - migrate - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# The revision whose schema matches what create_all() produced.
BASELINE_REVISION = "001"
def _alembic_config() -> Config:
cfg = Config("alembic.ini")
cfg.set_main_option("sqlalchemy.url", str(engine.url))
return cfg
def run() -> None:
inspector = inspect(engine)
tables = set(inspector.get_table_names())
cfg = _alembic_config()
if "alembic_version" in tables:
logger.info("Database is under Alembic control — upgrading to head")
command.upgrade(cfg, "head")
elif "documents" in tables:
# Case 2: pre-Alembic schema built by create_all().
logger.warning(
"Found an existing schema with no alembic_version table. "
"Adopting it at revision %s, then applying newer revisions.",
BASELINE_REVISION,
)
command.stamp(cfg, BASELINE_REVISION)
command.upgrade(cfg, "head")
# Stamping marks the baseline as applied without running it, so the
# indexes that revision creates are skipped. create_all() never made
# them either — which is precisely why production has been doing
# sequential scans instead of using idx_embeddings_hnsw. Create them
# explicitly for adopted databases.
logger.info("Ensuring vector/trigram indexes on the adopted schema")
ensure_indexes()
else:
logger.info("Fresh database — creating schema from the baseline")
command.upgrade(cfg, "head")
logger.info("Migrations complete")
if __name__ == "__main__":
try:
run()
except Exception as e:
logger.error(f"Migration failed: {e}", exc_info=True)
sys.exit(1)
|