Parameters:
  SageMakerUserProfileName:
    Type: String
    Description: The user profile name for the SageMaker workshop
    Default: SageMakerDemoUser
  SageMakerDomainName:
    Type: String
    Description: The domain name of the Sagemaker studio instance
    Default: SageMakerDemoDomain
  RedshiftSecret:
    Description: Please enter Secret name for redshift database
    Type: String
    Default: SecretForRedshiftSageMakerDemo2023
  PrivateSubnetCIDR:
    Description: Please enter IPv4 CIDR for Private Subnet
    Type: String
    Default: 172.31.96.0/24
  glueConnectionName:
    Description: Please enter Glue connection name for Redshift DB
    Type: String
    Default: RS-SM-Demo-glue-Connection
Mappings:
  RegionMap:
    us-east-1:
      datascience: 'arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-east-1:663277389841:image/sagemaker-data-wrangler-1.0
    us-east-2:
      datascience: 'arn:aws:sagemaker:us-east-2:429704687514:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-east-2:415577184552:image/sagemaker-data-wrangler-1.0
    us-west-1:
      datascience: 'arn:aws:sagemaker:us-west-1:742091327244:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-west-1:926135532090:image/sagemaker-data-wrangler-1.0
    us-west-2:
      datascience: 'arn:aws:sagemaker:us-west-2:236514542706:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-west-2:174368400705:image/sagemaker-data-wrangler-1.0
    af-south-1:
      datascience: 'arn:aws:sagemaker:af-south-1:559312083959:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:af-south-1:143210264188:image/sagemaker-data-wrangler-1.0
    ap-east-1:
      datascience: 'arn:aws:sagemaker:ap-east-1:493642496378:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-east-1:707077482487:image/sagemaker-data-wrangler-1.0
    ap-south-1:
      datascience: 'arn:aws:sagemaker:ap-south-1:394103062818:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-south-1:089933028263:image/sagemaker-data-wrangler-1.0
    ap-northeast-2:
      datascience: 'arn:aws:sagemaker:ap-northeast-2:806072073708:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-northeast-2:131546521161:image/sagemaker-data-wrangler-1.0
    ap-southeast-1:
      datascience: 'arn:aws:sagemaker:ap-southeast-1:492261229750:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-southeast-1:119527597002:image/sagemaker-data-wrangler-1.0
    ap-southeast-2:
      datascience: 'arn:aws:sagemaker:ap-southeast-2:452832661640:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-southeast-2:422173101802:image/sagemaker-data-wrangler-1.0
    ap-northeast-1:
      datascience: 'arn:aws:sagemaker:ap-northeast-1:102112518831:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-northeast-1:649008135260:image/sagemaker-data-wrangler-1.0
    ca-central-1:
      datascience: 'arn:aws:sagemaker:ca-central-1:310906938811:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ca-central-1:557239378090:image/sagemaker-data-wrangler-1.0
    eu-central-1:
      datascience: 'arn:aws:sagemaker:eu-central-1:936697816551:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-central-1:024640144536:image/sagemaker-data-wrangler-1.0
    eu-west-1:
      datascience: 'arn:aws:sagemaker:eu-west-1:470317259841:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-1:245179582081:image/sagemaker-data-wrangler-1.0
    eu-west-2:
      datascience: 'arn:aws:sagemaker:eu-west-2:712779665605:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-2:894491911112:image/sagemaker-data-wrangler-1.0
    eu-west-3:
      datascience: 'arn:aws:sagemaker:eu-west-3:615547856133:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-3:807237891255:image/sagemaker-data-wrangler-1.0
    eu-north-1:
      datascience: 'arn:aws:sagemaker:eu-north-1:243637512696:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-north-1:054986407534:image/sagemaker-data-wrangler-1.0
    eu-south-1:
      datascience: >-
        arn:aws:sagemaker:eu-south-1:488287956546:image/sagemaker-data-wrangler-1.0
      datawrangler: 'arn:aws:sagemaker:eu-south-1:592751261982:image/datascience-1.0'
    sa-east-1:
      datascience: 'arn:aws:sagemaker:sa-east-1:782484402741:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:sa-east-1:424196993095:image/sagemaker-data-wrangler-1.0
