AWSTemplateFormatVersion: 2010-09-09
Description: Exivity AWS CUR Solution
Resources:
  LambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code:
        ZipFile: |
          import time
          import boto3
          import botocore
          RETRY_COUNT = 10
          def lambda_handler(event, context):
              query = event['query']
              S3_OUTPUT = event['s3output']
              print(query)
              client = boto3.client('athena')
              response = client.start_query_execution(QueryString=query,ResultConfiguration={'OutputLocation': S3_OUTPUT,})
              query_execution_id = response['QueryExecutionId']
              print(query_execution_id)
              s3 = boto3.client('s3')
              S3_OUTPUT.replace("/","")
              print(S3_OUTPUT)
              for i in range(1, 1 + RETRY_COUNT):
                  query_status = client.get_query_execution(QueryExecutionId=query_execution_id)
                  query_execution_status = query_status['QueryExecution']['Status']['State']
                  if query_execution_status == 'SUCCEEDED':
                      print("STATUS:" + query_execution_status)
                      url = s3.generate_presigned_url(ClientMethod='get_object',Params={'Bucket': S3_OUTPUT[5:],'Key': query_execution_id + '.csv'})
                      print(url)
                      return({'url': url})
                      break
                  if query_execution_status == 'FAILED':
                      raise Exception("STATUS:" + query_execution_status)
                  else:
                      print("STATUS:" + query_execution_status)
                      time.sleep(i)
              else:
                  client.stop_query_execution(QueryExecutionId=query_execution_id)
                  raise Exception('TIME OVER')
              result = client.get_query_results(QueryExecutionId=query_execution_id)
      Description: Lambda Function to query Athena DB
      Handler: index.lambda_handler
      Runtime: python2.7
      FunctionName: QueryAthena2
      MemorySize: 512
      Role: !GetAtt 
        - ApigwRole
        - Arn
      Timeout: 300
  ApigwRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
                - apigateway.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Policies:
        - PolicyName: MyNewSuperLambdaPolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - '*'
                Resource: '*'
      RoleName: MyNewSuperLambdaRole
  MyRestApi:
    Type: 'AWS::ApiGateway::RestApi'
    Properties:
      Name: QueryAthena2
      EndpointConfiguration:
        Types:
          - REGIONAL
  UserModel:
    Type: 'AWS::ApiGateway::Model'
    Properties:
      ContentType: application/json
      Name: Query
      RestApiId: !Ref MyRestApi
      Schema:
        $schema: 'http://json-schema.org/draft-04/schema#'
        title: Query
        type: object
        additionalProperties: false
        properties:
          query:
            type: string
          s3output:
            type: string
        required:
          - query
          - s3output
  RequestValidator:
    Type: 'AWS::ApiGateway::RequestValidator'
    Properties:
      Name: Validator
      RestApiId: !Ref MyRestApi
      ValidateRequestBody: 'true'
      ValidateRequestParameters: 'true'
  UsersResource:
    Type: 'AWS::ApiGateway::Resource'
    Properties:
      RestApiId: !Ref MyRestApi
      ParentId: !GetAtt 
        - MyRestApi
        - RootResourceId
      PathPart: QueryAthena2
  QueryAthena:
    Type: 'AWS::ApiGateway::Method'
    Properties:
      HttpMethod: GET
      AuthorizationType: NONE
      RequestValidatorId: !Ref RequestValidator
      Integration:
        Type: AWS
        IntegrationHttpMethod: POST
        Uri: !Join 
          - ''
          - - 'arn:aws:apigateway:'
            - !Ref 'AWS::Region'
            - ':lambda:path/2015-03-31/functions/'
            - !GetAtt 
              - LambdaFunction
              - Arn
            - /invocations
        IntegrationResponses:
          - StatusCode: 200
        PassthroughBehavior: WHEN_NO_MATCH
        Credentials: !GetAtt 
          - ApigwRole
          - Arn
        RequestTemplates:
          application/json: >-
            {"query": "$input.params('query')", "s3output":
            "$input.params('s3output')"}
      OperationName: QueryAthena2
      RequestParameters:
        method.request.querystring.query: true
        method.request.querystring.s3output: true
      MethodResponses:
        - StatusCode: 200
      ResourceId: !Ref UsersResource
      RestApiId: !Ref MyRestApi
  RestApiDeployment:
    Type: 'AWS::ApiGateway::Deployment'
    Properties:
      RestApiId: !Ref MyRestApi
      StageName: v1
    DependsOn:
      - QueryAthena
Outputs:
  Endpoint:
    Description: Exivity Endpoint
    Value: !Sub 
      - 'https://${API}.execute-api.${AWS::Region}.amazonaws.com/v1/QueryAthena2'
      - API: !Ref MyRestApi
