Lambda Durable Functionsを利用してRedshiftに対しクエリ実行と結果のポーリングをしてみた

.
├── handler
│   └── sample
│       ├── query.sql
│       └── redshift_batch_executor.py
├── layers
│   └── dependencies
│       └── requirements.txt
├── samconfig.toml
└── template.yaml

version = 0.1

[default]
region = "ap-northeast-1"

[default.build.parameters]
debug = true

[default.deploy.parameters]
stack_name = "<YOUR_STACK_NAME>"
s3_bucket = "<YOUR_S3_BUCKET>"
s3_prefix = "<YOUR_S3_PREFIX>"
capabilities = "CAPABILITY_NAMED_IAM"
confirm_changeset = true

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: "Redshift Query Executor with Lambda Durable Functions"

Parameters:
  RedshiftWorkgroupName:
    Type: String
    Description: "Redshift Serverless Workgroup Name"
    Default: "<YOUR_WORKGROUP_NAME>"
  RedshiftDatabaseName:
    Type: String
    Description: "Redshift Database Name"
    Default: "<YOUR_DB_NAME>"
  DurableExecutionTimeout:
    Type: Number
    Description: "Durable Functions Execution Timeout (seconds)"
    Default: 3700

Globals:
  Function:
    Timeout: 180
    MemorySize: 128

Resources:
  # Lambda Layer
  DependenciesLayer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: "uehara-redshift-durable-functions-layer"
      Description: "Dependencies for Redshift Durable Functions"
      ContentUri: layers/dependencies/
      CompatibleRuntimes:
        - python3.14
      CompatibleArchitectures:
        - x86_64
    Metadata:
      BuildMethod: python3.14
      BuildArchitecture: x86_64

  # Lambda
  RedshiftQueryFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: "uehara-redshift-durable-functions"
      Role: !GetAtt RedshiftQueryFunctionRole.Arn
      CodeUri: handler/sample/
      Handler: redshift_batch_executor.lambda_handler
      Runtime: python3.14
      Architectures:
        - x86_64
      Timeout: 60
      Layers:
        - !Ref DependenciesLayer
      Environment:
        Variables:
          WORKGROUP_NAME: !Ref RedshiftWorkgroupName
          DATABASE_NAME: !Ref RedshiftDatabaseName
      DurableConfig:
        ExecutionTimeout: !Ref DurableExecutionTimeout

  # Lambda Role
  RedshiftQueryFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: "uehara-redshift-durable-functions-role"
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action: sts:AssumeRole
            Principal:
              Service:
                - lambda.amazonaws.com
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/AmazonRedshiftFullAccess
        - arn:aws:iam::aws:policy/AmazonRedshiftDataFullAccess
      Policies:
        - PolicyName: "DurableFunctionsPolicy"
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - lambda:CheckpointDurableExecution
                  - lambda:SendDurableExecutionCallbackSuccess
                  - lambda:SendDurableExecutionCallbackFailure
                Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:uehara-redshift-durable-functions*"

DELETE FROM public.mart_customer_summary;

INSERT INTO public.mart_customer_summary (
    customer_id,
    total_orders,
    total_amount,
    last_order_date,
    updated_at
)
SELECT
    customer_id,
    COUNT(DISTINCT order_id) AS total_orders,
    SUM(quantity * unit_price) AS total_amount,
    MAX(order_date) AS last_order_date,
    CURRENT_TIMESTAMP AS updated_at
FROM
    orders
GROUP BY
    customer_id;

import os
from pathlib import Path

import boto3
import sqlparse
from aws_durable_execution_sdk_python import (
    DurableContext,
    StepContext,
    durable_execution,
    durable_step,
)
from aws_durable_execution_sdk_python.config import Duration

# 環境変数
WORKGROUP_NAME = os.environ.get("WORKGROUP_NAME")
DATABASE_NAME = os.environ.get("DATABASE_NAME")

# Redshift Data APIクライアント
redshift_data = boto3.client("redshift-data")

def load_sql_statements(sql_file_path: str) -> list[str]:
    """SQLファイルを読み込み、個別のステートメントに分割する

    Args:
        sql_file_path (str): SQLファイルのパス

    Returns:
        list[str]: 分割されたSQLステートメントのリスト

    Raises:
        FileNotFoundError: SQLファイルが存在しない場合
    """
    with open(sql_file_path, "r", encoding="utf-8") as f:
        sql_content = f.read()

    # sqlparseでSQLを分割
    statements = sqlparse.split(sql_content)

    # 空のステートメントを除外
    return [stmt.strip() for stmt in statements if stmt.strip()]

@durable_step
def execute_batch_query(step_context: StepContext, statements: list) -> str:
    """Redshiftにバッチクエリを実行する

    Args:
        step_context (StepContext): Durable Functionsのステップコンテキスト
        statements (list): 実行するSQLステートメントのリスト

    Returns:
        str: クエリ実行ID
    """
    step_context.logger.info(f"Executing batch query with {len(statements)} statements")
    response = redshift_data.batch_execute_statement(
        WorkgroupName=WORKGROUP_NAME,
        Database=DATABASE_NAME,
        Sqls=statements,
    )
    query_id = response["Id"]
    step_context.logger.info(f"Batch query started with ID: {query_id}")
    return query_id

@durable_step
def check_query_status(step_context: StepContext, query_id: str) -> dict:
    """クエリの実行状態を確認する

    Args:
        step_context (StepContext): Durable Functionsのステップコンテキスト
        query_id (str): クエリ実行ID

    Returns:
        dict: クエリ状態情報
    """
    step_context.logger.info(f"Checking query status for ID: {query_id}")
    response = redshift_data.describe_statement(Id=query_id)
    status = response["Status"]
    step_context.logger.info(f"Query status: {status}")

    if status in ("FAILED", "ABORTED"):
        error_message = response.get("Error", "Unknown error")
        raise Exception(f"Query failed with status {status}: {error_message}")

    return {"status": status, "completed": status == "FINISHED"}

@durable_execution
def lambda_handler(event: dict, context: DurableContext) -> dict:
    """SQLファイル内の全ステートメントをバッチ実行するLambdaハンドラー

    Args:
        event (dict): Lambdaイベント
        context (DurableContext): Durable Functionsのコンテキスト

    Returns:
        dict: 実行結果を含むレスポンス
    """
    # SQLファイルのパスを取得（デフォルトは同ディレクトリのquery.sql）
    sql_file_path = event.get(
        "sql_file_path",
        str(Path(__file__).parent / "query.sql")
    )

    statements = load_sql_statements(sql_file_path)
    context.logger.info(f"Loaded {len(statements)} statements from {sql_file_path}")

    # バッチでクエリ実行
    query_id = context.step(execute_batch_query(statements))
    context.logger.info(f"Batch query submitted with ID: {query_id}")

    # ポーリングでクエリ完了を待機（10秒間隔、最大360回 = 1時間）
    max_attempts = 360
    for attempt in range(max_attempts):
        status_result = context.step(check_query_status(query_id))

        if status_result["completed"]:
            context.logger.info("Batch query completed successfully")
            break

        # 10秒待機
        context.wait(Duration.from_seconds(10))
    else:
        raise Exception(f"Query timed out after {max_attempts * 10} seconds")

    return {
        "statusCode": 200,
        "body": {
            "message": "All queries executed successfully",
            "query_id": query_id,
            "statement_count": len(statements),
            "status": status_result["status"],
        },
    }

boto3==1.42.23
sqlparse==0.5.0
aws-durable-execution-sdk-python==1.1.0

$ sam build --use-container

$ sam deploy