[アップデート] RDS ブルー/グリーンデプロイの RDS Proxy サポートを検証してみた

import json
import os
import time
from datetime import datetime, timezone

import pymysql

_SSL_CA = "/etc/pki/tls/certs/ca-bundle.crt"

conn: pymysql.Connection | None = None


def _get_connection() -> pymysql.Connection:
    """
    グローバルコネクションを返す。
    初回呼び出し時にコネクションを確立し、以降は再利用する。
    切断されている場合は _probe() 側で conn を None にリセットしてから呼ぶこと。
    """
    global conn
    if conn is None:
        conn = pymysql.connect(
            host=os.environ["RDS_PROXY_HOST"],
            port=int(os.environ.get("DB_PORT", "3306")),
            user=os.environ["DB_USER"],
            password=os.environ["DB_PASSWORD"],
            database=os.environ["DB_NAME"],
            connect_timeout=5,
            read_timeout=5,
            write_timeout=5,
            ssl={"ca": _SSL_CA},
        )
    return conn


def _probe() -> dict:
    """
    DB に SELECT NOW() を 1 回実行し、結果と所要時間を返す。
    エラー時は例外を握りつぶし ok=False のレコードを返す。
    """
    ts = datetime.now(timezone.utc).isoformat()
    t0 = time.perf_counter()
    try:
        global conn
        if conn is not None:
            try:
                conn.ping(reconnect=False)
            except Exception:
                conn = None
        c = _get_connection()
        with c.cursor() as cur:
            cur.execute("SELECT NOW()")
            db_now = str(cur.fetchone()[0])
        elapsed_ms = round((time.perf_counter() - t0) * 1000)
        return {"ts": ts, "ok": True, "db_now": db_now, "elapsed_ms": elapsed_ms}
    except Exception as e:
        elapsed_ms = round((time.perf_counter() - t0) * 1000)
        conn = None
        return {"ts": ts, "ok": False, "error": str(e), "elapsed_ms": elapsed_ms}


def _run_downtime_check(duration_seconds: int, interval_seconds: float) -> dict:
    """
    duration_seconds 秒間、interval_seconds 秒ごとに DB への接続を試みる。
    各接続テスト結果を CloudWatch Logs に出力し、最後にサマリーを返す。
    """
    records: list[dict] = []
    end_at = time.monotonic() + duration_seconds

    print(json.dumps({
        "event": "start",
        "duration_seconds": duration_seconds,
        "interval_seconds": interval_seconds,
    }))

    while time.monotonic() < end_at:
        result = _probe()
        print(json.dumps(result))
        records.append(result)
        time.sleep(interval_seconds)

    errors = [r for r in records if not r["ok"]]
    first_error = errors[0] if errors else None
    last_error = errors[-1] if errors else None

    downtime_ms: int | None = None
    if first_error and last_error and first_error is not last_error:
        t_first = datetime.fromisoformat(first_error["ts"])
        t_last = datetime.fromisoformat(last_error["ts"])
        downtime_ms = round((t_last - t_first).total_seconds() * 1000) + last_error["elapsed_ms"]
    elif first_error:
        downtime_ms = first_error["elapsed_ms"]

    summary = {
        "total_probes": len(records),
        "success": len(records) - len(errors),
        "errors": len(errors),
        "estimated_downtime_ms": downtime_ms,
        "first_error": first_error,
        "last_error": last_error,
    }
    print(json.dumps({"event": "summary", **summary}))
    return summary


def lambda_handler(event: dict, context: object) -> dict:
    """
    duration_seconds 秒間、interval_seconds 秒ごとに DB への接続を試みる。
    Lambda のタイムアウトを duration_seconds + 30 秒以上に設定すること。
    """
    duration = int(event.get("duration_seconds", 120))
    interval = float(event.get("interval_seconds", 1))
    summary = _run_downtime_check(duration, interval)
    return {"statusCode": 200, "body": json.dumps(summary, ensure_ascii=False)}