[python/boto3] Amazon Athenaの同時実行数制御のため、実行中件数を取得してみた

Amazon Athenaは素晴らしいサービスです。
もうSQLはAthena以外では書きたくない。。。

だって、index貼らなくていいし、重い重い処理かけてもDBサーバーごと落ちたりしない（注意：タイムアウトはあります。）から、効率考えずに書きたいクエリをかける！
MySQLだと、2000万件とかセレクトしに行ったらどうしよう、とか。
スロークエリになって他の処理に影響与えてしまわないかな、とか。心配事が尽きないです。。。

その点、Athenaはインフラのリソース気にしなくてクエリ投げ放題なんて、天国のSQLです！！
だから、ぜひ、使い倒したいところです。

ただ、同時実行数がデフォルトだと、たった20しかありません！（BigQueryは100なのに！上限解除もできるみたいですが）
だから、実行しようとした時に、20以上動いていたら、StepFunctionsとかでWait処理かけたいですよね？

その時に使える、boto3で実行中のクエリ数（State = Queued , Running)を取得するコードを書きました。Queuedも一応入れときました。
20件のStateを取れればいいんですが、一応、何件でも遡って取得できるようにしています。（上限解除している人のためにも）

import boto3

# Athenaの実行中の件数を取得する関数
def count_athena_workings(works:int=30):
    client = boto3.client('athena')
    ret = []
    def get_list(works,ret,next_token=''):
        max_num = 50 #2021年10月現在、list_query_executionsの上限が50
        work_nums =  max_num if works > max_num else works
        works -= max_num # 実行した分を減算
        
        if next_token: #2ページ目以降を取得する場合
            res = client.list_query_executions(NextToken=next_token,MaxResults=work_nums,WorkGroup='primary')
        else: #1ページ目を取得する場合
            res = client.list_query_executions(MaxResults=work_nums,WorkGroup='primary')
        next_token = res['NextToken'] #次のページ実行するためのnext_tokenを取得

        res2 = ''
        if res['QueryExecutionIds']:
            res2 = client.batch_get_query_execution(QueryExecutionIds=res['QueryExecutionIds'])
        
        for q in res2['QueryExecutions']:
            if q['Status']['State'] in ['QUEUED','RUNNING']:  # 'State': 'QUEUED'|'RUNNING'|'SUCCEEDED'|'FAILED'|'CANCELLED'
                ret.append(q['QueryExecutionId'])
                
        if works > 0: #2ページ目以降は再起的に取得
            ret = get_list(works,ret,next_token)
            
        return ret
            
    ret = get_list(works,ret) #初期実行
    
    return len(ret)

これで、Athenaを存分に働かせられますね！
（もっと良い方法があれば教えてください！）

Athenaを思い通りに操れると、データ分析し放題です！
億単位のレコードでも、大丈夫！
何百万レコードとか、ちょっと少ないくらいに感じる、この優越感が最高です。