# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
# The following code provides the sample database schema used in this post. 
# Create these objects in the Aurora PostgreSQL database you plan to use for the benchmark test.
# CREATE SEQUENCE if not exists order_seq
# INCREMENT 1
# START 1
# MINVALUE 1
# MAXVALUE 9223372036854775807
# CACHE 1;
#
# CREATE SEQUENCE if not exists order_item_seq
# INCREMENT 1
# START 1
# MINVALUE 1
# MAXVALUE 9223372036854775807
# CACHE 1;
#
# CREATE SEQUENCE if not exists order_payment_seq
# INCREMENT 1
# START 1
# MINVALUE 1
# MAXVALUE 9223372036854775807
# CACHE 1;
#
# CREATE TABLE customer_order (
# order_id bigint NOT NULL,
# order_description varchar(256) NOT NULL,
# order_date timestamp(0) without time zone NOT NULL,
# CONSTRAINT order_pkey PRIMARY KEY (order_id)
# );
#
# CREATE TABLE order_item (
# order_id bigint NOT NULL references customer_order (order_id),
# order_item_id bigint NOT NULL DEFAULT nextval('order_item_seq'),
# item_description varchar(2048) NOT NULL,
# quantity integer NOT NULL,
# price numeric(12,2) NOT NULL,
# CONSTRAINT order_item_pkey PRIMARY KEY (order_id,order_item_id)
# );
#
# CREATE TABLE order_payment (
# order_id bigint NOT NULL references customer_order (order_id),
# payment_id bigint NOT NULL DEFAULT nextval('order_payment_seq'),
# amount numeric(12,2) NOT NULL ,
# payment_date timestamp(0) without time zone,
# CONSTRAINT order_payment_pkey PRIMARY KEY (order_id,payment_id)
# );
---
AWSTemplateFormatVersion: '2010-09-09'

Description: |
  # Auotmation for the blog post https://aws.amazon.com/blogs/database/automate-benchmark-tests-for-amazon-aurora-postgresql/
  # Run a benchmark across Amazon Elastic Compute Cloud (Amazon EC2) instances.
  # Multiple benchmarks may run concurrently. 
  # One benchmark performs select, update, insert, and delete. 
  # At the same time, another benchmark performs select statements. 
  # **WARNING** This template creates
      AWS::EC2::LaunchTemplate
      AWS::IAM::Role
      AWS::IAM::InstanceProfile 
      AWS::AutoScaling::AutoScalingGroup 
      AWS::SSM::Document
  # You will be billed for the AWS resources used if you create a stack from this template.  
Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
    - Label:
        default: Database Parameters
      Parameters:
      - BenchmarkDB
      - BenchmarkSecretArn  
      - BenchmarkReadHost  
    - Label:
        default: Launch EC2 Benchmark Parameters
      Parameters:
      - InstanceType
      - KeyName
      - SecurityGroupId
      - SubnetID
      - NumberOfReaders
      - NumberOfWriters
    - Label:
        default: Benchmark EC2 Configuration
      Parameters:
      - BenchmarkConnections
      - BenchmarkThreads
      - BenchmarkTime
    ParameterLabels:
      BenchmarkDB:
        default: Database Name
      BenchmarkSecretArn:
        default: Database Secrets Manager Secret ARN
      BenchmarkReadHost:
        default: Database Reader Host Endpoint        
      BenchmarkConnections:
        default: Number of Client Connections to Database
      BenchmarkThreads:
        default: Number of Threads per Connection
      BenchmarkTime:
        default: Benchmark Duration
      InstanceType:
        default: EC2 Instance Type
      KeyName:
        default: EC2 Instance Security Key Name
      SecurityGroupId:
        default: Security Group ID for the Benchmark EC2 Instances
      SubnetID:
        default: SubnetID for the Benchmark EC2 instances
      NumberOfReaders:
        default: Number of Benchmark Instances for Reads
      NumberOfWriters:
        default: Number of Benchmark Instances for Writes
