Setup Migrations environment correctly

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.2rc7
+current_version = 0.0.2rc11
 commit = False
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\-(?P<release>[a-z]+)(?P<build>\d+))?

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "0.0.2rc7"
+__version__ = "0.0.2rc11"
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.settings import app_settings, oauth2_settings

orchestrator/app.py

@@ -74,7 +74,6 @@ def __init__(  # type: ignore
         initialise_logging()
 
         self.include_router(api_router, prefix="/api")
-        self.base_settings = base_settings
 
         init_database(base_settings)
 
@@ -99,7 +98,6 @@ def _index() -> str:
 
     def instrument_app(self) -> None:
         logger.info("Activating Opentelemetry tracing to app", app=self.title)
-        self.base_settings.TRACING_ENABLED = True
         trace.set_tracer_provider(tracer_provider)
         FastAPIInstrumentor.instrument_app(self)
         RequestsInstrumentor().instrument()
@@ -127,7 +125,7 @@ def add_sentry(
         self.add_middleware(SentryAsgiMiddleware)
 
     @staticmethod
-    def register_subscription_model(product_to_subscription_model_mapping: Dict[str, Type[SubscriptionModel]]) -> None:
+    def register_subscription_models(product_to_subscription_model_mapping: Dict[str, Type[SubscriptionModel]]) -> None:
         """
         Register your subscription models.
 

orchestrator/cli/database.py

@@ -1,8 +1,11 @@
 import os
 from shutil import copyfile
+from typing import List, Optional
 
 import jinja2
 import typer
+from alembic import command
+from alembic.config import Config
 from structlog import get_logger
 
 import orchestrator
@@ -20,13 +23,25 @@
 )
 
 
-@app.command(name="init")
+def alembic_cfg() -> Config:
+    cfg = Config("alembic.ini")
+    version_locations = cfg.get_main_option("version_locations")
+    cfg.set_main_option(
+        "version_locations", f"{version_locations} {orchestrator_module_location}/{migration_dir}/versions/schema"
+    )
+    logger.info("Version Locations", locations=cfg.get_main_option("version_locations"))
+    return cfg
+
+
+@app.command(
+    help="Initialize an empty migrations environment. This command will throw an exception when it detects conflicting files and directories."
+)
 def init() -> None:
     """
-    Run the migrations.
+    Initialise the migrations directory.
 
-    This command will run the migrations for initialization of the database. If you have extra migrations that need to be run,
-    add this to the
+    This command will initialize a migration directory for the orchestrator core application and setup a correct
+    migration environment.
 
     Returns:
         None
@@ -45,7 +60,7 @@ def init() -> None:
     logger.info("Creating directory", directory=os.path.abspath(versions_schema))
     os.makedirs(versions_schema)
 
-    source_env_py = os.path.join(orchestrator_module_location, f"{migration_dir}/env.py")
+    source_env_py = os.path.join(orchestrator_module_location, f"{migration_dir}/templates/env.py.j2")
     env_py = os.path.join(migration_dir, "env.py")
     logger.info("Creating file", file=os.path.abspath(env_py))
     copyfile(source_env_py, env_py)
@@ -65,13 +80,83 @@ def init() -> None:
     if not os.access(os.path.join(os.getcwd(), "alembic.ini"), os.F_OK):
         logger.info("Creating file", file=os.path.join(os.getcwd(), "alembic.ini"))
         with open(os.path.join(os.getcwd(), "alembic.ini"), "w") as alembic_ini:
-            alembic_ini.write(
-                template.render(
-                    migrations_dir=migration_dir,
-                    module_migrations_dir=os.path.join(
-                        orchestrator_module_location, f"{migration_dir}/versions/schema"
-                    ),
-                )
-            )
+            alembic_ini.write(template.render(migrations_dir=migration_dir))
     else:
         logger.info("Skipping Alembic.ini file. It already exists")
+
+
+@app.command(help="Get the database heads")
+def heads():
+    command.heads(alembic_cfg())
+
+
+@app.command(help="Merge database revisions.")
+def merge(
+    revisions: Optional[List[str]] = typer.Argument(
+        None, help="Add the revision you would like to merge to this command."
+    ),
+    message: str = typer.Option(None, "--message", "-m", help="The revision message"),
+) -> None:
+    """
+    Merge database revisions.
+
+    Args:
+        revisions: List of revisions to merge
+        message: Optional message for the revision.
+
+    Returns:
+        None
+
+    """
+    command.merge(alembic_cfg(), revisions, message=message)
+
+
+@app.command()
+def upgrade(revision: Optional[str] = typer.Argument(None, help="Rev id to upgrade to")) -> None:
+    """
+    Upgrade the database.
+
+    Args:
+        revision: Optional argument to indicate where to upgrade to.
+
+    Returns:
+        None
+
+    """
+    command.upgrade(alembic_cfg(), revision)
+
+
+@app.command()
+def downgrade(revision: Optional[str] = typer.Argument(None, help="Rev id to upgrade to")) -> None:
+    """
+    Downgrade the database.
+
+    Args:
+        revision: Optional argument to indicate where to downgrade to.
+
+    Returns:
+        None
+
+    """
+    command.downgrade(alembic_cfg(), revision)
+
+
+@app.command()
+def revision(
+    message: str = typer.Option(None, "--message", "-m", help="The revision message"),
+    autogenerate: bool = typer.Option(False, help="Detect schema changes and add migrations"),
+    head: str = typer.Option(None, help="Determine the head the head you need to add your migration to."),
+) -> None:
+    """
+    Create a new revision file.
+
+    Args:
+        message: The revision message
+        autogenerate: Whether to detect schema changes.
+        head: To which head the migration applies
+
+    Returns:
+        None
+
+    """
+    command.revision(alembic_cfg(), message, autogenerate=autogenerate, head=head)

orchestrator/cli/main.py

@@ -11,14 +11,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import typer
 
 from orchestrator.cli import database, scheduler
 
 app = typer.Typer()
 app.add_typer(scheduler.app, name="scheduler", help="Access all the scheduler functions")
-app.add_typer(database.app, name="db", help="interact with the database")
+app.add_typer(database.app, name="db", help="Interact with the application database")
+
 
 if __name__ == "__main__":
     app()

orchestrator/migrations/templates/alembic.ini.j2

@@ -8,7 +8,7 @@ file_template = %%(year)d-%%(month).2d-%%(day).2d_%%(rev)s_%%(slug)s
 # the 'revision' command, regardless of autogenerate
 # revision_environment = false
 script_location = {{ migrations_dir }}
-version_locations = {{ module_migrations_dir }}/versions/schema {{ migrations_dir }}/versions/general
+version_locations = {{ migrations_dir }}/versions/general
 # Logging configuration
 [loggers]
 keys = root,sqlalchemy,alembic

orchestrator/migrations/templates/env.py.j2

@@ -0,0 +1,100 @@
+import logging
+import os
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+import orchestrator
+from orchestrator.db.database import BaseModel
+from orchestrator.settings import app_settings
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+logger = logging.getLogger("alembic.env")
+
+config.set_main_option("sqlalchemy.url", app_settings.DATABASE_URI)
+version_locations = config.get_main_option("version_locations")
+config.set_main_option(
+    "version_locations", f"{os.path.dirname(orchestrator.__file__)}/migrations/versions/schema {version_locations}"
+)
+context.script.version_locations = [
+    f"{os.path.dirname(orchestrator.__file__)}/migrations/versions/schema",
+    f"{version_locations}",
+]
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+target_metadata = BaseModel.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url, target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+
+    # this callback is used to prevent an auto-migration from being generated
+    # when there are no changes to the schema
+    # reference: http://alembic.zzzcomputing.com/en/latest/cookbook.html
+    def process_revision_directives(context, revision, directives):  # type: ignore
+        if getattr(config.cmd_opts, "autogenerate", False):
+            script = directives[0]
+            if script.upgrade_ops.is_empty():
+                directives[:] = []
+                logger.info("No changes in schema detected.")
+
+    engine = engine_from_config(
+        config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool
+    )
+
+    connection = engine.connect()
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        process_revision_directives=process_revision_directives,
+        compare_type=True,
+    )
+    try:
+        with context.begin_transaction():
+            context.run_migrations()
+    finally:
+        connection.close()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()