Create RDS database snapshot before executing alembic migrations

backend/dbmigrations_handler.py

@@ -4,15 +4,67 @@
 
 import logging
 import os
-
+import datetime
+import boto3
+import time
 from alembic import command
+from alembic.script import ScriptDirectory
+from alembic.migration import MigrationContext
 from alembic.config import Config
+from dataall.base.db.connection import ENVNAME, get_engine
 
 logger = logging.getLogger()
 logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))
 
 
 def handler(event, context) -> None:
+    """
+    This function will be called once upon every deployment.
+    It checks if there are any alembic migration scripts to execute.
+    If there are, it will create a snapshot of the database.
+    It executes the alembic migration scripts.
+    """
     alembic_cfg = Config('alembic.ini')
     alembic_cfg.set_main_option('script_location', './migrations')
+
+    # Get head version
+    script = ScriptDirectory.from_config(alembic_cfg)
+    head_rev = script.get_current_head()
+
+    # Get current version from database
+    engine = get_engine(ENVNAME)
+    with engine.engine.connect() as connection:
+        context = MigrationContext.configure(connection)
+        current_rev = context.get_current_revision()
+
+    if head_rev != current_rev:
+        snapshot_id = f'dataall-migration-{head_rev}-{datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")}'
+        cluster_id = engine.dbconfig.host.split('.')[0]
+        logger.info(
+            f'Creating RDS snapshot for cluster {cluster_id}, head revision {head_rev} is ahead of {current_rev}...'
+        )
+        try:
+            rds_client = boto3.client('rds', region_name=os.getenv('AWS_REGION'))
+            cluster_id = 'res-fr-db'
+            cluster_status = ''
+            while cluster_status != 'available':
+                # Edge case in which the cluster is performing backup and/or maintenance operations.
+                # If it times out the CICD pipeline fails and needs to be retried.
+                logger.info(f'Waiting while the cluster is available, status={cluster_status}')
+                response = rds_client.describe_db_clusters(DBClusterIdentifier=cluster_id)
+                cluster_status = response['DBClusters'][0]['Status']
+                time.sleep(30)
+
+            rds_client.create_db_cluster_snapshot(
+                DBClusterSnapshotIdentifier=snapshot_id,
+                DBClusterIdentifier=cluster_id,
+                Tags=[
+                    {'Key': 'Application', 'Value': 'dataall'},
+                ],
+            )
+        except Exception as e:
+            logger.exception(f'Failed to create RDS snapshot: {e}')
+            raise Exception(f'Failed to create RDS snapshot: {e}')
+
+    # Execute the alembic migration scripts
     command.upgrade(alembic_cfg, 'head')  # logging breaks after this command

deploy/stacks/backend_stack.py

@@ -322,6 +322,16 @@ def __init__(
             ecr_repository=repo,
             execute_after=[aurora_stack.cluster],
             connectables=[aurora_stack.cluster],
+            additional_policy_statements=[
+                iam.PolicyStatement(
+                    effect=iam.Effect.ALLOW,
+                    actions=['rds:AddTagsToResource', 'rds:CreateDBClusterSnapshot', 'rds:DescribeDBClusters'],
+                    resources=[
+                        f'arn:aws:rds:*:{self.account}:snapshot:dataall*',
+                        f'arn:aws:rds:*:{self.account}:cluster:dataall*',
+                    ],
+                )
+            ],
             **kwargs,
         )
 

deploy/stacks/trigger_function_stack.py

@@ -27,6 +27,7 @@ def __init__(
         vpce_connection: ec2.IConnectable = None,
         connectables: List[ec2.IConnectable] = [],
         execute_after: List[Construct] = [],
+        additional_policy_statements: List[iam.PolicyStatement] = [],
         **kwargs,
     ):
         super().__init__(scope, id, **kwargs)
@@ -38,13 +39,14 @@ def __init__(
         env = {'envname': envname, 'LOG_LEVEL': 'INFO'}
 
         function_sgs = self.create_lambda_sgs(envname, handler, resource_prefix, vpc)
+        policy_statements = self.get_policy_statements(resource_prefix).append(additional_policy_statements)
 
         self.trigger_function = TriggerFunction(
             self,
             f'TriggerFunction-{handler}',
             function_name=f'{resource_prefix}-{envname}-{handler.replace(".", "_")}',
             description=f'dataall {handler} trigger function',
-            initial_policy=self.get_policy_statements(resource_prefix),
+            initial_policy=policy_statements,
             code=_lambda.Code.from_ecr_image(repository=ecr_repository, tag=image_tag, cmd=[handler]),
             vpc=vpc,
             security_groups=[function_sgs],