Parameters:
  InstanceType:
    Type: String
    Description: Enter EC2 Instance Type. Default is t3.micro.
    AllowedValues:
      - t3.micro
      - t3.small
      - t3.medium
      - t3.large
    Default: t3.micro
  KeyName:
    Type: AWS::EC2::KeyPair::KeyName
    Description: Amazon EC2 Key Pair
  SecurityGroupId:
    Type: AWS::EC2::SecurityGroup::Id
    Description: Enter Security Group ID for the EC2 instance
  SubnetID:
    Type: AWS::EC2::Subnet::Id
    Description: Enter the VPC Subnet ID to launch the EC2 instances.
  NumberOfReaders:
    Type: String
    Description: Enter the number of EC2 instances to launch for read benchmark. Default is 1.
    MinLength: 1
    AllowedValues:
      - '1'
      - '2'
      - '3'
      - '4'
      - '5'
      - '6'
      - '7'
      - '8'
      - '9'
      - '10'
    Default: '1'
  NumberOfWriters:
    Type: String
    Description: Enter the number of EC2 instances to launch for write benchmark. Default is 1.  
    MinLength: 1
    AllowedValues:
      - '1'
      - '2'
      - '3'
      - '4'
      - '5'
      - '6'
      - '7'
      - '8'
      - '9'
      - '10'
    Default: '1'
  BenchmarkDB:
    Type: String
    Description: Enter the database name for benchmark. 
    MinLength: 1
    Default: 'postgres'
  BenchmarkSecretArn:
    Type: String
    Description: Enter Secrets Manager secret ARN used for database credentials. The host, port, username, and password are retrieved from secret. 
    MinLength: 1   
  BenchmarkReadHost:
    Type: String
    Description: Enter the database reader endpoint.  If not provided, defaults to the writer host defined in Secrets Manager secret. 
  BenchmarkConnections:
    Type: Number
    Description: Enter number of client connections for benchmark. Default is 16.
    MinValue: 1
    MaxValue: 5000
    Default: '16'
  BenchmarkThreads:
    Type: Number
    Description: Enter number of threads for benchmark. Default is 8. 
    Default: '8'
    MinValue: 1
    MaxValue: 32
  BenchmarkTime:
    Type: Number
    Description: Enter time in seconds to run the benchmark. Default is 300.
    Default: 300
    MinValue: 1
    MaxValue: 86400
Rules:
  KeyNameRule:
    Assertions:
      - Assert:
          'Fn::Not':
            - 'Fn::Equals':
              - !Ref KeyName
              - ""
        AssertDescription: 'EC2 Instance Security Key Name is a required parameter'
  SubnetRule:
    Assertions:
      - Assert:
          'Fn::Not':
            - 'Fn::Equals':
              - !Ref SubnetID
              - ""
        AssertDescription: 'SubnetID is a required parameter'
  SecurityGroupRule:
    Assertions:
      - Assert:
          'Fn::Not':
            - 'Fn::Equals':
              - !Ref SecurityGroupId
              - ""
        AssertDescription: 'Security group ID is a required parameter'        
