AWSTemplateFormatVersion: "2010-09-09"
Description: "Stack for Redshift Demos"
Parameters:
  Cluster:
    Type: String
    AllowedPattern : .+
  Db:
    Type: String
    AllowedPattern : .+
  DbUser:
    Type: String
    AllowedPattern : .+

Resources:
  Bucket:
    Type: 'AWS::S3::Bucket'
  RedshiftRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: redshift.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: S3Access
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - s3:Get*
                  - s3:List*
                Resource:
                  - !Sub 'arn:aws:s3:::${Bucket}'
                  - !Sub 'arn:aws:s3:::${Bucket}/*'
        - PolicyName: GlueAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - glue:CreateDatabase
                  - glue:CreateTable
                  - glue:GetTable
                  - glue:DeleteTable
                Resource:
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:catalog'
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:database/myspectrum-db'
                  - !Sub 'arn:aws:glue:${AWS::Region}:${AWS::AccountId}:table/myspectrum-db/*'
  LambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: RedshiftAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: redshift:GetClusterCredentials
                Resource:
                  - !Sub 'arn:aws:redshift:${AWS::Region}:${AWS::AccountId}:dbuser:${Cluster}/${DbUser}'
                  - !Sub 'arn:aws:redshift:${AWS::Region}:${AWS::AccountId}:dbname:${Cluster}/${Db}'
              - Effect: Allow
                Action:
                  - redshift:ModifyClusterIamRoles
                  - redshift-data:ExecuteStatement
                  - redshift:DescribeClusters
                Resource: !Sub 'arn:aws:redshift:${AWS::Region}:${AWS::AccountId}:cluster:${Cluster}'
              - Effect: Allow
                Action: redshift-data:DescribeStatement
                Resource: '*'
              - Effect: Allow
                Action:
                  - iam:PassRole
                Resource: !GetAtt 'RedshiftRole.Arn'
        - PolicyName: S3Access
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: s3:*
                Resource:
                  - !Sub 'arn:aws:s3:::${Bucket}/*'
                  - !Sub 'arn:aws:s3:::${Bucket}'
              - Effect: Allow
                Action: s3:Get*
                Resource: 'arn:aws:s3:::redshift-downloads/*'
              - Effect: Allow
                Action: s3:ListBucket
                Resource: 'arn:aws:s3:::redshift-downloads'
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
  LambdaFunction:
    Type: "AWS::Lambda::Function"
    Properties:
      Timeout: 300
      Code:
        ZipFile: |
          import json
          import boto3
          import logging
          import time
          import urllib3

          logging.basicConfig()
          logging.basicConfig(format='%(asctime)s,%(msecs)d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s', datefmt='%Y-%m-%d:%H:%M:%S',level=logging.DEBUG)

          logger = logging.getLogger(__name__)
          logger.setLevel(logging.INFO)

          http = urllib3.PoolManager()
          redshiftdata_client = boto3.client('redshift-data')

          def execQuery(cluster, db, user, statement):
            try:
              response = redshiftdata_client.execute_statement(ClusterIdentifier=cluster, Database=db, DbUser=user, Sql=statement)
              response = redshiftdata_client.describe_statement(Id=response["Id"])

              while response['Status'] in ['SUBMITTED','PICKED','STARTED']:
                time.sleep(2)
                response = redshiftdata_client.describe_statement(Id=response["Id"])
              if response['HasResultSet']:
                response['Records'] = redshiftdata_client.get_statement_result(Id=response['Id'])['Records']
              print ('Finished: '+statement)
              return response
            except Exception as e:
              print("Failed Exec Query:", e)
              raise e

          def send(event, context, responseStatus, responseData):
            responseUrl = event['ResponseURL']
            responseBody = {
                'Status' : responseStatus,
                'Reason' : "See the details in CloudWatch Log Stream: {}".format(context.log_stream_name),
                'PhysicalResourceId' : context.log_stream_name,
                'StackId' : event['StackId'],
                'RequestId' : event['RequestId'],
                'LogicalResourceId' : event['LogicalResourceId'],
                'NoEcho' : False,
                'Data' : responseData
            }
            json_responseBody = json.dumps(responseBody)
            headers = {
                'content-type' : '',
                'content-length' : str(len(json_responseBody))
            }
            try:
                response = http.request('PUT', responseUrl, headers=headers, body=json_responseBody)
            except Exception as e:
                print("send(..) failed executing http.request(..):", e)

          def handler(event, context):
            logger.info(json.dumps(event))
            try:
              redshift_client = boto3.client('redshift')
              s3_client = boto3.client('s3')
              role=event['ResourceProperties']['RedshiftRole']
              cluster=event['ResourceProperties']['Cluster']
              bucket=event['ResourceProperties']['Bucket']
              db=event['ResourceProperties']['Db']
              dbUser=event['ResourceProperties']['DbUser']
              region=event['ResourceProperties']['Region']

              if event['RequestType'] == 'Delete':
                try:
                  s3 = boto3.resource('s3')
                  bucket = s3.Bucket(bucket)
                  bucket.objects.all().delete()
                except:
                  print('Bucket could not be emptied')

                try:
                  execQuery(cluster, db, dbUser, "drop table myspectrum_schema.sales;")
                  execQuery(cluster, db, dbUser, "drop schema if exists myspectrum_schema;")
                  execQuery(cluster, db, dbUser, "drop schema if exists mylocal_schema cascade;")
                except:
                  print('Table & schema could not be dropped.')

                try:
                  redshift_client.modify_cluster_iam_roles(ClusterIdentifier=cluster,RemoveIamRoles=[role])
                  waiter = redshift_client.get_waiter('cluster_available')
                  waiter.wait(ClusterIdentifier=cluster)
                except:
                  print('Role could not be removed')

                send(event, context, 'SUCCESS', {'Data': 'Delete complete'})

              else:
                try:
                  redshift_client.modify_cluster_iam_roles(ClusterIdentifier=cluster,AddIamRoles=[role])
                  waiter = redshift_client.get_waiter('cluster_available')
                  waiter.wait(ClusterIdentifier=cluster)

                  srcBucket='redshift-downloads'
                  try:
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/allevents_pipe.txt'}, bucket, 'tickit/allevents_pipe.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/allusers_pipe.txt'}, bucket, 'tickit/allusers_pipe.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/category_pipe.txt'}, bucket, 'tickit/category_pipe.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/date2008_pipe.txt'}, bucket, 'tickit/date2008_pipe.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/listings_pipe.txt'}, bucket, 'tickit/listings_pipe.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/sales_tab.txt'}, bucket, 'tickit/sales_tab.txt')
                    s3_client.copy({'Bucket': srcBucket,'Key': 'tickit/venue_pipe.txt'}, bucket, 'tickit/venue_pipe.txt')
                    body='{"entries": [{"url":"s3://'+bucket+'/tickit/sales_tab.txt","meta": {"content_length": 11260097}}]}'
                    s3_client.put_object(Bucket=bucket, Key='tickit/sales_tab.json', Body=body.encode('ascii'))
                  except:
                    print("bad object")

                  execQuery(cluster, db, dbUser, "drop table myspectrum_schema.sales;")
                  execQuery(cluster, db, dbUser, "drop schema if exists myspectrum_schema;")
                  execQuery(cluster, db, dbUser, "drop schema if exists mylocal_schema cascade;")

                  statement = "create external schema myspectrum_schema \
                      from data catalog database 'myspectrum-db'\
                      iam_role '" + role + "'\
                      create external database if not exists;"

                  result = execQuery(cluster, db, dbUser, statement)
                  if result['Status'] != 'FINISHED':
                    print(result)
                    send(event, context, 'FAILED', {'Data': 'Failed SQL: '+result['Error']})
                    return

                  statement = "create external table myspectrum_schema.sales(\
                      salesid integer,\
                      listid integer,\
                      sellerid integer,\
                      buyerid integer,\
                      eventid integer,\
                      dateid smallint,\
                      qtysold smallint,\
                      pricepaid decimal(8,2),\
                      commission decimal(8,2),\
                      saletime timestamp) \
                    row format delimited \
                    fields terminated by '\\t' \
                    stored as textfile \
                    location 's3://" + bucket + "/tickit/sales_tab.json';"

                  result = execQuery(cluster, db, dbUser, statement)
                  if result['Status'] != 'FINISHED':
                    print(result)
                    send(event, context, 'FAILED', {'Data': 'Failed SQL: '+result['Error']})
                    return

                  statement = "create schema mylocal_schema;\
                    create table mylocal_schema.event(\
                      eventid integer not null distkey,\
                      venueid smallint not null,\
                      catid smallint not null,\
                      dateid smallint not null sortkey,\
                      eventname varchar(200),\
                      starttime timestamp);\
                    \
                    copy mylocal_schema.event from 's3://" + bucket + "/tickit/allevents_pipe.txt' \
                      iam_role '" + role + "'\
                      delimiter '|' timeformat 'YYYY-MM-DD HH:MI:SS';"

                  result = execQuery(cluster, db, dbUser, statement)
                  if result['Status'] != 'FINISHED':
                    print(result)
                    send(event, context, 'FAILED', {'Data': 'Failed SQL: '+result['Error']})
                    return

                  send(event, context, 'SUCCESS', {'Data': 'Create complete'})
                except Exception as e:
                  logger.error(e)
                  send(event, context, 'FAILED', {'Data': 'Create failed'})
            except:
              send(event, context, 'FAILED', {'Data': 'Something failed'})
      Handler: index.handler
      Role: !GetAtt 'LambdaRole.Arn'
      Runtime: python3.9
  InitLambda:
    Type: Custom::LambdaDataLake
    Properties:
      ServiceToken: !GetAtt 'LambdaFunction.Arn'
      RedshiftRole: !GetAtt 'RedshiftRole.Arn'
      Bucket: !Ref Bucket
      Cluster: !Ref Cluster
      Db: !Ref Db
      DbUser: !Ref DbUser
      Region: !Sub ${AWS::Region}
Outputs:
  S3Bucket:
    Description: S3 Bucket contains source files.
    Value: !Ref Bucket
  RedshiftRole:
    Description: Name of Redshift Role added to your Redshift cluster.
    Value: !Ref RedshiftRole
  LambdaRole:
    Description: Name of Role attached to the Lambda function used to build resources.
    Value: !Ref LambdaRole
  LambdaFunction:
    Description: Name of Lambda function used to build resources.
    Value: !Ref LambdaFunction
