[AWS CDK] Athenaクエリを実行するLambda関数の権限を設定するカスタムConstructを作ってみた

#AWS CDK

#Amazon Athena

若槻龍太

2023.01.31

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんにちは、CX事業本部 Delivery部の若槻です。

Amazon Athenaでクエリの実行開始および実行結果取得を行いたい際には、AthenaだけでなくGlueやS3リソースへのアクセス権限が必要となります。

そのため、クエリを実行するリソース（Lambda関数）などに必要最低限のアクセス権限を付与したいとなるとIAMポリシー設定の記述が煩雑となってしまいます。

そこで今回は、Athenaクエリ実行に必要なパーミッションをLambda関数に設定するカスタムConstructをAWS CDKで作ってみたので、サンプルコード共有します。

カスタムConstructを使用することにより、Stack側の記述をシンプルにすることを目指します。

サンプルコード

次のようなカスタムConstructを作成しました。引数で渡されるLambda関数オブジェクトに必要なアクセス権限を付与しています。

import { Construct } from 'constructs';
import { aws_lambda, aws_iam, aws_athena, aws_s3, Stack } from 'aws-cdk-lib';
import * as glue from '@aws-cdk/aws-glue-alpha';

interface AthenaQueryLambdaFuncPermissionConstructProps {
  lambdaFunc: aws_lambda.Function; // Athenaクエリを実行するLambda関数
  glueDatabase: glue.Database; // Athenaクエリ実行対象Glueデータベース
  glueTable: glue.Table; // Athenaクエリ実行対象Glueテーブル
  athenaWorkGroup: aws_athena.CfnWorkGroup; // Athenaワークグループ
  athenaQueryResultBucket: aws_s3.Bucket; // Athenaクエリ実行結果保管Bucket
}

// Lambda関数に権限を付与するカスタムConstruct
export class AthenaQueryLambdaFuncPermissionConstruct extends Construct {
  constructor(
    scope: Construct,
    id: string,
    props: AthenaQueryLambdaFuncPermissionConstructProps
  ) {
    super(scope, id);

    const accountId = Stack.of(this).account;
    const region = Stack.of(this).region;

    // Glueテーブルの読み取り権限付与
    props.glueTable.grantRead(props.lambdaFunc);

    // Glueデータベースおよびデータカタログの読み取り権限付与
    props.lambdaFunc.addToRolePolicy(
      new aws_iam.PolicyStatement({
        actions: ['glue:Get*'],
        effect: aws_iam.Effect.ALLOW,
        resources: [
          props.glueDatabase.databaseArn,
          `arn:aws:glue:${region}:${accountId}:catalog`,
        ],
      })
    );

    // Athenaクエリ実行権限付与
    props.lambdaFunc.addToRolePolicy(
      new aws_iam.PolicyStatement({
        actions: [
          'athena:getDataCatalog',
          'athena:getQueryExecution',
          'athena:startQueryExecution',
          'athena:GetQueryResults',
        ],
        effect: aws_iam.Effect.ALLOW,
        resources: [
          `arn:aws:athena:${region}:${accountId}:workgroup/${props.athenaWorkGroup.name}`,
          `arn:aws:athena:${region}:${accountId}:datacatalog/AwsDataCatalog`,
        ],
      })
    );

    // Athenaクエリ実行結果保管バケットの書き込み権限付与
    props.athenaQueryResultBucket.grantReadWrite(props.lambdaFunc);
  }
}

上記のうち、arn:aws:glue:${region}:${accountId}:catalogは馴染みの無いリソースでしたが、クエリ時にこのリソースに対してglue:GetTableアクションが実行されるため権限付与が必要でした。

動作確認

Lambda関数ハンドラーのコードです。Amazon Athenaのクエリ実行開始およびクエリ結果取得のシンプルな記述を可能にするライブラリであるAthena-Queryを使ってSELECTクエリを実行しています。

import { Athena } from '@aws-sdk/client-athena';
import AthenaQuery from '@classmethod/athena-query';

