chatgpt-register/ops/d1_cleanup.py

"""Periodic Cloudflare D1 cleanup for cfmail storage tables."""

from __future__ import annotations

import argparse
import json
import os
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from .common import now

DEFAULT_D1_DATABASE_ID = ""
DEFAULT_D1_MAIL_RETENTION_HOURS = 2
DEFAULT_D1_ADDRESS_RETENTION_HOURS = 24
DEFAULT_D1_CLEANUP_BATCH_SIZE = 5000
_QUERY_TIMEOUT_SECONDS = 30
_CLOUDFLARE_API_BASE = "https://api.cloudflare.com/client/v4"

_missing_credentials_warned = False


class D1CleanupError(RuntimeError):
    """Raised when the D1 query API returns a non-recoverable error."""


class D1TableMissingError(D1CleanupError):
    """Raised when the target D1 table does not exist."""


def _warning(message: str) -> None:
    print(f"[{now()}] [d1_cleanup] warning: {message}")


def _credentials_from_env() -> tuple[str, str, str] | None:
    global _missing_credentials_warned

    auth_email = str(os.getenv("ZHUCE6_CFMAIL_CF_AUTH_EMAIL", "")).strip()
    auth_key = str(os.getenv("ZHUCE6_CFMAIL_CF_AUTH_KEY", "")).strip()
    account_id = str(os.getenv("ZHUCE6_CFMAIL_CF_ACCOUNT_ID", "")).strip()
    if auth_email and auth_key and account_id:
        _missing_credentials_warned = False
        return auth_email, auth_key, account_id
    if not _missing_credentials_warned:
        _warning(
            "missing Cloudflare credentials, skip cleanup "
            "(need ZHUCE6_CFMAIL_CF_AUTH_EMAIL / ZHUCE6_CFMAIL_CF_AUTH_KEY / ZHUCE6_CFMAIL_CF_ACCOUNT_ID)"
        )
        _missing_credentials_warned = True
    return None


def _error_messages(payload: dict[str, Any]) -> list[str]:
    messages: list[str] = []
    for bucket in ("errors", "messages"):
        items = payload.get(bucket)
        if not isinstance(items, list):
            continue
        for item in items:
            if isinstance(item, dict):
                message = str(item.get("message") or "").strip()
                if message:
                    messages.append(message)
    result_items = payload.get("result")
    if isinstance(result_items, list):
        for result in result_items:
            if not isinstance(result, dict):
                continue
            if bool(result.get("success", True)):
                continue
            message = str(result.get("error") or result.get("message") or "").strip()
            if message:
                messages.append(message)
    return messages


def _raise_for_payload(payload: dict[str, Any]) -> None:
    messages = _error_messages(payload)
    text = " | ".join(messages) if messages else json.dumps(payload, ensure_ascii=False)
    lowered = text.lower()
    if "no such table" in lowered or "sqlite_error" in lowered:
        raise D1TableMissingError(text)
    raise D1CleanupError(text)


def _query(database_id: str, sql: str, params: list[Any] | None = None) -> dict[str, Any]:
    credentials = _credentials_from_env()
    if credentials is None:
        raise D1CleanupError("missing_cloudflare_credentials")
    auth_email, auth_key, account_id = credentials
    url = f"{_CLOUDFLARE_API_BASE}/accounts/{account_id}/d1/database/{database_id}/query"
    body = {"sql": sql}
    if params:
        body["params"] = params
    request = Request(
        url,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Content-Type": "application/json",
            "X-Auth-Email": auth_email,
            "X-Auth-Key": auth_key,
        },
        method="POST",
    )
    try:
        with urlopen(request, timeout=_QUERY_TIMEOUT_SECONDS) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            detail = str(exc)
        try:
            payload = json.loads(detail)
        except json.JSONDecodeError:
            payload = None
        if isinstance(payload, dict):
            _raise_for_payload(payload)
        raise D1CleanupError(f"http {exc.code}: {detail}") from exc
    except URLError as exc:
        raise D1CleanupError(f"network error: {exc}") from exc
    except json.JSONDecodeError as exc:
        raise D1CleanupError(f"invalid json response: {exc}") from exc

    if not bool(payload.get("success", False)):
        _raise_for_payload(payload)
    return payload


def _first_result(payload: dict[str, Any]) -> dict[str, Any]:
    results = payload.get("result")
    if not isinstance(results, list) or not results or not isinstance(results[0], dict):
        raise D1CleanupError("missing result payload")
    result = results[0]
    if not bool(result.get("success", True)):
        _raise_for_payload(payload)
    return result


def _query_once(database_id: str, sql: str, params: list[Any] | None = None) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    payload = _query(database_id, sql, params=params)
    result = _first_result(payload)
    rows = result.get("results")
    if not isinstance(rows, list):
        rows = []
    meta = result.get("meta")
    if not isinstance(meta, dict):
        meta = {}
    normalized_rows = [row for row in rows if isinstance(row, dict)]
    return normalized_rows, meta


def _count_rows(database_id: str, table: str) -> tuple[int, int | None]:
    rows, meta = _query_once(database_id, f"SELECT COUNT(*) AS count FROM {table}")
    count = 0
    if rows:
        try:
            count = int(rows[0].get("count") or 0)
        except Exception:
            count = 0
    size_after = meta.get("size_after")
    try:
        return count, int(size_after) if size_after is not None else None
    except Exception:
        return count, None


