1  """Alembic migration runner.
 2  
 3  Stdlib-only on purpose: this file runs as the very first step of
 4  `make update`, before `uv sync` has had a chance to install anything
 5  new from `pyproject.toml`. Adding a third-party CLI dep here means an
 6  old instance whose .venv is missing that dep can't even reach the
 7  migrations to fix the upgrade. argparse keeps this script bootable
 8  under any Python 3 install.
 9  """
10  import argparse
11  
12  from alembic.config import Config
13  from alembic import command
14  
15  from restai.config import POSTGRES_URL, MYSQL_URL, POSTGRES_HOST, MYSQL_HOST
16  
17  
18  def get_database_url():
19      if POSTGRES_HOST:
20          return POSTGRES_URL
21      elif MYSQL_HOST:
22          return MYSQL_URL
23      else:
24          return "sqlite:///./restai.db"
25  
26  
27  def _alembic_cfg():
28      cfg = Config("alembic.ini")
29      cfg.set_main_option("sqlalchemy.url", get_database_url())
30      return cfg
31  
32  
33  def upgrade():
34      """Run database migrations to upgrade the database schema."""
35      command.upgrade(_alembic_cfg(), "head")
36  
37  
38  def downgrade():
39      """Run database migrations to downgrade the database schema."""
40      command.downgrade(_alembic_cfg(), "-1")
41  
42  
43  def main():
44      parser = argparse.ArgumentParser(
45          description="Run Alembic migrations against the configured database.",
46      )
47      sub = parser.add_subparsers(dest="cmd", required=True)
48      sub.add_parser("upgrade", help="Upgrade to the latest revision (alembic head).")
49      sub.add_parser("downgrade", help="Downgrade by one revision (alembic -1).")
50      args = parser.parse_args()
51  
52      if args.cmd == "upgrade":
53          upgrade()
54      elif args.cmd == "downgrade":
55          downgrade()
56  
57  
58  if __name__ == "__main__":
59      main()