Resources:
  SageMakerDemoSG:
    Type: 'AWS::EC2::SecurityGroup'
    Properties:
      VpcId: !GetAtt 
        - DefaultVpcFinder
        - VpcId
      GroupDescription: !Sub Default Security group for demo
  ConsumerSGIngress:
    Type: 'AWS::EC2::SecurityGroupIngress'
    DependsOn: SageMakerDemoSG
    Properties:
      GroupId: !Ref SageMakerDemoSG
      IpProtocol: -1
      SourceSecurityGroupId: !Ref SageMakerDemoSG
  ConsumerSGEgress:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      IpProtocol: -1
      GroupId: !Ref SageMakerDemoSG
      CidrIp: 0.0.0.0/0
      
  PrivateSubnet:
    Type: 'AWS::EC2::Subnet'
    Properties:
      VpcId: !GetAtt 
        - DefaultVpcFinder
        - VpcId
      AvailabilityZone: !Select 
        - 0
        - !GetAZs ''
      CidrBlock: !Ref PrivateSubnetCIDR
      Tags:
        - Key: Name
          Value: Private Subnet for Redshift/Sagemaker
  NatGatewayEIP:
    Type: 'AWS::EC2::EIP'
    Properties:
      Domain: vpc
  NatGateway:
    Type: 'AWS::EC2::NatGateway'
    Properties:
      AllocationId: !GetAtt 
        - NatGatewayEIP
        - AllocationId
      SubnetId: !Select 
        - 0
        - !GetAtt 
          - DefaultVpcFinder
          - Subnets
  PrivateRouteTable:
    Type: 'AWS::EC2::RouteTable'
    Properties:
      VpcId: !GetAtt 
        - DefaultVpcFinder
        - VpcId
      Tags:
        - Key: Name
          Value: Private RouteTable
  DefaultPrivateRoute:
    Type: 'AWS::EC2::Route'
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NatGateway
  PrivateSubnetRouteTableAssociation:
    Type: 'AWS::EC2::SubnetRouteTableAssociation'
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      SubnetId: !Ref PrivateSubnet
  DemoRedshiftSecret:
    Type: 'AWS::SecretsManager::Secret'
    Properties:
      Name: !Ref RedshiftSecret
      Description: This is my redshift instance secret
      GenerateSecretString:
        SecretStringTemplate: '{"username": "admin"}'
        GenerateStringKey: password
        PasswordLength: 16
        ExcludePunctuation: true
  DemoRedshiftCluster:
    Type: 'AWS::Redshift::Cluster'
    Properties:
      IamRoles:
        - !GetAtt 
          - SageMakerExecutionRole
          - Arn
      DBName: dev
      Encrypted: true
      NodeType: dc2.large
      ClusterType: single-node
      ClusterSubnetGroupName: !Ref RedshiftSubnetGroup
      MasterUsername: !Sub '{{resolve:secretsmanager:${DemoRedshiftSecret}::username}}'
      MasterUserPassword: !Sub '{{resolve:secretsmanager:${DemoRedshiftSecret}::password}}'
      PubliclyAccessible: false
      VpcSecurityGroupIds:
        - !Ref SageMakerDemoSG
  RedshiftSubnetGroup:
    Type: 'AWS::Redshift::ClusterSubnetGroup'
    Properties:
      Description: Demo Subnet Group
      SubnetIds:
        - !Ref PrivateSubnet
  SecretRedshiftAttachment:
    Type: 'AWS::SecretsManager::SecretTargetAttachment'
    DependsOn: DemoRedshiftCluster
    Properties:
      SecretId: !Ref DemoRedshiftSecret
      TargetId: !Ref DemoRedshiftCluster
      TargetType: 'AWS::Redshift::Cluster'
  glueConnectionforRedshift:
    Type: 'AWS::Glue::Connection'
    DependsOn: SecretRedshiftAttachment
    Properties:
      CatalogId: !Ref 'AWS::AccountId'
      ConnectionInput:
        Name: !Ref glueConnectionName
        Description: Redshift JDBC Connection
        PhysicalConnectionRequirements:
          SecurityGroupIdList:
            - !Ref SageMakerDemoSG
          SubnetId: !Ref PrivateSubnet
          AvailabilityZone: !GetAtt 
            - PrivateSubnet
            - AvailabilityZone
        ConnectionType: JDBC
        ConnectionProperties:
          JDBC_ENFORCE_SSL: false
          JDBC_CONNECTION_URL: !Join 
            - ''
            - - 'jdbc:redshift://'
              - !Sub '{{resolve:secretsmanager:${DemoRedshiftSecret}::host}}'
              - ':'
              - !Sub '{{resolve:secretsmanager:${DemoRedshiftSecret}::port}}'
              - /
              - !Sub '{{resolve:secretsmanager:${DemoRedshiftSecret}::dbname}}'
          SECRET_ID: !Ref RedshiftSecret
          KAFKA_SSL_ENABLED: false
  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess'
        - 'arn:aws:iam::aws:policy/AmazonS3FullAccess'
        - 'arn:aws:iam::aws:policy/IAMFullAccess'
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
  LambdaExecutionPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      Path: /
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: CloudWatchLogsPermissions
            Effect: Allow
            Action:
              - 'logs:CreateLogGroup'
              - 'logs:CreateLogStream'
              - 'logs:PutLogEvents'
            Resource: !Sub 'arn:${AWS::Partition}:logs:*:*:*'
          - Sid: SageMakerDomainPermission
            Effect: Allow
            Action:
              - 'sagemaker:ListDomains'
              - 'sagemaker:CreateDomain'
              - 'sagemaker:DescribeDomain'
              - 'sagemaker:DeleteDomain'
              - 'sagemaker:UpdateDomain'
              - 'sagemaker:ListUserProfiles'
              - 'sagemaker:CreateUserProfile'
              - 'sagemaker:UpdateUserProfile'
              - 'sagemaker:DeleteUserProfile'
              - 'sagemaker:DescribeUserProfile'
              - 'sagemaker:ListApps'
              - 'sagemaker:CreateApp'
              - 'sagemaker:DescribeApp'
              - 'sagemaker:DeleteApp'
              - 'sagemaker:UpdateApp'
            Resource:
              - !Sub 'arn:${AWS::Partition}:sagemaker:*:*:domain/*'
              - !Sub 'arn:${AWS::Partition}:sagemaker:*:*:user-profile/*'
              - !Sub 'arn:${AWS::Partition}:sagemaker:*:*:app/*'
          - Sid: SageMakerProjectsPermission
            Effect: Allow
            Action:
              - 'servicecatalog:AcceptPortfolioShare'
              - 'sagemaker:EnableSagemakerServicecatalogPortfolio'
              - 'sagemaker:DisableSagemakerServicecatalogPortfolio'
            Resource: '*'
          - Sid: SageMakerExecPassRole
            Effect: Allow
            Action:
              - 'iam:PassRole'
            Resource: !GetAtt 
              - SageMakerExecutionRole
              - Arn
      Roles:
        - !Ref LambdaExecutionRole
  SageMakerExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
                - redshift.amazonaws.com
                - glue.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/AmazonSageMakerFullAccess'
        - 'arn:aws:iam::aws:policy/service-role/AwsGlueSessionUserRestrictedServiceRole'
        - 'arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole'
        - 'arn:aws:iam::aws:policy/AmazonRedshiftReadOnlyAccess'
      Path: /
      Policies:
        - PolicyName: !Sub 'SageMaker-secretsmanager-policy-2023-${AWS::StackName}'
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'secretsmanager:GetSecretValue'
                Resource:
                  - !Sub '${DemoRedshiftSecret}'
        - PolicyName: !Sub 'SageMaker-iam-policy-2023-${AWS::StackName}'
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'iam:GetRole'
                  - 'iam:PassRole'
                  - 'sts:GetCallerIdentity'
                Resource:
                  - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*'
        - PolicyName: !Sub 'SageMaker-redshift-policy-2023-${AWS::StackName}'
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'redshift:getclustercredentials'
                Resource: '*'
  DomainDeploymentDelay:
    Type: 'Custom::Delay'
    DependsOn: SageMakerExecutionRole
    Properties:
      ServiceToken: !GetAtt 
        - DelayLambda
        - Arn
      TimeToWait: 10
  DelayLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      Handler: index.handler
      Timeout: 660
      Role: !GetAtt 
        - LambdaExecutionRole
        - Arn
      Runtime: python3.7
      Code:
        ZipFile: |
          import json
          import cfnresponse
          import time
          def handler(event, context):
            time_to_wait = int(event['ResourceProperties']['TimeToWait'])
            print(f'Waiting for {time_to_wait} seconds')
            time.sleep(time_to_wait)
            print(f'Waiting finished')
            cfnresponse.send(event, context, cfnresponse.SUCCESS, {},'')
  
  DefaultVpcLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: CFGetDefaultVpcId
      Code:
        ZipFile: |
          import json
          import boto3
          import cfnresponse

          ec2 = boto3.client('ec2')

          def lambda_handler(event, context):              
              if 'RequestType' in event and event['RequestType'] == 'Create':
                  vpc_id = get_default_vpc_id()
                  subnets =  get_subnets_for_vpc(vpc_id)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {'VpcId': vpc_id , "Subnets" : subnets}, '')
              else:
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {},'')

          def get_default_vpc_id():
              vpcs = ec2.describe_vpcs(Filters=[{'Name': 'is-default', 'Values': ['true']}])
              vpcs = vpcs['Vpcs']
              vpc_id = vpcs[0]['VpcId']
              return vpc_id


          def get_subnets_for_vpc(vpcId):
              response = ec2.describe_subnets(
                  Filters=[
                      {
                          'Name': 'vpc-id',
                          'Values': [vpcId]
                      }
                  ]
              )
              subnet_ids = []
              for subnet in response['Subnets']:
                  subnet_ids.append(subnet['SubnetId'])
              return subnet_ids 
      Description: Return default VPC ID and Subnets
      Handler: index.lambda_handler
      MemorySize: 512
      Role: !GetAtt 
        - LambdaExecutionRole
        - Arn
      Runtime: python3.7
      Timeout: 5
  DefaultVpcFinder:
    Type: 'Custom::ResourceForFindingDefaultVpc'
    Properties:
      ServiceToken: !GetAtt 
        - DefaultVpcLambda
        - Arn
  
  StudioDomain:
    Type: 'AWS::SageMaker::Domain'
    DependsOn: DomainDeploymentDelay
    Properties:
      AppNetworkAccessType: VpcOnly
      AuthMode: IAM
      DefaultUserSettings:
        ExecutionRole: !GetAtt 
          - SageMakerExecutionRole
          - Arn
        SecurityGroups:
          - !Ref SageMakerDemoSG
      DomainName: !Ref SageMakerDomainName
      SubnetIds:
        - !Ref PrivateSubnet
      VpcId: !GetAtt 
        - DefaultVpcFinder
        - VpcId
  UserProfile:
    Type: 'AWS::SageMaker::UserProfile'
    Properties:
      DomainId: !GetAtt 
        - StudioDomain
        - DomainId
      UserProfileName: !Ref SageMakerUserProfileName
      UserSettings:
        ExecutionRole: !GetAtt 
          - SageMakerExecutionRole
          - Arn
  JupyterApp:
    Type: 'AWS::SageMaker::App'
    DependsOn: UserProfile
    Properties:
      AppName: default
      AppType: JupyterServer
      DomainId: !GetAtt 
        - StudioDomain
        - DomainId
      UserProfileName: !Ref SageMakerUserProfileName
  DataScienceApp:
    Type: 'AWS::SageMaker::App'
    DependsOn: UserProfile
    Properties:
      AppName: Sagemaker-demo-datascience-ml-t3-medium
      AppType: KernelGateway
      DomainId: !GetAtt 
        - StudioDomain
        - DomainId
      ResourceSpec:
        InstanceType: ml.t3.medium
        SageMakerImageArn: !FindInMap 
          - RegionMap
          - !Ref 'AWS::Region'
          - datascience
      UserProfileName: !Ref SageMakerUserProfileName
Outputs:
  DemoRedshiftCluster:
    Description: Redshift cluster
    Value: !Sub >-
      ${DemoRedshiftCluster.Endpoint.Address}:${DemoRedshiftCluster.Endpoint.Port}
  SecretforRedshift:
    Description: AWS Secrets Manager secret for Redshift cluster
    Value: !Ref RedshiftSecret
  StudioDomainID:
    Description: The domain name of the Sagemaker studio instance
    Value: !Ref StudioDomain
  UserProfileID:
    Description: The user profile name for the SageMaker workshop
    Value: !Ref UserProfile