def _delete_in_batches(database_id: str, table: str, retention_hours: int, batch_size: int) -> tuple[int, int | None]:
    total_deleted = 0
    latest_size_after: int | None = None
    safe_retention = max(0, int(retention_hours))
    safe_batch_size = max(1, int(batch_size))
    sql = (
        f"DELETE FROM {table} "
        f"WHERE created_at < datetime('now', '-{safe_retention} hours') "
        f"LIMIT {safe_batch_size}"
    )
    while True:
        _rows, meta = _query_once(database_id, sql)
        changes_raw = meta.get("changes")
        try:
            changes = int(changes_raw or 0)
        except Exception:
            changes = 0
        size_after = meta.get("size_after")
        try:
            latest_size_after = int(size_after) if size_after is not None else latest_size_after
        except Exception:
            pass
        if changes <= 0:
            break
        total_deleted += changes
    return total_deleted, latest_size_after


def _final_size_after(database_id: str) -> int | None:
    _rows, meta = _query_once(database_id, "SELECT 1 AS ok")
    size_after = meta.get("size_after")
    try:
        return int(size_after) if size_after is not None else None
    except Exception:
        return None


def d1_cleanup_once(
    database_id: str = DEFAULT_D1_DATABASE_ID,
    mail_retention_hours: int = DEFAULT_D1_MAIL_RETENTION_HOURS,
    address_retention_hours: int = DEFAULT_D1_ADDRESS_RETENTION_HOURS,
    batch_size: int = DEFAULT_D1_CLEANUP_BATCH_SIZE,
) -> dict[str, object]:
    summary: dict[str, object] = {
        "deleted_mails": 0,
        "deleted_addresses": 0,
        "deleted_senders": 0,
        "size_after_bytes": None,
        "skipped_reason": None,
    }
    normalized_database_id = str(database_id or "").strip()
    if not normalized_database_id:
        summary["skipped_reason"] = "missing_database_id"
        return summary

    if _credentials_from_env() is None:
        summary["skipped_reason"] = "missing_cloudflare_credentials"
        return summary

    size_after_bytes: int | None = None
    try:
        raw_mails_count, size_after_bytes = _count_rows(normalized_database_id, "raw_mails")
    except D1TableMissingError:
        _warning("table raw_mails not found, skip count")
        raw_mails_count = 0
    try:
        address_count, count_size_after = _count_rows(normalized_database_id, "address")
        if count_size_after is not None:
            size_after_bytes = count_size_after
    except D1TableMissingError:
        _warning("table address not found, skip count")
        address_count = 0

    if raw_mails_count == 0 and address_count == 0:
        summary["size_after_bytes"] = size_after_bytes
        summary["skipped_reason"] = "nothing_to_clean"
        print(f"[{now()}] [d1_cleanup] nothing to clean")
        return summary

    try:
        deleted_mails, delete_size_after = _delete_in_batches(
            normalized_database_id,
            "raw_mails",
            retention_hours=mail_retention_hours,
            batch_size=batch_size,
        )
        summary["deleted_mails"] = deleted_mails
        if delete_size_after is not None:
            size_after_bytes = delete_size_after
    except D1TableMissingError:
        _warning("table raw_mails not found, skip cleanup")

    try:
        deleted_addresses, delete_size_after = _delete_in_batches(
            normalized_database_id,
            "address",
            retention_hours=address_retention_hours,
            batch_size=batch_size,
        )
        summary["deleted_addresses"] = deleted_addresses
        if delete_size_after is not None:
            size_after_bytes = delete_size_after
    except D1TableMissingError:
        _warning("table address not found, skip cleanup")

    try:
        deleted_senders, delete_size_after = _delete_in_batches(
            normalized_database_id,
            "address_sender",
            retention_hours=address_retention_hours,
            batch_size=batch_size,
        )
        summary["deleted_senders"] = deleted_senders
        if delete_size_after is not None:
            size_after_bytes = delete_size_after
    except D1TableMissingError:
        _warning("table address_sender not found, skip cleanup")

    try:
        final_size_after = _final_size_after(normalized_database_id)
        if final_size_after is not None:
            size_after_bytes = final_size_after
    except D1CleanupError as exc:
        _warning(f"final size check failed: {exc}")

    summary["size_after_bytes"] = size_after_bytes
    size_mb_text = "unknown"
    if isinstance(size_after_bytes, int):
        size_mb_text = f"{size_after_bytes / (1024 * 1024):.1f}MB"
    print(
        f"[{now()}] [d1_cleanup] 清理完成 | raw_mails=-{summary['deleted_mails']} "
        f"| address=-{summary['deleted_addresses']} | address_sender=-{summary['deleted_senders']} "
        f"| size={size_mb_text}"
    )
    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="清理 cfmail Cloudflare D1 历史数据")
    parser.add_argument("--once", action="store_true", help="只执行一轮")
    parser.add_argument("--interval", type=int, default=1800, help="清理间隔秒数")
    parser.add_argument("--database-id", default=DEFAULT_D1_DATABASE_ID, help="Cloudflare D1 database id")
    parser.add_argument(
        "--mail-retention-hours",
        type=int,
        default=DEFAULT_D1_MAIL_RETENTION_HOURS,
        help="raw_mails 保留小时数",
    )
    parser.add_argument(
        "--address-retention-hours",
        type=int,
        default=DEFAULT_D1_ADDRESS_RETENTION_HOURS,
        help="address / address_sender 保留小时数",
    )
    args = parser.parse_args()

    interval = max(1, int(args.interval))
    while True:
        started_at = time.time()
        try:
            d1_cleanup_once(
                database_id=str(args.database_id).strip(),
                mail_retention_hours=int(args.mail_retention_hours),
                address_retention_hours=int(args.address_retention_hours),
            )
        except Exception as exc:
            _warning(f"cleanup failed: {exc}")
        elapsed = time.time() - started_at
        if args.once:
            break
        time.sleep(max(0, interval - elapsed))


if __name__ == "__main__":
    main()