Athenaのクエリ課金額をSlack通知する

はじめに

こんにちは！DR&MLOps グループの阿部です。

Gunosyには社内警察と呼ばれる人がおり、たとえばデータ可視化の際に円グラフを使うと正しい使い方を教えてくれる、母数という言葉の使い方を正してくれる、方々がいます。今回はAthenaで課金額の高いクエリを投げるとSlackで警告してくれる、Athena課金警察というボットを紹介します。そういえばこんな記事もありました。

data.gunosy.io

~~さあ、高額課金者を晒し上げにしてやりましょう！~~

どうやってやるか

最近 CloudWatch Events で Athenaの Query State 監視に対応したようです。

aws.amazon.com

今回はその機能を使い、Query State に変化があったときに Lambda を呼び出します。ドキュメントはこちら

docs.aws.amazon.com

CloudWatchイベントのルールはこのようにします。

{
  "detail": {
    "currentState": [
      "SUCCEEDED"
    ]
  },
  "detail-type": [
    "Athena Query State Change"
  ],
  "source": [
    "aws.athena"
  ]
}

また、ターゲットを Lambda にしておきます。こうすることで、Athenaのクエリが正常に実行されたときに Lambda が実行されます。

Lambdaの中身

イベントにはQueryIDが入っていて、このIDをもとにクエリの情報を取得することができます。今回はGoで実装されています。

Athenaクエリの状態変更を受け取るためのイベントが必要になります。こちらは今の所 aws-lambda-go に実装がないので、自分で実装する必要があります。

package events

import (
    "time"
)

const (
    AthenaEventSource                = "aws.athena"
    AthenaQueryStateChangeDetailType = "Athena Query State Change"
)

type AthenaQueryState string

const (
    AthenaQueryStateQueued    AthenaQueryState = "QUEUED"
    AthenaQueryStateRunning   AthenaQueryState = "RUNNING"
    AthenaQueryStateSucceeded AthenaQueryState = "SUCCEEDED"
    AthenaQueryStateFailed    AthenaQueryState = "FAILED"
    AthenaQueryStateCancelled AthenaQueryState = "CANCELLED"
)

type AthenaQueryStatementType string

const (
    AthenaQueryStatementTypeDdl     AthenaQueryStatementType = "DDL"
    AthenaQueryStatementTypeDml     AthenaQueryStatementType = "DML"
    AthenaQueryStatementTypeUtility AthenaQueryStatementType = "UTILITY"
)

// AthenaQueryStateChangeEvent is documented at:
// https://docs.aws.amazon.com/athena/latest/ug/athena-cloudwatch-events.html
type AthenaQueryStateChangeEvent struct {
    AccountID  string                            `json:"account"`
    Region     string                            `json:"region"`
    DetailType string                            `json:"detail-type"`
    Source     string                            `json:"source"`
    Version    string                            `json:"version"`
    Time       time.Time                         `json:"time"`
    ID         string                            `json:"id"`
    Resources  []string                          `json:"resources"`
    Detail     AthenaQueryStateChangeEventDetail `json:"detail"`
}

type AthenaQueryStateChangeEventDetail struct {
    VersionID        string                   `json:"versionId"`
    CurrentState     AthenaQueryState         `json:"currentState"`
    PreviousState    AthenaQueryState         `json:"previousState"`
    StatementType    AthenaQueryStatementType `json:"statementType"`
    QueryExecutionID string                   `json:"queryExecutionId"`
    WorkGroupName    string                   `json:"workgroupName"`
    SequenceNumber   string                   `json:"sequenceNumber"`
}

このイベントのハンドラーを実装します。

func handleEvent(ctx context.Context, event events.AthenaQueryStateChangeEvent) error {

    sess, _ := session.NewSession(
        &aws.Config{
            Region: aws.String(event.Region),
        },
    )

    svc := athena.New(sess)
    output, err := svc.GetQueryExecutionWithContext(
        ctx,
        &athena.GetQueryExecutionInput{
            QueryExecutionId: aws.String(event.Detail.QueryExecutionID),
        },
    )
...

上のスニペットを例に取ると、 output からクエリ内容: output.QueryExecution.Query スキャン量: output.QueryExecution.Statistics.DataScannedInBytes などを取得できます。

中身はクエリの中身、データスキャン量などが入っていて、予め決められた閾値を上回るときにSlack通知が飛ぶようにしてあります。また、スキャン量を元に課金額も通知しています。

結果

来ました！すごいスキャン量ですね。。

なんか見覚えがあるクエリだと思ったら、犯人は自分でした 😨 犯人第一号となってしまいました。

悩み

今の所通知にWorkGroupは出すようにしていますが、そのクエリを誰が書いたのかまでは特定できないのでクエリの中身を見てそれっぽい人に教えて上げるという運用になっています。何らかの方法でクエリを書いた人がわかる仕組みがあると良いなーと思っています。

まとめ

Gunosyでは課金額を通知して可視化することで必要以上にコストかからないようにしています。 Athenaを使っているとパーティションが効いていない、同じテーブルを繰り返しスキャンしてしまう等で必要以上にスキャン量が増えてしまうこともあると思います。今回は課金警察を導入し、そのようなクエリの発見・改善やバッチなどで繰り返しスキャンされる際の集計テーブル作成のための気付きにしています。また、Slackに通知されることによって活発にクエリチューニングの話し合いが行われるという利点もあります。分析者にとってもスキャン量に怯えることなく、安心して分析できるようになりそうです。