きゅぶろぐ

きゅぶんずの ぶろぐができて べんりだな

C#でLambdaからAthenaのクエリを叩く

この世界にC#でLambdaやAthena叩いてる人が少なすぎてネット上に全然コードが落ちてないので、
未来のためにコードの断片を残しておきます。

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using Amazon.Athena;
using Amazon.Athena.Model;
using Amazon.Lambda.Core;
using Newtonsoft.Json;

// Assembly attribute to enable the Lambda function's JSON input to be converted into a .NET class.
[assembly: LambdaSerializer(typeof(Amazon.Lambda.Serialization.SystemTextJson.LambdaJsonSerializer))]

namespace HelloWorld
{
    public class Function
    {
        /// <summary>
        /// A simple function that takes a string and does a ToUpper
        /// </summary>
        /// <param name="input"></param>
        /// <param name="context"></param>
        /// <returns></returns>
        public async Task<string> FunctionHandler(string input, ILambdaContext context)
        {
            var athenaClient = new AmazonAthenaClient(Amazon.RegionEndpoint.APNortheast1);

            var startQueryExecutionResponse = await athenaClient.StartQueryExecutionAsync(new StartQueryExecutionRequest
            {
                QueryString = @"select count(*) from sample_database;",
                WorkGroup = "primary",
                ResultConfiguration = new ResultConfiguration
                {
                    OutputLocation = "s3://bucket-name/"
                },
                QueryExecutionContext = new QueryExecutionContext
                {
                    Database = "default"
                }
            });

            while (true)
            {
                await Task.Delay(1000);

                var getQueryExecutionResponse = await athenaClient.GetQueryExecutionAsync(new GetQueryExecutionRequest
                {
                    QueryExecutionId = startQueryExecutionResponse.QueryExecutionId
                });

                var queryExecutionState = getQueryExecutionResponse.QueryExecution.Status.State;
                if (queryExecutionState == QueryExecutionState.RUNNING || queryExecutionState == QueryExecutionState.QUEUED) continue;

                if (queryExecutionState == QueryExecutionState.FAILED)
                {
                    return JsonConvert.SerializeObject(getQueryExecutionResponse);
                }
                break;
            }

            var getQueryResultsResponse = athenaClient.GetQueryResultsAsync(new GetQueryResultsRequest
            {
                QueryExecutionId = startQueryExecutionResponse.QueryExecutionId
            });

            return JsonConvert.SerializeObject(getQueryResultsResponse.Result.ResultSet.Rows);
        }
    }
}