---
AWSTemplateFormatVersion: 2010-09-09

Description: >
    This template deploys a VPC and a Amazon SageMaker Notebook Instance

Parameters:
  VpcName:
    Default: EC2SpotSageMakerWorkshop
    Type: String

  VpcCIDR:
    Default: 10.215.0.0/16
    Type: String

  Subnet1CIDR:
    Default: 10.215.10.0/24
    Type: String

  Subnet2CIDR:
    Default: 10.215.20.0/24
    Type: String

  Subnet3CIDR:
    Default: 10.215.30.0/24
    Type: String

  Subnet4CIDR:
    Default: 10.215.40.0/24
    Type: String

  DefaultCodeRepository:
    Default: https://github.com/aws-samples/rds-postgresql-pgvector.git
    Type: String

  DBEngineVersion:
    Type: String
    Default: 15.2
    AllowedValues:
      - 15.2
      - 14.7
      - 13.10

Resources:

# VPC ----------------------------------------------------------
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: !Ref VpcCIDR
      Tags:
        - Key: Name
          Value: !Ref VpcName

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: !Ref VpcName

  InternetGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      InternetGatewayId: !Ref InternetGateway
      VpcId: !Ref VPC

  Subnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 0, !GetAZs ]
      MapPublicIpOnLaunch: true
      CidrBlock: !Ref Subnet1CIDR
      Tags:
        - Key: Name
          Value: !Sub ${VpcName} (Public)

  Subnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 1, !GetAZs ]
      MapPublicIpOnLaunch: true
      CidrBlock: !Ref Subnet2CIDR
      Tags:
        - Key: Name
          Value: !Sub ${VpcName} (Public)

  Subnet3:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 0, !GetAZs ]
      MapPublicIpOnLaunch: false
      CidrBlock: !Ref Subnet3CIDR
      Tags:
        - Key: Name
          Value: !Sub ${VpcName} (Private)

  Subnet4:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 1, !GetAZs ]
      MapPublicIpOnLaunch: false
      CidrBlock: !Ref Subnet4CIDR
      Tags:
        - Key: Name
          Value: !Sub ${VpcName} (Private)

  RouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Ref VpcName

  DefaultRoute:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref RouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway

  Subnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref RouteTable
      SubnetId: !Ref Subnet1

  Subnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref RouteTable
      SubnetId: !Ref Subnet2

  Subnet3RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref RouteTable
      SubnetId: !Ref Subnet3

  Subnet4RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref RouteTable
      SubnetId: !Ref Subnet4


# END VPC ------------------------------------------------------

# NOTEBOOK -----------------------------------------------------
  NotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      InstanceType: "ml.t2.large"
      RoleArn: !GetAtt ExecutionRole.Arn
      SubnetId: !Ref Subnet1
      SecurityGroupIds:
        - !Ref SecurityGroup
      DefaultCodeRepository: !Ref DefaultCodeRepository
      VolumeSizeInGB: 20
  
  SecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Notebook Instance Security Group
      VpcId: !Ref VPC
      SecurityGroupEgress:
        - IpProtocol: "-1"
          CidrIp: 0.0.0.0/0
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 5432
          ToPort: 5432
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 443
          ToPort: 443
          CidrIp: 0.0.0.0/0

  ExecutionRole: 
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
          - 
            Effect: "Allow"
            Principal: 
              Service: 
                - "sagemaker.amazonaws.com"
                - "translate.amazonaws.com"
            Action: 
              - "sts:AssumeRole"
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/TranslateFullAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AdministratorAccess
      Policies:
        - 
          PolicyName: "s3_access"
          PolicyDocument: 
            Version: "2012-10-17"
            Statement: 
              - 
                Effect: "Allow"
                Action: 
                  - "s3:PutBucketPolicy"
                  - "s3:DeleteBucket"
                Resource: "arn:aws:s3:::sagemaker-*"                

# END NOTEBOOK -------------------------------------------------

