Merge pull request #155 from ecmwf-projects/COPDS-2511-alembic

custom alembic cli for upgrade/downgrade

Author: alex75

README.md

@@ -28,7 +28,7 @@ Before pushing to GitHub, run the following commands:
 1. Run the static type checker: `make type-check`
 1. Build the documentation (see [Sphinx tutorial](https://www.sphinx-doc.org/en/master/tutorial/)): `make docs-build`
 
-### Instructions for database updating
+### Instructions for creating a new database version
 
 The package `cads-broker` comes with its 'broker' database.
 In case of database structure upgrade, developers must follow these steps:
@@ -46,6 +46,34 @@ In case of database structure upgrade, developers must follow these steps:
    Similarly, do the same with the `downgrade` function.
 1. Commit and push the modifications and the new file.
 
+### Instructions for moving between different database versions
+
+The package comes with its own 'broker-alembic-cli' script in order to move between different
+database versions. This script is a slight modified version of the 'alembic' script, overriding
+default config path used ([/alembic.ini](/alembic.ini)) and the sqlalchemy.url used, that is
+automatically computed by the environment and not read from any ini file.
+
+All the database releases where you can migrate up and down must be defined by files contained inside
+the folder [/alembic/versions](/alembic/versions). All these files are in a version queue: each file has
+link to its revision hash (variable 'revision', the prefix of the file name) and to the next older one
+(variable 'down_revision'), and contains code to step up and down that database version.\
+Some useful commands are listed below.
+
+- To migrate to the newest version, type:\
+  `broker-alembic-cli upgrade head`
+- To upgrade to a specific version hash, for example 8ccbe515155c, type:\
+  `broker-alembic-cli upgrade 8ccbe515155c`
+- To downgrade to a specific version hash, for example 8ccbe515155c, type:\
+  `broker-alembic-cli downgrade 8ccbe515155c`
+- To get the current version hash of the database, type:\
+  `broker-alembic-cli current`
+
+Database migration changes could be applied to the cacholote component of the database, too. In such case,
+migrate the cacholote component after the migration by the 'broker-alembic-cli' tool.
+
+Other details are the same of the standard alembic migration tool,
+see the [Alembic tutorial](https://alembic.sqlalchemy.org/en/latest/tutorial.html).
+
 For details about the alembic migration tool, see the [Alembic tutorial](https://alembic.sqlalchemy.org/en/latest/tutorial.html).
 
 ## Quality of Service rules examples

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]

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,

alembic/versions/a4e8be715296_add_deleted_as_new_status.py

@@ -46,4 +46,9 @@ def upgrade() -> None:
 def downgrade() -> None:
     # Remove the new status from the enum
     # this doesn't work
-    op.execute("ALTER TYPE status DELETE VALUE 'deleted'")
+    #op.execute("ALTER TYPE status DELETE VALUE 'deleted'")
+    op.execute("CREATE TYPE status_old AS ENUM ('accepted','running','failed','successful','dismissed')")
+    op.execute("DELETE FROM system_requests where status='deleted'")
+    op.execute("ALTER TABLE system_requests ALTER COLUMN status TYPE status_old USING (status::text::status_old)")
+    op.execute("DROP TYPE status")
+    op.execute("ALTER TYPE status_old RENAME TO status")

alembic/versions/d5d4afc97d40_more_indexes_on_system_requests.py

@@ -26,6 +26,9 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
-    alembic.op.drop_index("idx_system_requests_status")
-    alembic.op.drop_index("idx_system_requests_created_at")
-    alembic.op.drop_index("idx_system_requests_finished_at")
+    alembic.op.drop_index("idx_system_requests_status", if_exists=True)
+    alembic.op.drop_index("idx_system_requests_created_at", if_exists=True)
+    alembic.op.drop_index("idx_system_requests_finished_at", if_exists=True)
+    alembic.op.drop_index("ix_system_requests_status", if_exists=True)
+    alembic.op.drop_index("ix_system_requests_created_at", if_exists=True)
+    alembic.op.drop_index("ix_system_requests_finished_at", if_exists=True)

cads_broker/alembic_cli.py

@@ -0,0 +1,35 @@
+"""Custom alembic CLI with new default config path + db url by environment."""
+
+import os
+from typing import Optional, Sequence
+
+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: Optional[Sequence[str]] = None) -> 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() -> None:
+    cli = MyCommandLine(prog="broker-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()

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

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