template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  Serverless stack for automated data processing tasks

Metadata:
  AWS::ServerlessRepo::Application:
    Name: Athena Query execution
    Description: Implementation of Serverless Stack to schedule queries
    Author: Jose Arevalo
    Labels:
      - state-machine
      - athena-query
    SemanticVersion: 0.1.0

Parameters:
  DeploymentId:
    Type: String
    Description: Unique identifier for each deployment

  ScheduleGroupName:
    Type: String
    Default: "AthenaScheduledQueries"
    Description: "SCHEDULE GROUP NAME created"

  QueryOutputLocation:
    Type: String
    Description: "S3 URI to put Athena Query results location, like (s3://aws-athena-query-results-{AWS_ACCOUNT_ID}-eu-west-1/)"

Resources:
##########################################################################
# IAM Role for StateMachine
##########################################################################
  DataProcessingStateMachineRole:
    Type: AWS::IAM::Role
    Properties:
      #RoleName: latana.role.serverless-data-pipe.state-machine
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: states.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonAthenaFullAccess
        - arn:aws:iam::aws:policy/AmazonS3FullAccess

##########################################################################
# IAM Role for Lambda Function
##########################################################################
  EventBridgeSyncFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      #RoleName: latana.role.lambda.EventBridgeSyncRole
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/CloudWatchFullAccess
        - arn:aws:iam::aws:policy/AmazonEventBridgeFullAccess

##########################################################################
# IAM Role for Scheduler Execution
##########################################################################
  SchedulerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      #RoleName: latana.role.lambda.SchedulerExecutionRole
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: scheduler.amazonaws.com
            Action: sts:AssumeRole
      Path: "/"
      Policies:
        - PolicyName: "scheduler_start_execution_state_function_athena_queries"
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - states:StartExecution
                Resource:
                  - !Ref AthenaQueriesStateMachine

##########################################################################
# State Machine
##########################################################################
  AthenaQueriesStateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      StateMachineName: StateMachine-Athena-QueryExecution
      RoleArn: !GetAtt DataProcessingStateMachineRole.Arn
      DefinitionString: |-
        {
          "Comment": "A description of my state machine",
          "StartAt": "Athena StartQueryExecution",
          "States": {
            "Athena StartQueryExecution": {
              "Type": "Task",
              "Resource": "arn:aws:states:::athena:startQueryExecution",
              "Parameters": {
                "QueryString.$": "$.QueryString",
                "WorkGroup": "primary",
                "ResultConfiguration": {
                  "OutputLocation.$": "$.OutputLocation"
                }
              },
              "Next": "Wait",
              "ResultSelector": {
                "QueryExecutionId.$": "$.QueryExecutionId"
              },
              "ResultPath": "$.QueryExecution"
            },
            "Wait": {
              "Type": "Wait",
              "Next": "Athena GetQueryExecution",
              "SecondsPath": "$.WaitSeconds"
            },
            "Athena GetQueryExecution": {
              "Type": "Task",
              "Resource": "arn:aws:states:::athena:getQueryExecution",
              "Parameters": {
                "QueryExecutionId.$": "$.QueryExecution.QueryExecutionId"
              },
              "Next": "Choice"
            },
            "Choice": {
              "Type": "Choice",
              "Choices": [
                {
                  "Variable": "$.QueryExecution.Status.State",
                  "StringEquals": "SUCCEEDED",
                  "Next": "Success"
                },
                {
                  "Variable": "$.QueryExecution.Status.State",
                  "StringEquals": "FAILED",
                  "Next": "Fail"
                }
              ],
              "Default": "Wait"
            },
            "Success": {
              "Type": "Succeed"
            },
            "Fail": {
              "Type": "Fail"
            }
          }
        }

##########################################################################
# Lambda Function
##########################################################################
  EventBridgeSyncFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: main.lambda_handler
      Runtime: python3.10
      Timeout: 180
      Role: !GetAtt EventBridgeSyncFunctionRole.Arn
      Architectures:
        - x86_64
      Environment:
        Variables:
          STM_ARN: !Ref AthenaQueriesStateMachine
          ROLE_ARN: !GetAtt SchedulerExecutionRole.Arn
          SCHEDULE_GROUP_NAME: !Ref ScheduleGroupName
          OUTPUT_LOCATION: !Ref QueryOutputLocation
      Layers:
        - !Ref libs

##########################################################################
# Layer to the Lambda Function
##########################################################################
  libs:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: lambda-dependencies
      Description: Dependencies for Lambda, include YAML
      ContentUri: libs/

##########################################################################
# Invoke the Lambda Function to sync the schedules
##########################################################################
  EventBridgeSyncFunctionCustom:
    Type: Custom::EventBridgeSyncFunction
    Properties:
      ServiceToken: !GetAtt EventBridgeSyncFunction.Arn
      DeploymentId: !Ref DeploymentId  # This line uses the DeploymentId parameter