# RDS PostgreSQL -----------------------------------------------

  EncryptionKey:
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain
    Type: AWS::KMS::Key
    Properties:
      EnableKeyRotation: true
      KeyPolicy:
        Version: 2012-10-17
        Id: !Ref AWS::StackName
        Statement:
          - Effect: Allow
            Principal:
              AWS:
                - !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Action: "kms:*"
            Resource: "*"
      Tags:
        - Key: Name
          Value: !Ref AWS::StackName

  EncryptionKeyAlias:
    Type: AWS::KMS::Alias
    Properties:
      AliasName: !Sub "alias/${AWS::StackName}"
      TargetKeyId: !Ref EncryptionKey

  DBSubnetGroup:
    Type: AWS::RDS::DBSubnetGroup
    Properties:
      DBSubnetGroupDescription: "RDS DB Subnet Group"
      SubnetIds: [!Ref Subnet3, !Ref Subnet4 ]

        # shared_preload_libraries: "vector, pg_stat_statements"
  RDSDBParameterGroup:
    Type: AWS::RDS::DBParameterGroup
    Properties:
      Description: "RDS PostgreSQL Custom Cluster parameter group"
      Family: postgres15
      Parameters:
        shared_preload_libraries: "pg_stat_statements"
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}"

  RDSSecrets:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: "rdspg-vector-secret"
      Description: 'This is the secret for RDS instance'
      GenerateSecretString:
        SecretStringTemplate: '{"username": "postgres" }'
        GenerateStringKey: 'password'
        PasswordLength: 16
        ExcludeCharacters: '"@/\'

  VPCSecurityGroup:
    Type: 'AWS::EC2::SecurityGroup'
    Properties:
      GroupDescription: !Ref 'AWS::StackName'
      SecurityGroupEgress:
      - IpProtocol: -1
        CidrIp: "0.0.0.0/0"
      SecurityGroupIngress:
      - IpProtocol: tcp
        FromPort: 5432
        ToPort: 5432
        CidrIp: !Ref VpcCIDR
        Description: 'Access to AppServer Host Security Group'
      VpcId: !Ref VPC
      Tags:
      - Key: Name
        Value: !Sub '${AWS::StackName}-DBSecurityGroup'

  DBInstance:
    Type: AWS::RDS::DBInstance
    Properties:
      DBInstanceIdentifier: "rdspg-vector"
      AllocatedStorage: "20"
      DBInstanceClass: db.t3.micro
      Engine: postgres
      MasterUsername: !Join ['', ['{{resolve:secretsmanager:', !Ref RDSSecrets, ':SecretString:username}}' ]]
      MasterUserPassword: !Join ['', ['{{resolve:secretsmanager:', !Ref RDSSecrets, ':SecretString:password}}' ]]
      DBParameterGroupName: !Ref RDSDBParameterGroup
      DBSubnetGroupName: !Ref DBSubnetGroup
      VPCSecurityGroups: 
        - !Ref VPCSecurityGroup
      AllowMajorVersionUpgrade: false
      AutoMinorVersionUpgrade: true
      EngineVersion: "15.2"
      KmsKeyId: !Ref EncryptionKey
      MultiAZ: false
      StorageType: gp3
      StorageEncrypted: true
      BackupRetentionPeriod: 0
      DeletionProtection: false
      PubliclyAccessible: false
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}"

  SecretPostgreSQLAttachment:
    Type: AWS::SecretsManager::SecretTargetAttachment
    Properties:
      SecretId: !Ref RDSSecrets
      TargetId: !Ref DBInstance
      TargetType: AWS::RDS::DBInstance

# end RDS PostgreSQL -----------------------------------------------

Outputs:

  NotebookInstanceURL:
    Description: SageMaker Notebook Instance URL
    Value: !Join
      - ''
      - - !Sub 'https://console.aws.amazon.com/sagemaker/home?region=${AWS::Region}#/notebook-instances/openNotebook/'
        - !GetAtt NotebookInstance.NotebookInstanceName
        - '?view=classic'

  DBEndpoint:
    Description: 'RDS PostgreSQL Endpoint'
    Value: !GetAtt 'DBInstance.Endpoint.Address'
    Export:
      Name:
        'Fn::Sub': '${AWS::StackName}-DBEndPoint'

  DBSecret:
    Description: Database Secret'
    Value: !Ref RDSSecrets
    Export:
      Name:
        'Fn::Sub': '${AWS::StackName}-DBSecrets'