const GLUE_DATABASE_NAME = process.env.GLUE_DATABASE_NAME || '';
const GLUE_TABLE_NAME = process.env.GLUE_TABLE_NAME || '';
const ATHENA_WORK_GROUP_NAME = process.env.ATHENA_WORK_GROUP_NAME || '';

const athena = new Athena({});
const athenaQuery = new AthenaQuery(athena, {
  db: GLUE_DATABASE_NAME,
  workgroup: ATHENA_WORK_GROUP_NAME,
});

export const handler = async () => {
  const items = [];

  for await (const item of athenaQuery.query(
    `SELECT * FROM ${GLUE_TABLE_NAME};`
  )) {
    items.push(item);
  }

  return items;
};

Stack Constructです。先程のカスタムConstructを呼び出して作成した各種リソースを引数として指定してます。

import { Construct } from 'constructs';
import {
  aws_lambda_nodejs,
  aws_s3,
  aws_athena,
  RemovalPolicy,
  Stack,
  StackProps,
} from 'aws-cdk-lib';
import * as glue_alpha from '@aws-cdk/aws-glue-alpha';

import { AthenaQueryLambdaFuncPermissionConstruct } from './constructs/athena-query-lambda-func-permission-construct';

export class CdkSampleStack extends Stack {
  constructor(scope: Construct, id: string, props: StackProps) {
    super(scope, id, props);

    // データソース格納バケット
    const dataBucket = new aws_s3.Bucket(this, 'dataBucket', {
      bucketName: `data-${this.account}-${this.region}`,
      removalPolicy: RemovalPolicy.DESTROY,
    });

    // Athenaクエリ実行結果保管Bucket
    const athenaQueryResultBucket = new aws_s3.Bucket(
      this,
      'athenaQueryResultBucket',
      {
        bucketName: `athena-query-result-${this.account}-${this.region}`,
        removalPolicy: RemovalPolicy.DESTROY,
      }
    );

    // Athenaクエリ実行対象Glueデータベース
    const glueDataBase = new glue_alpha.Database(this, 'glueDataBase', {
      databaseName: 'glue_data_base',
    });

    // Athenaクエリ実行対象Glueテーブル
    const glueTable = new glue_alpha.Table(this, 'glueTable', {
      tableName: 'glue_table',
      database: glueDataBase,
      bucket: dataBucket,
      s3Prefix: 'data/',
      dataFormat: glue_alpha.DataFormat.JSON,
      columns: [
        {
          name: 'amount',
          type: glue_alpha.Schema.INTEGER,
        },
        {
          name: 'deviceid',
          type: glue_alpha.Schema.STRING,
        },
      ],
    });

    // Athenaワークグループ
    const athenaWorkGroup = new aws_athena.CfnWorkGroup(
      this,
      'athenaWorkGroup',
      {
        name: 'athenaWorkGroup',
        workGroupConfiguration: {
          resultConfiguration: {
            outputLocation: `s3://${athenaQueryResultBucket.bucketName}/result-data`,
          },
        },
        recursiveDeleteOption: true,
      }
    );

    // Athenaクエリを実行するLambda関数
    const queryDevicesFunc = new aws_lambda_nodejs.NodejsFunction(
      this,
      'queryDevicesFunc',
      {
        environment: {
          GLUE_DATABASE_NAME: glueDataBase.databaseName,
          GLUE_TABLE_NAME: glueTable.tableName,
          ATHENA_WORK_GROUP_NAME: athenaWorkGroup.name,
        },
      }
    );

    // Lambda関数に権限を付与するカスタムConstruct
    new AthenaQueryLambdaFuncPermissionConstruct(
      this,
      'AthenaQueryLambdaFuncPermission',
      {
        lambdaFunc: queryDevicesFunc,
        glueDatabase: glueDataBase,
        glueTable: glueTable,
        athenaWorkGroup: athenaWorkGroup,
        athenaQueryResultBucket: athenaQueryResultBucket,
      }
    );
  }
}

CDKデプロイをしてリソース作成後に、Lambda関数を実行するとAthenaクエリが実行でき結果のデータを取得することができました。