[AWS] Managed KB ACLをRDBと自動同期させる

📝 TL;DR

• Managed KBのACLメタデータは、リアルタイム反映を頑張るより定期的バッチ反映がお勧め
• どうしても即時に近い反映が必要ならセルフでメタデータフィルタ実装がお勧め

👋 はじめに

AWSのManaged　KBに文書単位のアクセス制御を効かせるACL awarenessという機能があります（詳細はこちらの素晴らしい記事をご参照ください）

この仕組みを実運用に乗せようとすると「その accessControlList は誰がいつ書くの？」という疑問にぶつかります。今回はこのアーキテクチャを考えてみました。

ACLには2種類ありますが、今回はグローバルACL（global-acl.json）の更新を対象としております

🗄️ ユーザー情報は大抵RDBにいる

ユーザー情報は大抵AuroraのようなリレーショナルDBの users テーブルで管理することが多いです。ACL のメタデータをこれと手動で二重管理するのは正直しんどいですし、かといって「ユーザー登録・更新のアプリケーションコードの中でついでにS3のglobal-acl.jsonも書き換える」という実装はあまり気が進みません。

というわけでAuroraの変更を検知してS3のglobal-acl.jsonを書き換えManaged KBの同期をキックするというアーキテクチャにしました。

💡 案１:Aurora変更イベントをキャッチしてLambdaで同期

AuroraMySQLにはトリガー関数から直接Lambdaを呼び出せるネイティブの lambda_async() という仕組みがあります。IAMロールをDBクラスタに関連付けておけばAFTER INSERTやAFTER UPDATEのトリガーの中からこの関数を呼ぶだけで、変更のあった行の情報を Lambda に渡せます。

💥 問題発生

グローバルACLのメタデータはデータソース単位で1つです。
つまり営業部に100人追加しても本来書き換えるべきglobal-acl.jsonは1個のはずなのにトリガーは100回発火しそのたびに独立したLambda実行が同じファイルを読み込み、書き換え、書き込みしてしまいます。同じファイルに対して100回分の書き換えが走ってしまったわけです。(終いには100回同期するというお粗末な展開に...)

💡 案2:SQSで重複排除を入れる

案１を改造してAuroraのトリガーとLambdaの間にSQSを挟んで重複を排除しました。
MessageDeduplicationIdに部署名を明示的に指定するようし、同じ部署への変更は5分以内なら1件に収束するようにしました。（この5分は変更不可）

💥 問題発生

１クエリで複数登録、更新のパターンには対処できましたが
SQSの重複排除は同じ部署宛のメッセージを5分間覚えていて弾きます。営業部に1人目、2人目、3人目と時間差で登録すると、1人目の登録でLambdaが動いてACLに反映されますが、2人目・3人目は同じ重複排除IDに引っかかって弾かれ、Auroraには存在するのにACLには反映されないまま取り残されます。
一括登録をさばくための重複排除が、今度は時間差登録を握りつぶす副作用を生んでしまいました。

💡 案3:定期バッチ処理（採用）

RDSトリガー方式は諦めEventBridgeで1時間ごとにLambdaを起動しAuroraのusersテーブルを丸ごと読みに行ってフォルダごとのACLと比較してglobal-acl.jsonを書き換えるというバッチ処理に切り替えました。

ALC作成lambda（クリックで展開）

sample.py

# 部署ごとにAuroraの現在のメンバーを取得し、グローバルACLファイル内の
# 対応するkeyPrefixのaclEntriesと比較する。差分がなければ何もしない。
# 差分があれば書き換えて、いずれかの部署で変更があった場合のみ最後に
# StartIngestionJobを1回呼ぶ。
import json
import os

import boto3
from botocore.exceptions import ClientError

s3 = boto3.client("s3")
bedrock_agent = boto3.client("bedrock-agent")
rds_data = boto3.client("rds-data")

BUCKET = os.environ["BUCKET"]
KB_ID = os.environ["KB_ID"]
DATA_SOURCE_ID = os.environ["DATA_SOURCE_ID"]
GLOBAL_ACL_KEY = os.environ["GLOBAL_ACL_KEY"]
DB_CLUSTER_ARN = os.environ["DB_CLUSTER_ARN"]
DB_SECRET_ARN = os.environ["DB_SECRET_ARN"]
DATABASE_NAME = os.environ["DATABASE_NAME"]

# 部署名 -> S3上のフォルダ(グローバルACLファイルのkeyPrefix)の対応表
DEPARTMENT_PREFIXES = {
    "sales": f"s3://{BUCKET}/acl-sync-test/sales/",
}


def fetch_department_emails(department):
    resp = rds_data.execute_statement(
        resourceArn=DB_CLUSTER_ARN,
        secretArn=DB_SECRET_ARN,
        database=DATABASE_NAME,
        sql="SELECT email FROM users WHERE department = :department",
        parameters=[
            {"name": "department", "value": {"stringValue": department}}
        ],
    )
    return {r[0]["stringValue"] for r in resp["records"]}


def lambda_handler(event, context):
    try:
        acl_list = json.loads(s3.get_object(Bucket=BUCKET, Key=GLOBAL_ACL_KEY)["Body"].read())
    except ClientError as e:
        if e.response["Error"]["Code"] not in ("NoSuchKey", "404"):
            raise
        acl_list = []

    changed = False

    for department, key_prefix in DEPARTMENT_PREFIXES.items():
        entry = next((e for e in acl_list if e["keyPrefix"] == key_prefix), None)
        if entry is None:
            entry = {"keyPrefix": key_prefix, "aclEntries": []}
            acl_list.append(entry)

        current = fetch_department_emails(department)
        existing = {e["Name"] for e in entry["aclEntries"]}

        if current == existing:
            continue  # 差分なし。書き込みもsyncもしない

        entry["aclEntries"] = [
            {"Name": email, "Type": "USER", "Access": "ALLOW"}
            for email in sorted(current)
        ]
        changed = True

    if changed:
        s3.put_object(
            Bucket=BUCKET,
            Key=GLOBAL_ACL_KEY,
            Body=json.dumps(acl_list, ensure_ascii=False, indent=2).encode(),
            ContentType="application/json",
        )
        try:
            bedrock_agent.start_ingestion_job(knowledgeBaseId=KB_ID, dataSourceId=DATA_SOURCE_ID)
        except ClientError as e:
            # 直前のsyncが進行中だとConflictExceptionになりうる。S3の書き込み自体は
            # 成功しているので、次回のスケジュール実行時に自然にリトライされる。
            print(f"start_ingestion_job に失敗した(書き込み自体は成功済み): {e}")

    return {"changed": changed}

※部署固定読み込みのため実運用ではもう少し工夫する必要があります

締め〆

Managed Knowledge Base の同期はそもそも非同期で、反映までに数分~のラグがあります。しかもグローバルACLファイルの変更は、公式ドキュメントによると対象プレフィックス配下をまるごと再インデックスするため頻繁な書き換えには向きません。
Auroraの変更を1件も取りこぼさず即座に反映したいというニーズ自体がManaged KBと噛み合っていませんでした。即時反映が必要ならベクトルストアを自分で握れるセルフマネージドを選ぶべきでしょう。

しかしこうやってアーキテクチャを考えるのは、シミュレーションRPGの戦略を立ててる時の感覚に近いワクワク感がありますね！