Add database interaction commands

orchestrator/app.py

@@ -12,10 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, Type
+from typing import Dict, Optional, Type
 
 import sentry_sdk
 import structlog
+import typer
 from fastapi.applications import FastAPI
 from fastapi_etag.dependency import add_exception_handler
 from nwastdlib.logging import initialise_logging
@@ -34,8 +35,10 @@
 
 from orchestrator.api.api_v1.api import api_router
 from orchestrator.api.error_handling import ProblemDetailException
+from orchestrator.cli.main import app as cli_app
 from orchestrator.db import db, init_database
 from orchestrator.db.database import DBSessionMiddleware
+from orchestrator.domain import SUBSCRIPTION_MODEL_REGISTRY, SubscriptionModel
 from orchestrator.exception_handlers import form_error_handler, problem_detail_handler
 from orchestrator.forms import FormException
 from orchestrator.settings import AppSettings, app_settings, tracer_provider
@@ -120,3 +123,33 @@ def add_sentry(
             integrations=[SqlalchemyIntegration(), RedisIntegration()],
         )
         self.add_middleware(SentryAsgiMiddleware)
+
+    @staticmethod
+    def register_subscription_model(product_to_subscription_model_mapping: Dict[str, Type[SubscriptionModel]]) -> None:
+        """
+        Register your subscription models.
+
+        This method is needed to register your subscription models inside the orchestrator core.
+
+        Args:
+            product_to_subscription_model_mapping: The dictionary should contain a mapping of products to SubscriptionModels.
+                The selection will be done depending on the name of the product.
+
+        Returns:
+            None:
+
+        Examples:
+            product_to_subscription_model_mapping = {
+                "Generic Product One": GenericProductModel,
+                "Generic Product Two": GenericProductModel,
+            }
+
+        """
+        SUBSCRIPTION_MODEL_REGISTRY.update(product_to_subscription_model_mapping)
+
+
+main_typer_app = typer.Typer()
+main_typer_app.add_typer(cli_app, name="orchestrator", help="The are the orchestrator cli commands")
+
+if __name__ == "__main__":
+    main_typer_app()

orchestrator/cli/database.py

@@ -0,0 +1,104 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+import typer
+from alembic import command
+from alembic.config import Config
+from alembic.util import CommandError
+from structlog import get_logger
+
+import orchestrator
+
+logger = get_logger(__name__)
+
+app: typer.Typer = typer.Typer()
+
+this_location = os.path.join(os.path.dirname(os.path.realpath(__file__)))
+orchestrator_module_location = os.path.dirname(orchestrator.__file__)
+alembic_cfg = Config(file_=os.path.join(orchestrator_module_location, "migrations/alembic.ini"))
+
+
+@app.command(name="upgrade")
+def run_migrations(
+    custom_migration_directory: Optional[Path] = typer.Option(None, help="The path towards the migration directory")
+) -> None:
+    """
+    Run the migrations.
+
+    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
+
+    Args:
+        custom_migration_directory: Path to the migration directory.
+
+    Returns:
+        None
+    """
+    logger.info("Running migrations on the database", extra_migration_directory=str(custom_migration_directory))
+    if custom_migration_directory:
+        alembic_cfg.set_main_option(
+            "version_locations",
+            f"{os.path.join(orchestrator_module_location, 'migrations/versions/schema')}, {os.path.join(this_location)}/{custom_migration_directory}",
+        )
+    try:
+        command.upgrade(alembic_cfg, "heads")
+    except CommandError:
+        logger.error(
+            "Unable to run the migrations, no revisions found",
+            path=f"{os.path.join(this_location)}/{custom_migration_directory}",
+        )
+
+
+@app.command(name="heads", help="List heads")
+def list_heads() -> None:
+    """
+    List heads of the database.
+
+    Returns:
+        Heads of the database.
+
+    """
+    command.heads(alembic_cfg)
+
+
+@app.command(name="downgrade", help="Downgrade database")
+def downgrade(revision: Optional[str] = typer.Option(None, help="The revision to downgrade to")) -> None:
+    """
+    Downgrade the Database to a certain revision.
+
+    Args:
+        revision: The revision to downgrade to.
+
+    Returns:
+        None
+
+    """
+    command.downgrade(alembic_cfg, revision)
+
+
+@app.command(name="migrate", help="Migrate the database")
+def migrate(
+    custom_migration_directory: Path = typer.Argument(..., help="The path towards the migration directory"),
+    message: Optional[str] = typer.Option(None, help="The revision message"),
+    autogenerate: Optional[bool] = typer.Option(
+        False, help="Detect model changes and automatically generate migrations."
+    ),
+) -> None:
+    """
+    Migrate the database.
+
+    Args:
+        custom_migration_directory: The migration directory.
+        message: The message of the migration
+        autogenerate: whether to automatically generate schema change migrations.
+
+    Returns:
+        None
+
+    """
+    alembic_cfg.set_main_option(
+        "version_locations",
+        f"{os.path.join(orchestrator_module_location, 'migrations/versions/schema')}, {os.path.join(this_location)}/{custom_migration_directory}",
+    )
+    command.revision(alembic_cfg, message, autogenerate=autogenerate, version_path=custom_migration_directory)

orchestrator/cli/main.py

@@ -13,14 +13,12 @@
 
 
 import typer
-from surf import load_surf_cli
 
-from orchestrator.cli import scheduler
+from orchestrator.cli import database, scheduler
 
 app = typer.Typer()
-app.add_typer(scheduler.app, name="scheduler")
-
-load_surf_cli(app)
+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")
 
 if __name__ == "__main__":
     app()

orchestrator/db/__init__.py

@@ -86,4 +86,5 @@ def init_database(settings: AppSettings) -> Database:
     "UtcTimestamp",
     "UtcTimestampException",
     "db",
+    "init_database",
 ]

orchestrator/domain/__init__.py

@@ -20,4 +20,4 @@
 
 __doc__ = make_product_type_index_doc(SUBSCRIPTION_MODEL_REGISTRY)
 
-__all__ = ["SubscriptionModel"]
+__all__ = ["SubscriptionModel", "SUBSCRIPTION_MODEL_REGISTRY"]

orchestrator/forms/__init__.py

@@ -26,7 +26,7 @@
 logger = structlog.get_logger(__name__)
 
 
-__all__ = ("generate_form", "post_process", "FormNotCompleteError", "FormValidationError")
+__all__ = ("generate_form", "post_process", "FormNotCompleteError", "FormValidationError", "FormException")
 
 
 class FormException(Exception):

setup.cfg

@@ -13,6 +13,7 @@ per-file-ignores =
 	# Allow first argument to be cls instead of self for pydantic validators
 	orchestrator/*: B902
 	orchestrator/api/*: B008
+	orchestrator/cli/*: B008
 
 [tool:pytest]
 addopts=--doctest-modules --doctest-ignore-import-errors