AWSTemplateFormatVersion: 2010-09-09
Description: >-
  Deploy Lambda Function that kicks of a SageMaker Pipeline.
Parameters:
  StaticCodeBucket:
    Type: String
    Default: "aws-hcls-ml"
  StaticCodeKey:
    Type: String
    Default: 'workshop/immersion_day_workshop_data_DO_NOT_DELETE/kick_off_sagemaker_pipelines_lambda/other_material/lambda.zip'
  SageMakerPipelineName:
    Type: String
    Default: "ObesityModelRetrainingPipeLine"

Resources:
  # Create the buckets that we need
  NotificationBucket:
    Type: 'AWS::S3::Bucket'
    
  LambdaTrigger:
    Type: 'Custom::LambdaTrigger'
    DependsOn: LambdaInvokePermission
    Properties:
      ServiceToken: !GetAtt CustomResourceLambdaFunction.Arn
      LambdaArn: !GetAtt CallMyLambdaCode.Arn
      Bucket: !Ref NotificationBucket

  # This is the lambda function that does the work
  CallMyLambdaCode: 
    Type: 'AWS::Lambda::Function'
    Properties:
      Code:
        S3Bucket: !Ref StaticCodeBucket
        S3Key: !Ref StaticCodeKey
      Handler: kick_off_pipeline_lambda.lambda_handler
      Role: !GetAtt LambdaIAMRole.Arn
      Runtime: python3.9
      Timeout: 300
      MemorySize: 128
      Environment:
        Variables:
          NOTIFICATION_BUCKET_NAME: !Ref NotificationBucket
          SAGEMAKER_PIPELINE_NAME: !Ref SageMakerPipelineName

  LambdaInvokePermission:
      Type: 'AWS::Lambda::Permission'
      Properties:
        FunctionName: !GetAtt CallMyLambdaCode.Arn
        Action: 'lambda:InvokeFunction'
        Principal: s3.amazonaws.com
        SourceAccount: !Ref 'AWS::AccountId'
        SourceArn: !Sub 'arn:aws:s3:::${NotificationBucket}'

  CustomResourceLambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt LambdaIAMRole.Arn
      Code:
        ZipFile: 
            |
            from __future__ import print_function
            import json
            import boto3
            import cfnresponse
            import os
            
            SUCCESS = "SUCCESS"
            FAILED = "FAILED"
            
            print('Loading function')
            s3 = boto3.resource('s3')
           
            def lambda_handler(event, context):
                print("Received event: " + json.dumps(event, indent=2))
                responseData={}
                try:
                    if event['RequestType'] == 'Delete':
                        print("Request Type:",event['RequestType'])
                        Bucket=event['ResourceProperties']['Bucket']
                        delete_notification(Bucket)
                        print("Sending response to custom resource after Delete")
                    elif event['RequestType'] == 'Create' or event['RequestType'] == 'Update':
                        print("Request Type:",event['RequestType'])
                        LambdaArn=event['ResourceProperties']['LambdaArn']
                        Bucket=event['ResourceProperties']['Bucket']
                        add_notification(LambdaArn, Bucket)
                        responseData={'Bucket':Bucket}
                        print("Sending response to custom resource")
                    responseStatus = 'SUCCESS'
                except Exception as e:
                    print('Failed to process:', e)
                    responseStatus = 'FAILURE'
                    responseData = {'Failure': 'Something bad happened.'}
                cfnresponse.send(event, context, responseStatus, responseData)

            def add_notification(LambdaArn, Bucket):
                bucket_notification = s3.BucketNotification(Bucket)
                response = bucket_notification.put(
                  NotificationConfiguration={
                    'LambdaFunctionConfigurations': [
                      {
                          'LambdaFunctionArn': LambdaArn,
                          'Events': [
                              's3:ObjectCreated:*'
                          ]
                      }
                    ]
                  }
                )
                print("Put request completed....")
              
            def delete_notification(Bucket):
                bucket_notification = s3.BucketNotification(Bucket)
                response = bucket_notification.put(
                    NotificationConfiguration={}
                )
                print("Delete request completed....")


      Runtime: python3.9
      Timeout: 500
      MemorySize: 1000

  LambdaIAMRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
                - comprehendmedical.amazonaws.com
                - comprehend.amazonaws.com
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 's3:GetBucketNotification'
                  - 's3:PutBucketNotification'
                Resource: !Sub 'arn:aws:s3:::${NotificationBucket}'
              - Effect: Allow
                Action:
                  - 's3:*'
                Resource: !Sub 'arn:aws:s3:::${NotificationBucket}/*'

              - Effect: Allow
                Action:
                  - 's3:*'
                Resource: !Sub 'arn:aws:s3:::${StaticCodeBucket}/*'
              - Effect: Allow
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
              - Effect: Allow
                Action: 
                  - "comprehendmedical:*"
                Resource: '*'
              - Effect: Allow
                Action: 
                  - "comprehend:*"
                Resource: '*'
              - Effect: Allow
                Action: 
                  - "sagemaker:*"
                Resource: '*'              

Outputs:
  InputBucket:
    Description: S3 bucket where data scientists can upload new data
    Value: !Ref NotificationBucket