Resources:
  EC2LaunchTemplate:
    # The LaunchTemplate is used for the ec2 instances to run benchmark
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateData:
        UserData:
          Fn::Base64:
            Fn::Join:
            - "\n"
            - - "#!/bin/bash"
              - sudo yum -y update
              - sudo amazon-linux-extras enable postgresql12
              - sudo yum clean metadata
              - sudo yum install postgresql -y
              - sudo yum install postgresql-contrib -y
              - sudo yum install jq -y
              - mkdir /home/ec2-user/benchmark
              - cd /home/ec2-user/benchmark
              - touch benchmark.sh
              - echo "#!/bin/bash" > benchmark.sh
              - echo "# SYNTAX FOR RUNNING FOR A SPECIFIC AMOUNT OF TIME IN SECONDS"
                >> benchmark.sh
              - echo "pgbench -c \$BenchmarkConnections -j \$BenchmarkThreads -T \$BenchmarkTime
                -U \$BenchmarkUser -d \$BenchmarkDB -h \$BenchmarkHost -p \$BenchmarkPort
                -f \$BenchmarkSqlFile -n" >> benchmark.sh
              - chmod u+x benchmark.sh
              - touch write.sql
              - echo "\set order_item_quantity random(1,100)" >> write.sql
              - echo "\set order_item_price double(random(1,1000))" >>  write.sql
              - echo "\set order_payment double(random(1, :order_item_price))" >>  write.sql
              - echo "\set order_update_id random(1,10000000)" >>  write.sql
              - echo "\set order_delete_id random(1,10000000)" >>  write.sql
              - echo "BEGIN;">> write.sql
              - echo "-- get a unique number for customer order" >> write.sql
              - echo "SELECT nextval('order_seq') \gset p_" >> write.sql
              - echo "-- insert customer_order" >> write.sql
              - echo "INSERT INTO customer_order(order_id, order_description, order_date)
                VALUES (:p_nextval, concat('description for order',:p_nextval), CURRENT_TIMESTAMP);"
                >> write.sql
              - echo "-- insert order item" >> write.sql
              - echo "INSERT INTO order_item(order_id, item_description, quantity,
                price) VALUES (:p_nextval, concat('item count is ',:order_item_quantity)
                ,:order_item_quantity, :order_item_price);" >> write.sql
              - echo "-- insert order payment" >> write.sql
              - echo "INSERT INTO order_payment(order_id, amount, payment_date) VALUES
                (:p_nextval, :order_payment, CURRENT_TIMESTAMP);" >> write.sql
              - echo "-- update random order item" >> write.sql
              - echo "UPDATE order_item set item_description = concat(item_description,'
                random update') WHERE order_id = :order_update_id;" >> write.sql
              - echo "-- delete random order item" >> write.sql
              - echo "DELETE FROM order_item WHERE order_id = :order_delete_id;" >>
                write.sql
              - echo "END;" >> write.sql
              - touch read.sql
              - 'echo "\set order_id random(1,10000000)" > read.sql '
              - echo "BEGIN;" >> read.sql
              - 'echo "-- select random orders" >> read.sql '
              - 'echo "SELECT SUM(order_item.PRICE * order_item.quantity) FROM customer_order
                , order_item WHERE customer_order.order_id = :order_id AND customer_order.order_id
                = order_item.order_id;" >> read.sql '
              - 'echo "--select random order payment" >> read.sql  '
              - echo "SELECT customer_order.order_id, order_item.item_description,
                order_item.quantity, order_payment.amount FROM customer_order , order_item
                , order_payment WHERE customer_order.order_id = :order_id AND order_item.order_id
                = customer_order.order_id AND order_payment.order_id = customer_order.order_id;"
                >> read.sql
              - 'echo "END;" >> read.sql '
        ImageId: '{{resolve:ssm:/aws/service/ami-amazon-linux-latest/amzn2-ami-hvm-x86_64-gp2}}'    
        SecurityGroupIds:
        - Ref: SecurityGroupId
        KeyName:
          Ref: KeyName
        InstanceType:
          Ref: InstanceType
        IamInstanceProfile:
          Arn:
            Fn::GetAtt:
            - MyInstanceProfile
            - Arn
      LaunchTemplateName: !Sub 'EC2PostgresSQLBenchmark-${AWS::StackName}' 
  EC2AttachedRole:
    # The role allows AWS Systems Manager to use the ec2 instances
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - ec2.amazonaws.com
          Action:
          - sts:AssumeRole
      Description: EC2 Role with permissions to access AWS Systems Manager
      ManagedPolicyArns:
      - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
      - arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess
      Policies:
      - 
        PolicyName: EmbeddedReadSecretPolicy
        PolicyDocument:
          Version: 2012-10-17
          Id: SECRET_READ_POLICY
          Statement:
            - Sid: SecretReadPermission
              Effect: Allow              
              Action: 
                - 'secretsmanager:GetResourcePolicy'
                - 'secretsmanager:GetSecretValue'
                - 'secretsmanager:DescribeSecret'
                - 'secretsmanager:ListSecretVersionIds'
              Resource: !Sub '${BenchmarkSecretArn}'      
  MyInstanceProfile:
    # The InstanceProfile uses the EC2AttachedRole and is assigned to the benchmark ec2 instances
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
      - Ref: EC2AttachedRole
  EC2ScalingGroupReaders:
    # Used to scale the number of ec2 instances used for read operations
    Type: AWS::AutoScaling::AutoScalingGroup
    Properties:
      MaxSize: '10'
      MinSize: '1'
      DesiredCapacity:
        Ref: NumberOfReaders
      LaunchTemplate:
        LaunchTemplateId:
          Ref: EC2LaunchTemplate
        Version:
          Fn::GetAtt:
          - EC2LaunchTemplate
          - LatestVersionNumber
      VPCZoneIdentifier:
      - Ref: SubnetID
      Tags:
      - Key: !Sub 'benchmark-${AWS::StackName}' 
        PropagateAtLaunch: 'true'
        Value: read
  EC2ScalingGroupWriters:
    # Used to scale the number of ec2 instances used for write operations
    Type: AWS::AutoScaling::AutoScalingGroup
    Properties:
      MaxSize: '10'
      MinSize: '1'
      DesiredCapacity:
        Ref: NumberOfWriters
      LaunchTemplate:
        LaunchTemplateId:
          Ref: EC2LaunchTemplate
        Version:
          Fn::GetAtt:
          - EC2LaunchTemplate
          - LatestVersionNumber
      VPCZoneIdentifier:
      - Ref: SubnetID
      Tags:
      - Key: !Sub 'benchmark-${AWS::StackName}' 
        PropagateAtLaunch: 'true'
        Value: write
  SSMDocument:
    # SSM uses the SSMDocument to run the benchmark commands on the benchmark ec2 instances
    Type: AWS::SSM::Document
    Properties:
      Content:
        schemaVersion: '2.2'
        description: Run a script on Linux instances.
        parameters:
          duration:
            type: String        
            description: "Benchmark duration in seconds (1-86400)"
            allowedPattern: "^([1-9]|[1-9][0-9]|[1-9][0-9][0-9]|[1-9][0-9][0-9][0-9]|[1-7][0-9][0-9][0-9][0-9]|[8][0-6][0-3][0-9][0-9]|[8][6][4][0][0])$"
            default:  !Sub "${BenchmarkTime}"
          connections:
            type: String        
            description: "Number of connections to database (1-5000)"
            allowedPattern: "^([1-9]|[1-9][0-9]|[1-9][0-9][0-9]|[1-4][0-9][0-9][0-9]|[5][0-0][0-0][0-0])$"
            default:  !Sub "${BenchmarkConnections}"
          threads:
            type: String        
            description: "Number of threads per client connection (1-32)"
            allowedPattern: "^([1-9]|[0-2][0-9]|[3][0-2]|)$"
            default:  !Sub "${BenchmarkThreads}"
        mainSteps:
        - action: aws:runShellScript          
          name: runCommands
          inputs:            
            timeoutSeconds: '86400'
            runCommand:
              -
                Fn::Join:
                - "\n"
                - - "# PostgreSQL connection environment variables"
                  - yum update -y aws-cfn-bootstrap
                  - Fn::Sub: aws configure set default.region ${AWS::Region}
                  - Fn::Sub: export BenchmarkDB=${BenchmarkDB}
                  - Fn::Sub: export BenchmarkHost=$(aws secretsmanager get-secret-value --secret-id ${BenchmarkSecretArn} --region ${AWS::Region} | jq -rc .SecretString | jq -rc .host)
                  - Fn::Sub: export BenchmarkReadHost=${BenchmarkReadHost}
                  - Fn::Sub: export BenchmarkPort=$(aws secretsmanager get-secret-value --secret-id ${BenchmarkSecretArn} --region ${AWS::Region} | jq -rc .SecretString | jq -rc .port)
                  - Fn::Sub: export PGPASSWORD=$(aws secretsmanager get-secret-value --secret-id ${BenchmarkSecretArn} --region ${AWS::Region} | jq -rc .SecretString | jq -rc .password)
                  - Fn::Sub: export BenchmarkUser=$(aws secretsmanager get-secret-value --secret-id ${BenchmarkSecretArn} --region ${AWS::Region} | jq -rc .SecretString | jq -rc .username)
                  - "# pgbench control environment variables"
                  - Fn::Sub: export BenchmarkConnections={{connections}}
                  - Fn::Sub: export BenchmarkThreads={{threads}}
                  - Fn::Sub: export BenchmarkTime={{duration}}
                  - export instanceid=$(curl http://169.254.169.254/latest/meta-data/instance-id)
                  - Fn::Sub: export BenchmarkSqlFile=$(aws ec2 describe-tags --region ${AWS::Region} --filters "Name=key,Values=benchmark-${AWS::StackName}" "Name=resource-id,Values=$instanceid" --query 'Tags[0].Value' | tr -d '"').sql
                  - "# time in seconds to run benchmark"
                  - cd /home/ec2-user/benchmark/   
                  - if [ -z "$BenchmarkReadHost" ]; then export BenchmarkReadHost=$BenchmarkHost; fi;             
                  - if [ "$BenchmarkSqlFile" = "read.sql" ]; then export BenchmarkHost=$BenchmarkReadHost; fi;
                  - echo "pgbench -c $BenchmarkConnections -j $BenchmarkThreads -T $BenchmarkTime -U $BenchmarkUser -d $BenchmarkDB -h $BenchmarkHost -p $BenchmarkPort -f $BenchmarkSqlFile -n"
                  - pgbench -c $BenchmarkConnections -j $BenchmarkThreads -T $BenchmarkTime
                    -U $BenchmarkUser -d $BenchmarkDB -h $BenchmarkHost -p $BenchmarkPort
                    -f $BenchmarkSqlFile -n
      DocumentType: Command
      Name: !Sub 'PostgreSQLBenchmarkCommandDocument-${AWS::StackName}' 
Outputs:
  BenchmarkReadTag:
    Description: "EC2 instances generating read traffic use this tag"
    Value: !Sub "key=benchmark-${AWS::StackName} value=read"
  BenchmarkWriteTag:
    Description: "EC2 instances generating write traffic use this tag"
    Value: !Sub "key=benchmark-${AWS::StackName} value=write"
