Cognitoの全ユーザーを取得するPythonスクリプト

はじめに

Cognitoの全ユーザーの情報をまとめて取得したいときはありませんか？　ネットで検索するとちらほらと出てきますが、Generatorで実装したものがなかったので備忘録的に残しておきます。

コード

コードは以下の通りです。

from typing import Generator, List
import boto3

client = boto3.client('cognito-idp')


def cognito_list_all_users(user_pool_id: str) -> Generator[List[dict], None, None]:
    # 1回目の取得
    # 必ずPaginationTokenが無いため
    res = client.list_users(
        UserPoolId=user_pool_id
    )
    yield res["Users"]
    
    while pagination_token := res.get("PaginationToken"):
        # 2回目以降の取得
        res = client.list_users(
            UserPoolId=user_pool_id,
            PaginationToken=pagination_token
        )
        yield res["Users"]

PagenationTokenがない1回目の場合と2回目以降のある場合とで分けて処理しています。そして、得られたCognitoのユーザーの結果をyieldしています。一応の注意点として、Python3.8から導入されたセイウチ演算子:=を使っています。

上記の関数を利用するときは、以下のようになります。

for u in cognito_list_all_users(user_pool_id="ap-northeast-1_**"):
    print(len(u))

おわりに

たまに、Cognitoに存在する全ユーザーの情報が必要になりその都度、簡易的に作成していました。ユーザーが少ない場合は問題ないのですがページネーションが入ってくると地味に面倒なので記事としてまとめました。

備忘録的な側面が強いですが誰かの役に立てば嬉しいです。

ChatGPT4が出力したコード。

ChatGPT4を使って、以下の順序でコードの作成をしてみました。

1回目 「AWSのcognitoの、ユーザーを一覧する関数をPythonでboto3を使って出力できるようにして」

全件の取得を考慮できてない状態でした。

import boto3

# Cognitoのクライアントを作成
client = boto3.client('cognito-idp')

# ユーザープールのIDを設定
user_pool_id = 'your-user-pool-id'

def list_users():
    # ユーザーの一覧を取得
    try:
        response = client.list_users(UserPoolId=user_pool_id)
        return response['Users']
    except client.exceptions.ClientError as e:
        print(f"An error occurred: {e}")
        return None

# 関数を実行してユーザーの一覧を取得
users = list_users()
if users is not None:
    for user in users:
        print(user)

2回目 「全件を取得する関数も書いて」

これだけでもほとんど使えるのですが、関数の最後にまとめて結果が返ってくるので件数が多くなるとちょっと大変なことになりそうです。ネットで検索するとこのタイプの取得方法が出てきます。

import boto3

# Cognitoのクライアントを作成
client = boto3.client('cognito-idp')

# ユーザープールのIDを設定
user_pool_id = 'your-user-pool-id'

def list_all_users():
    # 全てのユーザーを格納するリスト
    all_users = []

    # 最初のページを取得
    try:
        response = client.list_users(UserPoolId=user_pool_id)
        all_users.extend(response['Users'])

        # 残りのページを取得
        while 'PaginationToken' in response:
            response = client.list_users(UserPoolId=user_pool_id, PaginationToken=response['PaginationToken'])
            all_users.extend(response['Users'])
    except client.exceptions.ClientError as e:
        print(f"An error occurred: {e}")
        return None

    return all_users

# 関数を実行して全てのユーザーを取得
users = list_all_users()
if users is not None:
    for user in users:
        print(user)

3回目 「Pythonのジェネレータを使った関数に変更して」

while文の中でclient.list_users()の処理が分岐されています。
while文の中で毎回評価していると、pagination_tokenがランダムにどちらかの状態になり得るように見えてしまいます。

import boto3

# Cognitoのクライアントを作成
client = boto3.client('cognito-idp')

# ユーザープールのIDを設定
user_pool_id = 'your-user-pool-id'

def list_users_generator():
    # ページネーショントークンの初期化
    pagination_token = None

    while True:
        # ページネーショントークンを使用してユーザーを取得
        if pagination_token:
            response = client.list_users(UserPoolId=user_pool_id, PaginationToken=pagination_token)
        else:
            response = client.list_users(UserPoolId=user_pool_id)

        # ユーザーの情報をyieldで返す
        for user in response.get('Users', []):
            yield user

        # 次のページがあるかチェック
        pagination_token = response.get('PaginationToken')
        if not pagination_token:
            break

# ジェネレータ関数を使用してユーザーを取得
for user in list_users_generator():
    print(user)

4回目 「上記の内容にさらに、while文のところでセイウチ演算子を使って書き直して」

4回目の結果だと、1行が長くなってしまっているためパッと見で何をしているのかわかりづらくなっている気がしました。ただ、やっぱりモック的な利用には最適だなと思いました（今回載せたコード作成時には利用していないですが…）。

import boto3

# Cognitoのクライアントを作成
client = boto3.client('cognito-idp')

# ユーザープールのIDを設定
user_pool_id = 'your-user-pool-id'

def list_users_generator():
    pagination_token = None

    while True:
        # ページネーショントークンを使用してユーザーを取得し、次のトークンがあれば更新
        response = client.list_users(UserPoolId=user_pool_id, PaginationToken=pagination_token) if pagination_token else client.list_users(UserPoolId=user_pool_id)

        for user in response.get('Users', []):
            yield user

        # セイウチ演算子を使用して、ページネーショントークンをチェックし、更新
        if not (pagination_token := response.get('PaginationToken')):
            break

一発で載せたようなコードが出てくれると、嬉しいなという思いもあってChatGPCの結果も記事に載せておきます。

参考

https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/cognito-idp.html