Merge pull request #374 from ral-facilities/add-migration-script-#339

Add `ims-migrate` script #339

Dockerfile

@@ -2,6 +2,9 @@ FROM python:3.12.5-alpine3.20@sha256:bb5d0ac04679d78a1258e7dfacdb4d9bdefe9a10480
 
 WORKDIR /inventory-management-system-api-run
 
+# Requirement when using a different workdir to get scripts to import correctly
+ENV PYTHONPATH="${PYTHONPATH}:/inventory-management-system-api-run"
+
 COPY pyproject.toml ./
 COPY inventory_management_system_api/ inventory_management_system_api/
 

Dockerfile.prod

@@ -2,6 +2,9 @@ FROM python:3.12.5-alpine3.20@sha256:bb5d0ac04679d78a1258e7dfacdb4d9bdefe9a10480
 
 WORKDIR /inventory-management-system-api-run
 
+# Requirement when using a different workdir to get scripts to import correctly
+ENV PYTHONPATH="${PYTHONPATH}:/inventory-management-system-api-run"
+
 COPY README.md pyproject.toml ./
 # Copy inventory_management_system_api source files
 COPY inventory_management_system_api/ inventory_management_system_api/

README.md

@@ -322,3 +322,33 @@ encoding the token. Once the JWT access token is decoded successfully, it checks
 payload, and it has not expired. This means that any microservice can be used to generate JWT access tokens so long as
 it meets the above criteria. The [LDAP-JWT Authentication Service](https://github.com/ral-facilities/ldap-jwt-auth) is
 a microservice that provides user authentication against an LDAP server and returns a JWT access token.
+
+### Migrations
+
+Migration scripts are located inside the `inventory_management_system/migrations/scripts`. See the
+`example_migration.py` for an example on how to implement one. Any new migrations added should be automatically picked
+up and shown via
+
+```bash
+ims-migrate list
+```
+
+or
+
+```bash
+docker exec -it inventory_management_system_api_container ims-migrate list
+```
+
+if running in Docker.
+
+To perform a migration you should use
+
+```bash
+ims-migrate forward <migration_name>
+```
+
+To revert the same migration use
+
+```bash
+ims-migrate backward <migration_name>
+```

inventory_management_system_api/migrations/migration.py

@@ -0,0 +1,160 @@
+"""Module for providing a migration script"""
+
+import argparse
+import importlib
+import logging
+from abc import ABC, abstractmethod
+
+from pymongo.client_session import ClientSession
+from pymongo.database import Database
+
+from inventory_management_system_api.core.database import get_database, mongodb_client
+
+logger = logging.getLogger()
+
+
+class BaseMigration(ABC):
+    """Base class for a migration with a forward and backward step"""
+
+    @abstractmethod
+    def __init__(self, database: Database):
+        pass
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Description of this migration"""
+        return ""
+
+    @abstractmethod
+    def forward(self, session: ClientSession):
+        """Method for executing the migration"""
+
+    def forward_after_transaction(self, session: ClientSession):
+        """Method called after the forward function is called to do anything that can't be done inside a transaction
+        (ONLY USE IF NECESSARY e.g. dropping a collection)"""
+
+    @abstractmethod
+    def backward(self, session: ClientSession):
+        """Method for reversing the migration"""
+
+    def backward_after_transaction(self, session: ClientSession):
+        """Method called after the backward function is called to do anything that can't be done inside a transaction
+        (ONLY USE IF NECESSARY e.g. dropping a collection)"""
+
+
+class SubCommand(ABC):
+    """Base class for a sub command"""
+
+    def __init__(self, help_message: str):
+        self.help_message = help_message
+
+    @abstractmethod
+    def setup(self, parser: argparse.ArgumentParser):
+        """Setup the parser by adding any parameters here"""
+
+    @abstractmethod
+    def run(self, args: argparse.Namespace):
+        """Run the command with the given parameters as added by 'setup'"""
+
+
+def load_migration(name: str) -> BaseMigration:
+    """Loads a migration script from the scripts module"""
+
+    migration_module = importlib.import_module(f"inventory_management_system_api.migrations.scripts.{name}")
+    migration_class = getattr(migration_module, "Migration", None)
+
+    database = get_database()
+    return migration_class(database)
+
+
+class CommandList(SubCommand):
+    """Command that lists available database migrations"""
+
+    def __init__(self):
+        super().__init__(help_message="List all available database migrations")
+
+    def setup(self, parser: argparse.ArgumentParser):
+        pass
+
+    def run(self, args: argparse.Namespace):
+        # Find a list of all available migration scripts
+        with importlib.resources.path("inventory_management_system_api.migrations.scripts", "") as scripts_path:
+            files_in_scripts = list(scripts_path.iterdir())
+            available_migrations = list(
+                filter(lambda name: "__" not in name, [file.name.replace(".py", "") for file in files_in_scripts])
+            )
+        for migration_name in available_migrations:
+            migration = load_migration(migration_name)
+
+            print(f"{migration_name} - {migration.description}")
+
+
+class CommandForward(SubCommand):
+    """Command that performs a forward database migration"""
+
+    def __init__(self):
+        super().__init__(help_message="Performs a forward database migration")
+
+    def setup(self, parser: argparse.ArgumentParser):
+        parser.add_argument("migration", help="Name of the migration to perform")
+
+    def run(self, args: argparse.Namespace):
+        migration_instance: BaseMigration = load_migration(args.migration)
+
+        # Run migration inside a session to lock writes and revert the changes if it fails
+        with mongodb_client.start_session() as session:
+            with session.start_transaction():
+                logger.info("Performing forward migration...")
+                migration_instance.forward(session)
+            # Run some things outside the transaction e.g. if needing to drop a collection
+            migration_instance.forward_after_transaction(session)
+            logger.info("Done!")
+
+
+class CommandBackward(SubCommand):
+    """Command that performs a backward database migration"""
+
+    def __init__(self):
+        super().__init__(help_message="Performs a backward database migration")
+
+    def setup(self, parser: argparse.ArgumentParser):
+        parser.add_argument("migration", help="Name of the migration to revert")
+
+    def run(self, args: argparse.Namespace):
+        migration_instance: BaseMigration = load_migration(args.migration)
+
+        # Run migration inside a session to lock writes and revert the changes if it fails
+        with mongodb_client.start_session() as session:
+            with session.start_transaction():
+                logger.info("Performing backward migration...")
+                migration_instance.backward(session)
+            # Run some things outside the transaction e.g. if needing to drop a collection
+            migration_instance.backward_after_transaction(session)
+            logger.info("Done!")
+
+
+# List of subcommands
+commands: dict[str, SubCommand] = {
+    "list": CommandList(),
+    "forward": CommandForward(),
+    "backward": CommandBackward(),
+}
+
+
+def main():
+    """Entrypoint for the ims-migrate script"""
+
+    parser = argparse.ArgumentParser()
+
+    subparser = parser.add_subparsers(dest="command")
+
+    for command_name, command in commands.items():
+        command_parser = subparser.add_parser(command_name, help=command.help_message)
+        command.setup(command_parser)
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    commands[args.command].run(args)

inventory_management_system_api/migrations/scripts/example_migration.py

@@ -0,0 +1,65 @@
+"""
+Module providing an example migration that does nothing
+"""
+
+import logging
+from typing import Collection
+
+from pymongo.client_session import ClientSession
+from pymongo.database import Database
+
+from inventory_management_system_api.migrations.migration import BaseMigration
+
+logger = logging.getLogger()
+
+# When the migration will modify database models by adding data may be a good idea to put the old here and pass any data
+# between them and the new ones before updating them in the database, to ensure all modifications are as expected
+# e.g.
+
+
+# class OldUnit(BaseModel):
+#     """
+#     Old database model for a Unit
+#     """
+
+#     value: str
+#     code: str
+
+
+class Migration(BaseMigration):
+    """Example migration that does nothing"""
+
+    description = "Example migration that does nothing"
+
+    def __init__(self, database: Database):
+        """Obtain any collections required for the migration here e.g."""
+
+        self._units_collection: Collection = database.units
+
+    def forward(self, session: ClientSession):
+        """This function should actually perform the migration
+
+        All database functions should be given the session in order to ensure all updates are done within a transaction
+        """
+
+        # Perform any database updates here e.g. for renaming a field
+
+        # self._units_collection.update_many(
+        #     {}, {"$rename": {"value": "renamed_value"}}, session=session
+        # )
+
+        logger.info("example_migration forward migration (that does nothing)")
+
+    def backward(self, session: ClientSession):
+        """This function should reverse the migration
+
+        All database functions should be given the session in order to ensure all updates are done within a transaction
+        """
+
+        # Perform any database updates here e.g. for renaming a field
+
+        # self._units_collection.update_many(
+        #     {}, {"$rename": {"renamed_value": "value"}}, session=session
+        # )
+
+        logger.info("example_migration backward migration (that does nothing)")

pyproject.toml

@@ -16,6 +16,9 @@ dependencies = [
 [project.urls]
 "Repository" = "https://github.com/ral-facilities/inventory-management-system-api"
 
+[project.scripts]
+"ims-migrate" = "inventory_management_system_api.migrations.migration:main"
+
 [project.optional-dependencies]
 code-analysis = [
     "black==24.8.0",

scripts/dev_cli.py

@@ -63,6 +63,7 @@ def run_mongodb_command(args: list[str], stdin: Optional[TextIOWrapper] = None,
 
 def add_mongodb_auth_args(parser: argparse.ArgumentParser):
     """Adds common arguments for MongoDB authentication"""
+
     parser.add_argument("-u", "--username", default="root", help="Username for MongoDB authentication")
     parser.add_argument("-p", "--password", default="example", help="Password for MongoDB authentication")
 
@@ -81,8 +82,8 @@ def get_mongodb_auth_args(args: argparse.Namespace):
 class SubCommand(ABC):
     """Base class for a sub command"""
 
-    def __init__(self, help: str):
-        self.help = help
+    def __init__(self, help_message: str):
+        self.help_message = help_message
 
     @abstractmethod
     def setup(self, parser: argparse.ArgumentParser):
@@ -103,7 +104,7 @@ class CommandDBInit(SubCommand):
     """
 
     def __init__(self):
-        super().__init__(help="Initialise database for development (using docker on linux)")
+        super().__init__(help_message="Initialise database for development (using docker on linux)")
 
     def setup(self, parser: argparse.ArgumentParser):
         add_mongodb_auth_args(parser)
@@ -168,7 +169,7 @@ class CommandDBImport(SubCommand):
     """Command that imports mock data into the database"""
 
     def __init__(self):
-        super().__init__(help="Imports database for development")
+        super().__init__(help_message="Imports database for development")
 
     def setup(self, parser: argparse.ArgumentParser):
         add_mongodb_auth_args(parser)
@@ -198,7 +199,7 @@ class CommandDBGenerate(SubCommand):
     """
 
     def __init__(self):
-        super().__init__(help="Generates new test data for the database and dumps it")
+        super().__init__(help_message="Generates new test data for the database and dumps it")
 
     def setup(self, parser: argparse.ArgumentParser):
         add_mongodb_auth_args(parser)
@@ -272,7 +273,7 @@ def main():
     subparser = parser.add_subparsers(dest="command")
 
     for command_name, command in commands.items():
-        command_parser = subparser.add_parser(command_name, help=command.help)
+        command_parser = subparser.add_parser(command_name, help=command.help_message)
         command.setup(command_parser)
 
     args = parser.parse_args()