From d19e07a64b19ba3c92c1149c65aab853adb7059f Mon Sep 17 00:00:00 2001 From: Peter Boers Date: Thu, 4 Mar 2021 20:01:22 +0100 Subject: [PATCH] Add database interaction commands --- .bumpversion.cfg | 2 +- orchestrator/__init__.py | 2 +- orchestrator/app.py | 35 ++++++++++- orchestrator/cli/database.py | 104 ++++++++++++++++++++++++++++++++ orchestrator/cli/main.py | 8 +-- orchestrator/db/__init__.py | 1 + orchestrator/domain/__init__.py | 2 +- orchestrator/forms/__init__.py | 2 +- setup.cfg | 1 + 9 files changed, 147 insertions(+), 10 deletions(-) create mode 100644 orchestrator/cli/database.py diff --git a/.bumpversion.cfg b/.bumpversion.cfg index de3e95107..5ad3b51b7 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.0.2rc2 +current_version = 0.0.2rc3 commit = False tag = False parse = (?P\d+)\.(?P\d+)\.(?P\d+)(\-(?P[a-z]+)(?P\d+))? diff --git a/orchestrator/__init__.py b/orchestrator/__init__.py index f95c06092..8674fa789 100644 --- a/orchestrator/__init__.py +++ b/orchestrator/__init__.py @@ -13,7 +13,7 @@ """This is the orchestrator workflow engine.""" -__version__ = "0.0.2rc2" +__version__ = "0.0.2rc3" from orchestrator.app import OrchestratorCore from orchestrator.settings import app_settings, oauth2_settings diff --git a/orchestrator/app.py b/orchestrator/app.py index 14310f5a7..7880bb18e 100644 --- a/orchestrator/app.py +++ b/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() diff --git a/orchestrator/cli/database.py b/orchestrator/cli/database.py new file mode 100644 index 000000000..ee61c9af8 --- /dev/null +++ b/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) diff --git a/orchestrator/cli/main.py b/orchestrator/cli/main.py index dc1db04de..532efedb7 100644 --- a/orchestrator/cli/main.py +++ b/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() diff --git a/orchestrator/db/__init__.py b/orchestrator/db/__init__.py index 4df8a4a5b..26b86d77a 100644 --- a/orchestrator/db/__init__.py +++ b/orchestrator/db/__init__.py @@ -86,4 +86,5 @@ def init_database(settings: AppSettings) -> Database: "UtcTimestamp", "UtcTimestampException", "db", + "init_database", ] diff --git a/orchestrator/domain/__init__.py b/orchestrator/domain/__init__.py index b1579b7d6..580b647b0 100644 --- a/orchestrator/domain/__init__.py +++ b/orchestrator/domain/__init__.py @@ -20,4 +20,4 @@ __doc__ = make_product_type_index_doc(SUBSCRIPTION_MODEL_REGISTRY) -__all__ = ["SubscriptionModel"] +__all__ = ["SubscriptionModel", "SUBSCRIPTION_MODEL_REGISTRY"] diff --git a/orchestrator/forms/__init__.py b/orchestrator/forms/__init__.py index 2d3bbd6c1..6947dc6ff 100644 --- a/orchestrator/forms/__init__.py +++ b/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): diff --git a/setup.cfg b/setup.cfg index 00624bed1..b66d6ac75 100644 --- a/setup.cfg +++ b/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