Amazon Athena の Iceberg テーブルへの同時 UPDATE で ICEBERG_COMMIT_ERROR になったクエリを再試行してみた

CREATE TABLE IF NOT EXISTS order_status (order_id int, status string, updated_at timestamp)
LOCATION 's3://iceberg-orders-lake-k7m2/order_status/'
TBLPROPERTIES ('table_type' = 'ICEBERG', 'format' = 'parquet')

INSERT INTO order_status
SELECT id, 'pending', current_timestamp FROM UNNEST(sequence(1, 1000)) AS t(id)

"""Athena の Iceberg テーブルへ 2 つのクライアントから同時に UPDATE し、ICEBERG_COMMIT_ERROR と再試行を確かめる。

使い方:
  .venv/bin/python logs/Y24/iceberg_concurrency.py setup            # テーブルを作り 1,000 行入れる（何度実行しても 1,000 行）
  .venv/bin/python logs/Y24/iceberg_concurrency.py run noretry 5    # 5 ラウンド、再試行なし
  .venv/bin/python logs/Y24/iceberg_concurrency.py run retry 5      # 5 ラウンド、再試行あり
  .venv/bin/python logs/Y24/iceberg_concurrency.py snapshots        # スナップショットの数

- 再試行（SDK）：boto3 の standard モード（API の呼び出しのスロットリングなどに対して、最大 5 回）。
- 再試行（クエリ）：GetQueryExecution の State が FAILED で、ErrorMessage に ICEBERG_COMMIT_ERROR を含むときだけ、
  フルジッター付きの指数バックオフ（基準 1 秒、上限 30 秒、最大 5 回）で同じ文をもう一度 StartQueryExecution する。
- 冪等性：UPDATE に `AND status = 'pending'` を付け、同じ文を 2 回実行しても 2 回目は 0 行の更新になるようにする。
"""
import json
import random
import sys
import threading
import time

import boto3
from botocore.config import Config

REGION = "ap-northeast-1"
WORKGROUP = "orders-iceberg-wg"
DATABASE = "orders_lake"
TABLE = "order_status"
BUCKET = "iceberg-orders-lake-k7m2"

athena = boto3.client("athena", region_name=REGION, config=Config(retries={"mode": "standard", "max_attempts": 5}))

def run_query(sql, label=""):
    """クエリを 1 回実行して完了まで待ち、(state, error_message, scanned_bytes, elapsed_s, query_id) を返す。"""
    t0 = time.time()
    qid = athena.start_query_execution(QueryString=sql, QueryExecutionContext={"Database": DATABASE}, WorkGroup=WORKGROUP)["QueryExecutionId"]
    wait = 0.5
    while True:
        q = athena.get_query_execution(QueryExecutionId=qid)["QueryExecution"]
        state = q["Status"]["State"]
        if state in ("SUCCEEDED", "FAILED", "CANCELLED"):
            ae = q["Status"].get("AthenaError", {})
            err = ae.get("ErrorMessage", "") or q["Status"].get("StateChangeReason", "")
            if state == "FAILED":
                print(json.dumps({"failed_query": qid, "AthenaError": ae, "StateChangeReason": q["Status"].get("StateChangeReason", "")}, ensure_ascii=False), file=sys.stderr)
            scanned = q.get("Statistics", {}).get("DataScannedInBytes", 0)
            return state, err, scanned, time.time() - t0, qid
        time.sleep(wait)
        wait = min(2.0, wait * 1.5)

def run_with_retry(sql, label, max_attempts=5, base=1.0, cap=30.0):
    """ICEBERG_COMMIT_ERROR のときだけ再試行する。試行ごとの結果の一覧を返す。"""
    attempts = []
    for i in range(1, max_attempts + 1):
        state, err, scanned, elapsed, qid = run_query(sql, label)
        attempts.append({"attempt": i, "state": state, "error": err[:120], "scanned": scanned, "elapsed": round(elapsed, 1)})
        if state == "SUCCEEDED":
            break
        if state == "FAILED" and "ICEBERG_COMMIT_ERROR" in err and i < max_attempts:
            sleep = random.uniform(0, min(cap, base * 2 ** (i - 1)))
            attempts[-1]["backoff"] = round(sleep, 2)
            time.sleep(sleep)
            continue
        break
    return attempts

