create_ssm_rds_patch_automation_document.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: 'CloudFormation template for RDS PostgreSQL Fleet Upgrade Automation Document'

Resources:

  AutomationServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: ssm.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonSSMAutomationRole

  RDSPostgreSQLFleetUpgrade:
    Type: AWS::SSM::Document
    Properties:
      DocumentType: Automation
      Content:
        schemaVersion: '0.3'
        description: Upgrade fleet or RDS PostgreSQL instances
        parameters:
          ScriptsDirOnEC2:
            type: String
            description: The directory where scripts will be stored
            default: /tmp/rds_patch
          AssumeIAMRole:
            type: String
            description: The ARN of the role that allows Automation to perform the actions on your behalf.
            default: !GetAtt AutomationServiceRole.Arn
          EC2InstanceId:
            type: String
            description: The ID of the existing EC2 Instance
          RunPreUpgradeTasks:
            type: String
            description: Whether to run pre-upgrade tasks (PREUPGRADE) or upgrade tasks (UPGRADE)
            default: PREUPGRADE
            allowedValues:
              - PREUPGRADE
              - UPGRADE
          SnsTopicArnEmail:
            type: String
            description: SNS Topic ARN for email notifications
            default: ""
          S3BucketName:
            type: String
            description: Name of the S3 bucket containing the shell script and related log files can be stored
            default: ""
          S3ScriptKey:
            type: String
            description: The S3 key (path) to the shell script
            default: rds_psql_patch.sh
          DBTagValue:
            type: String
            description: The tag value to identify RDS PostgreSQL instances
            default: 'Y'
          DBTagKey:
            type: String
            description: The tag key to identify RDS PostgreSQL instances
            default: UpgradeDB
          TargetEngineVersion:
            type: String
            description: Target Engine Version (Sample value 15.8)
        assumeRole: '{{ AssumeIAMRole }}'
        mainSteps:
          - name: DownloadAndSaveScript
            action: aws:runCommand
            nextStep: UpgradeInstances
            isEnd: false
            onFailure: Abort
            inputs:
              DocumentName: AWS-RunShellScript
              InstanceIds:
                - '{{ EC2InstanceId }}'
              Parameters:
                commands:
                  - |
                    #!/bin/bash
                    set -euo pipefail
                    # e: exits on error
                    # u: exits on undefined variables
                    # o pipefail: exits if any command in a pipe fails

                    # Debug information
                    # echo "Current user: $(whoami)"
                    # echo "Home directory: $HOME"

                    # Setup environment
                    export AWS_REGION="{{ global:REGION }}"
                    export AWS_DEFAULT_REGION="{{ global:REGION }}"

                    # Create directory
                    mkdir -p {{ ScriptsDirOnEC2 }}

                    # Download file
                    aws s3 cp s3://{{ S3BucketName }}/{{ S3ScriptKey }} {{ ScriptsDirOnEC2 }}/rds_psql_patch.sh

                    # Make script executable
                    chmod +x {{ ScriptsDirOnEC2 }}/rds_psql_patch.sh

                    # Verify final state
                    echo "Final script permissions:"
                    ls -la {{ ScriptsDirOnEC2 }}/rds_psql_patch.sh

          - name: UpgradeInstances
            action: aws:runCommand
            isEnd: true
            onFailure: Abort
            inputs:
              DocumentName: AWS-RunShellScript
              InstanceIds:
                - '{{ EC2InstanceId }}'
              Parameters:
                commands:
                  - |
                    #!/bin/bash
                    set -euo pipefail
                    # e: exits on error
                    # u: exits on undefined variables
                    # o pipefail: exits if any command in a pipe fails


                    # Debug information
                    # echo "Current user: $(whoami)"
                    # echo "Home directory: $HOME"

                    # Setup environment
                    export AWS_REGION="{{ global:REGION }}"
                    export AWS_DEFAULT_REGION="{{ global:REGION }}"

                    check_aws_cli_version() {
                        # Check if aws command exists
                        if ! command -v aws &> /dev/null; then
                            echo "Error: AWS CLI is not installed"
                            exit 1
                        fi

                        # Get version string
                        VERSION=$(aws --version 2>&1)
                        
                        # Extract major version number
                        if [[ $VERSION =~ aws-cli/([0-9]+)\. ]]; then
                            MAJOR_VERSION="${BASH_REMATCH[1]}"
                            
                            if [ "$MAJOR_VERSION" -lt 2 ]; then
                                echo "Error: AWS CLI version 2 or greater is required. Current version: $VERSION"
                                exit 1
                            else
                                echo "AWS CLI version check passed: $VERSION"
                            fi
                        else
                            echo "Error: Could not determine AWS CLI version from: $VERSION"
                            exit 1
                        fi
                    }

                    # Run the version check
                    check_aws_cli_version

                    # Create logs directory with timestamp suffix for master log
                    TIMESTAMP=$(date '+%Y%m%d-%H-%M-%S')
                    MASTER_LOG="{{ ScriptsDirOnEC2 }}/logs/upgrade_master-${TIMESTAMP}.log"
                    mkdir -p {{ ScriptsDirOnEC2 }}/logs

                    log_message() {
                      echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "${MASTER_LOG}"
                    }

                    export TAG_KEY="{{ DBTagKey }}"
                    export TAG_VALUE="{{ DBTagValue }}"
                    INSTANCES=$(aws rds describe-db-instances --filters "Name=engine,Values=postgres" --output json | \
                      jq -r --arg key "$TAG_KEY" --arg value "$TAG_VALUE" \
                        '.DBInstances[] | select(.TagList[] | select(.Key==$key and .Value==$value)) | .DBInstanceIdentifier'
                    )
                    log_message "Found instances to upgrade: ${INSTANCES}"

                    COUNT=0
                    declare -a PIDS=()

                    cleanup() {
                      log_message "Cleaning up background processes..."
                      for pid in "${PIDS[@]}"; do
                        kill -9 $pid 2>/dev/null || true
                      done
                      exit 1
                    }

                    trap cleanup SIGINT SIGTERM

                    export SCRIPTS_DIR={{ ScriptsDirOnEC2 }}
                    export S3_BUCKET_PATCH_LOGS={{ S3BucketName }}
                    export SNS_TOPIC_ARN_EMAIL={{ SnsTopicArnEmail }}

                    for INSTANCE in $INSTANCES; do
                      # Create instance-specific log directory
                      INSTANCE_LOG_DIR="{{ ScriptsDirOnEC2 }}/logs/${INSTANCE}"
                      mkdir -p "${INSTANCE_LOG_DIR}"
                      INSTANCE_LOG="${INSTANCE_LOG_DIR}/upgrade-${TIMESTAMP}.log"

                      (
                        cd {{ ScriptsDirOnEC2 }}
                        log_message "Starting upgrade for instance: ${INSTANCE}" >> "$INSTANCE_LOG"
                        if ./rds_psql_patch.sh "$INSTANCE" {{ TargetEngineVersion }} {{ RunPreUpgradeTasks }} >> "$INSTANCE_LOG" 2>&1; then
                          log_message "Completed upgrade for instance: ${INSTANCE}" >> "$INSTANCE_LOG"
                        else
                          log_message "Failed upgrade for instance: ${INSTANCE}" >> "$INSTANCE_LOG"
                          exit 1
                        fi
                      ) &
                      PIDS+=($!)
                      COUNT=$((COUNT+1))
                      log_message "Started upgrade process for ${INSTANCE} with PID ${PIDS[-1]}"
                    done

                    EXIT_STATUS=0
                    FAILED_INSTANCES=""

                    for pid in "${PIDS[@]}"; do
                      if ! wait $pid; then
                        EXIT_STATUS=1
                        FAILED_INSTANCES="${FAILED_INSTANCES} $(ps -p $pid -o args= 2>/dev/null | grep -o 'instance_[^ ]*' || echo 'unknown')"
                      fi
                    done

                    log_message "Upgrade initiated for $COUNT instances"

                    if [ $EXIT_STATUS -eq 0 ]; then
                      log_message "All upgrades completed successfully"
                    else
                      log_message "Some upgrades failed. Failed instances:${FAILED_INSTANCES}"
                    fi

                    if [ -n "{{ S3BucketName }}" ]; then
                      log_message "Uploading logs to S3..."
                      aws s3 cp {{ ScriptsDirOnEC2 }}/logs s3://{{ S3BucketName }}/logs/$(date '+%Y-%m-%d')/ --recursive
                    fi

                    exit $EXIT_STATUS

Outputs:
  SSMDocumentName:
    Description: Name of the SSM Automation Document
    Value: !Ref RDSPostgreSQLFleetUpgrade
  SSMDocumentAssumeRoleArn:
    Description: The ARN of the role that allows Automation to perform the actions on your behalf.
    Value: !GetAtt AutomationServiceRole.Arn