{
    "AWSTemplateFormatVersion": "2010-09-09",
    "Description": "Exivity AWS CUR Solution",
    "Resources": {
        "LambdaFunction": {
            "Type": "AWS::Lambda::Function",
            "Properties": {
                "Code": {
                    "ZipFile": "import time\nimport boto3\nimport botocore\nRETRY_COUNT = 10\ndef lambda_handler(event, context):\n    query = event['query']\n    S3_OUTPUT = event['s3output']\n    print(query)\n    client = boto3.client('athena')\n    response = client.start_query_execution(QueryString=query,ResultConfiguration={'OutputLocation': S3_OUTPUT,})\n    query_execution_id = response['QueryExecutionId']\n    print(query_execution_id)\n    s3 = boto3.client('s3')\n    S3_OUTPUT.replace(\"/\",\"\")\n    print(S3_OUTPUT)\n    for i in range(1, 1 + RETRY_COUNT):\n        query_status = client.get_query_execution(QueryExecutionId=query_execution_id)\n        query_execution_status = query_status['QueryExecution']['Status']['State']\n        if query_execution_status == 'SUCCEEDED':\n            print(\"STATUS:\" + query_execution_status)\n            url = s3.generate_presigned_url(ClientMethod='get_object',Params={'Bucket': S3_OUTPUT[5:],'Key': query_execution_id + '.csv'})\n            print(url)\n            return({'url': url})\n            break\n        if query_execution_status == 'FAILED':\n            raise Exception(\"STATUS:\" + query_execution_status)\n        else:\n            print(\"STATUS:\" + query_execution_status)\n            time.sleep(i)\n    else:\n        client.stop_query_execution(QueryExecutionId=query_execution_id)\n        raise Exception('TIME OVER')\n    result = client.get_query_results(QueryExecutionId=query_execution_id)\n"
                },
                "Description": "Lambda Function to query Athena DB",
                "Handler": "index.lambda_handler",
                "Runtime": "python2.7",
                "FunctionName": "QueryAthena2",
                "MemorySize": 512,
                "Role": {
                    "Fn::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": {
                    "Fn::GetAtt": [
                        "MyRestApi",
                        "RootResourceId"
                    ]
                },
                "PathPart": "QueryAthena2"
            }
        },
        "QueryAthena": {
            "Type": "AWS::ApiGateway::Method",
            "Properties": {
                "HttpMethod": "GET",
                "AuthorizationType": "NONE",
                "RequestValidatorId": {
                    "Ref": "RequestValidator"
                },
                "Integration": {
                    "Type": "AWS",
                    "IntegrationHttpMethod": "POST",
                    "Uri": {
                        "Fn::Join": [
                            "",
                            [
                                "arn:aws:apigateway:",
                                {
                                    "Ref": "AWS::Region"
                                },
                                ":lambda:path/2015-03-31/functions/",
                                {
                                    "Fn::GetAtt": [
                                        "LambdaFunction",
                                        "Arn"
                                    ]
                                },
                                "/invocations"
                            ]
                        ]
                    },
                    "IntegrationResponses": [
                        {
                            "StatusCode": 200
                        }
                    ],
                    "PassthroughBehavior": "WHEN_NO_MATCH",
                    "Credentials": {
                        "Fn::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": {
                "Fn::Sub": [
                    "https://${API}.execute-api.${AWS::Region}.amazonaws.com/v1/QueryAthena2",
                    {
                        "API": {
                            "Ref": "MyRestApi"
                        }
                    }
                ]
            }
        }
    }
}