AAI-119 set up migrations for aai-backend database

README.md

@@ -71,6 +71,23 @@ uv run -- ruff check . --fix
 pre-commit run --all-files
 ```
 
+# Database management
+
+The deployed service uses a Postgres database on AWS RDS.
+In order to generate migrations for the database locally,
+we use a Postgres docker container to generate migrations against.
+
+After making any changes to the database models, run the
+`generate_migrations.py` script to create migrations:
+
+```shell
+python generate_migrations.py -m migration_name
+```
+
+and commit them to git. Once your updated code has been
+deployed on AWS, you can use `aws ecs execute-command`
+to run the migrations.
+
 # Deployment
 
 Currently the service is deployed to AWS via the CDK scripts in `deploy/`,

deploy/aai_backend_deploy/aai_backend_deploy_stack.py

@@ -22,6 +22,7 @@
 from aws_cdk import (
     aws_elasticloadbalancingv2 as elbv2,
 )
+from aws_cdk import aws_iam as iam
 from aws_cdk import (
     aws_route53 as route53,
 )
@@ -60,6 +61,10 @@ def __init__(self, scope: Construct, construct_id: str, config: dict, **kwargs)
         task_definition = ecs.FargateTaskDefinition(self, "AaiBackendTaskDef",
                                                     memory_limit_mib=1024,
                                                     cpu=512)
+        # Allow executing comands in the ECS container
+        task_definition.task_role.add_managed_policy(
+            iam.ManagedPolicy.from_aws_managed_policy_name("AmazonSSMManagedInstanceCore")
+        )
 
         container = task_definition.add_container(
             "FastAPIContainer",
@@ -70,12 +75,12 @@ def __init__(self, scope: Construct, construct_id: str, config: dict, **kwargs)
             # Set an env variable to the current time to force redeploy -
             #   might be better to use an image tag in future
             environment={
-                "FORCE_REDEPLOY": str(datetime.datetime.now())
+                "FORCE_REDEPLOY": str(datetime.datetime.now()),
+                "DB_HOST": self.db_host,
             },
             secrets={
                 "DB_USER": ecs.Secret.from_secrets_manager(db_secret, field="username"),
                 "DB_PASSWORD": ecs.Secret.from_secrets_manager(db_secret, field="password"),
-                "DB_HOST": self.db_host,
             },
             logging=ecs.LogDrivers.aws_logs(stream_prefix="FastAPI"),
         )
@@ -101,6 +106,7 @@ def __init__(self, scope: Construct, construct_id: str, config: dict, **kwargs)
             domain_zone=route53.HostedZone.from_lookup(
                 self, "AaiBackendZone", domain_name=self.zone_domain
             ),
+            enable_execute_command=True
         )
 
         service.target_group.configure_health_check(path="/", healthy_http_codes="200-399")

generate_migrations.py

@@ -0,0 +1,62 @@
+import os
+import subprocess
+import time
+
+import click
+
+DB_CONTAINER_NAME = "temp_alembic_db"
+DEFAULT_POSTGRES_IMAGE = "postgres:17"
+DEFAULT_PORT = 5433
+POSTGRES_USER = "testuser"
+POSTGRES_PASSWORD = "testpass"
+
+
+def run(cmd, env=None):
+    print(f"> {cmd}")
+    subprocess.run(cmd, shell=True, check=True, env=env or os.environ)
+
+
+@click.command()
+@click.option('--revision-message', '-m', required=False, help="Message for Alembic revision.")
+@click.option('--check', is_flag=True, help="Only run 'alembic check' after DB container is up.")
+def generate_migrations(revision_message, check):
+    """Spin up a temp Postgres DB, apply migrations or run alembic check."""
+    if not check and not revision_message:
+        raise click.UsageError("Missing option '-m' / '--revision-message'. Required unless using --check.")
+
+    database_url = f"localhost:{DEFAULT_PORT}"
+    os.environ["DB_HOST"] = database_url
+    os.environ["DB_USER"] = POSTGRES_USER
+    os.environ["DB_PASSWORD"] = POSTGRES_PASSWORD
+
+    try:
+        print("🚀 Starting temporary Postgres container...")
+        run(
+            f"docker run --rm -d --name {DB_CONTAINER_NAME} "
+            f"-e POSTGRES_USER={POSTGRES_USER} "
+            f"-e POSTGRES_PASSWORD={POSTGRES_PASSWORD} "
+            f"-p {DEFAULT_PORT}:5432 {DEFAULT_POSTGRES_IMAGE}"
+        )
+
+        print("⏳ Waiting for database to be ready...")
+        time.sleep(5)  # Could be enhanced with pg_isready
+
+        if check:
+            print("🧱 Applying existing Alembic migrations...")
+            run("alembic upgrade head")
+            print("🔍 Running 'alembic check'...")
+            run("alembic check")
+        else:
+            print("🧱 Applying existing Alembic migrations...")
+            run("alembic upgrade head")
+
+            print("📝 Generating new Alembic revision...")
+            run(f"alembic revision --autogenerate -m \"{revision_message}\"")
+
+    finally:
+        print("🧹 Cleaning up: stopping container...")
+        subprocess.run(f"docker stop {DB_CONTAINER_NAME}", shell=True)
+
+
+if __name__ == "__main__":
+    generate_migrations()

migrations/versions/101f45395233_add_group_membership.py

@@ -0,0 +1,41 @@
+"""add_group_membership
+
+Revision ID: 101f45395233
+Revises:
+Create Date: 2025-07-01 16:29:48.072722
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+
+
+# revision identifiers, used by Alembic.
+revision: str = '101f45395233'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('groupmembership',
+    sa.Column('id', sa.Uuid(), nullable=False),
+    sa.Column('group', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('user_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('user_email', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('approval_status', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('updated_at', sa.DateTime(), nullable=False),
+    sa.Column('updated_by_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('updated_by_email', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('groupmembership')
+    # ### end Alembic commands ###