custom alembic cli for upgrade/downgrade

Author: alex75

Diff for: alembic.ini

@@ -2,7 +2,7 @@
 
 [alembic]
 # path to migration scripts
-script_location = alembic
+script_location = %(here)s/alembic
 
 # template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
 # Uncomment the line below if you want the files to be prepended with date and time
@@ -60,12 +60,6 @@ version_path_separator = os  # Use os.pathsep. Default configuration used for ne
 # are written from script.py.mako
 # output_encoding = utf-8
 
-drivername =
-username =
-password =
-host =
-port =
-database =
 sqlalchemy.url = %(drivername)s://%(username)s:%(password)s@%(host)s:%(port)s/%(database)s
 
 [post_write_hooks]

Diff for: alembic/env.py

@@ -19,8 +19,6 @@
 import alembic.context
 import cads_broker
 
-config = alembic.context.config
-
 
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.
@@ -33,11 +31,7 @@ def run_migrations_offline() -> None:
     Calls to alembic.context.execute() here emit the given string to the
     script output.
     """
-    url_props = dict()
-    for prop in ["drivername", "username", "password", "host", "port", "database"]:
-        url_props[prop] = config.get_main_option(prop)
-    url_props["port"] = url_props["port"] and int(url_props["port"]) or None  # type: ignore
-    url = sa.engine.URL.create(**url_props)  # type: ignore
+    url = alembic.context.config.get_main_option("sqlalchemy.url")
     alembic.context.configure(
         url=url,
         target_metadata=cads_broker.database.BaseModel.metadata,
@@ -54,12 +48,8 @@ def run_migrations_online() -> None:
     In this scenario we need to create an Engine
     and associate a connection with the alembic.context.
     """
-    url_props = dict()
-    for prop in ["drivername", "username", "password", "host", "port", "database"]:
-        url_props[prop] = config.get_main_option(prop)
-    url_props["port"] = url_props["port"] and int(url_props["port"]) or None  # type: ignore
-    url = sa.engine.URL.create(**url_props)  # type: ignore
-    engine = sa.create_engine(url, poolclass=sa.pool.NullPool)
+    url = alembic.context.config.get_main_option("sqlalchemy.url")
+    engine = sa.create_engine(url, poolclass=sa.pool.NullPool)  # type: ignore
     with engine.connect() as connection:
         alembic.context.configure(
             connection=connection,

Diff for: cads_broker/alembic_cli.py

@@ -0,0 +1,34 @@
+"""Custom alembic CLI with new default config path + db url by environment."""
+
+import os
+
+import cads_broker.config
+from alembic.config import CommandLine, Config
+
+alembic_ini_path = os.path.abspath(os.path.join(__file__, "..", "..", "alembic.ini"))
+
+
+class MyCommandLine(CommandLine):
+    def main(self, argv=None):
+        options = self.parser.parse_args(argv)
+        if not hasattr(options, "cmd"):
+            self.parser.error("too few arguments")
+        else:
+            cfg = Config(
+                file_=options.config,
+                ini_section=options.name,
+                cmd_opts=options,
+            )
+            url = cads_broker.config.ensure_settings().connection_string.replace(
+                "%", "%%"
+            )
+            cfg.set_main_option("sqlalchemy.url", url)
+            self.run_cmd(cfg, options)
+
+
+def main():
+    cli = MyCommandLine(prog="alembic-cli")
+    config_in_parser = [p for p in cli.parser._actions if p.dest == "config"][0]
+    config_in_parser.default = alembic_ini_path
+    config_in_parser.help = f'Alternate config file; defaults to "{alembic_ini_path}"'
+    cli.main()

Diff for: cads_broker/database.py

@@ -19,7 +19,7 @@
 
 import alembic.command
 import alembic.config
-from cads_broker import config
+from cads_broker import alembic_cli, config
 
 BaseModel = sa.orm.declarative_base()
 
@@ -932,15 +932,9 @@ def init_database(connection_string: str, force: bool = False) -> sa.engine.Engi
     :param force: if True, drop the database structure and build again from scratch
     """
     engine = sa.create_engine(connection_string)
-    migration_directory = os.path.abspath(os.path.join(__file__, "..", ".."))
-    os.chdir(migration_directory)
-    alembic_config_path = os.path.join(migration_directory, "alembic.ini")
-    alembic_cfg = alembic.config.Config(alembic_config_path)
-    for option in ["drivername", "username", "password", "host", "port", "database"]:
-        value = getattr(engine.url, option)
-        if value is None:
-            value = ""
-        alembic_cfg.set_main_option(option, str(value))
+    alembic_cfg = alembic.config.Config(alembic_cli.alembic_ini_path)
+    # passwords with special chars are urlencoded, but '%' must be escaped in ini files
+    alembic_cfg.set_main_option("sqlalchemy.url", connection_string.replace("%", "%%"))
     if not sqlalchemy_utils.database_exists(engine.url):
         sqlalchemy_utils.create_database(engine.url)
         # cleanup and create the schema

Diff for: pyproject.toml

@@ -31,6 +31,7 @@ readme = "README.md"
 
 [project.scripts]
 broker = "cads_broker.entry_points:main"
+broker-alembic-cli = "cads_broker.alembic_cli:main"
 
 [tool.coverage.run]
 branch = true