def fetch_scalar_rows(sql):
    state, err, _, _, qid = run_query(sql)
    if state != "SUCCEEDED":
        raise RuntimeError(f"{state}: {err}")
    rows = athena.get_query_results(QueryExecutionId=qid)["ResultSet"]["Rows"][1:]
    return [[c.get("VarCharValue") for c in r["Data"]] for r in rows]

def setup():
    run_query(f"""
        CREATE TABLE IF NOT EXISTS {TABLE} (order_id int, status string, updated_at timestamp)
        LOCATION 's3://{BUCKET}/{TABLE}/'
        TBLPROPERTIES ('table_type' = 'ICEBERG', 'format' = 'parquet')
    """)
    n = int(fetch_scalar_rows(f"SELECT count(*) FROM {TABLE}")[0][0])
    if n == 0:
        print(run_query(f"INSERT INTO {TABLE} SELECT id, 'pending', current_timestamp FROM UNNEST(sequence(1, 1000)) AS t(id)"))
    print("rows:", fetch_scalar_rows(f"SELECT count(*) FROM {TABLE}")[0][0])

def reset():
    state, err, *_ = run_query(f"UPDATE {TABLE} SET status = 'pending', updated_at = current_timestamp WHERE status <> 'pending'")
    assert state == "SUCCEEDED", err

def round_once(mode, rnd):
    # クライアント A は 1〜600、クライアント B は 401〜1000 を同時に更新する（同じデータファイルを対象にする）
    sqls = {
        "A": f"UPDATE {TABLE} SET status = 'packed', updated_at = current_timestamp WHERE order_id BETWEEN 1 AND 600 AND status = 'pending'",
        "B": f"UPDATE {TABLE} SET status = 'packed', updated_at = current_timestamp WHERE order_id BETWEEN 401 AND 1000 AND status = 'pending'",
    }
    results = {}

    def worker(name):
        if mode == "retry":
            results[name] = run_with_retry(sqls[name], name)
        else:
            state, err, scanned, elapsed, qid = run_query(sqls[name], name)
            results[name] = [{"attempt": 1, "state": state, "error": err[:120], "scanned": scanned, "elapsed": round(elapsed, 1)}]

    threads = [threading.Thread(target=worker, args=(n,)) for n in ("A", "B")]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    packed = int(fetch_scalar_rows(f"SELECT count(*) FROM {TABLE} WHERE status = 'packed'")[0][0])
    print(json.dumps({"round": rnd, "mode": mode, "results": results, "packed_rows": packed, "expected": 1000}, ensure_ascii=False))
    return results, packed

def main():
    cmd = sys.argv[1]
    if cmd == "setup":
        setup()
    elif cmd == "run":
        mode, rounds = sys.argv[2], int(sys.argv[3])
        summary = {"rounds": rounds, "mode": mode, "commit_errors": 0, "retries": 0, "incomplete_rounds": 0, "failed_scanned_bytes": 0}
        for r in range(1, rounds + 1):
            reset()
            results, packed = round_once(mode, r)
            for name, attempts in results.items():
                for a in attempts:
                    if a["state"] == "FAILED" and "ICEBERG_COMMIT_ERROR" in a["error"]:
                        summary["commit_errors"] += 1
                        summary["failed_scanned_bytes"] += a["scanned"]
                summary["retries"] += max(0, len(attempts) - 1)
            if packed != 1000:
                summary["incomplete_rounds"] += 1
        print("SUMMARY", json.dumps(summary))
    elif cmd == "snapshots":
        rows = fetch_scalar_rows(f'SELECT count(*) FROM "{TABLE}$snapshots"')
        print("snapshots:", rows[0][0])
    elif cmd == "status":
        print(fetch_scalar_rows(f"SELECT status, count(*) FROM {TABLE} GROUP BY status ORDER BY status"))

if __name__ == "__main__":
    main()

.venv/bin/python logs/Y24/iceberg_concurrency.py run noretry 5

ICEBERG_COMMIT_ERROR: Failed to commit Iceberg update to table: orders_lake.order_status

.venv/bin/python logs/Y24/iceberg_concurrency.py run retry 5

{"ErrorCategory": 2, "ErrorType": 233, "Retryable": false, "ErrorMessage": "ICEBERG_COMMIT_ERROR: Failed to commit Iceberg update to table: orders_lake.order_status"}