tencent-baidu-tracking/scripts/honor_backfill_oaid.py

#!/usr/bin/env python3
import argparse
import glob
import hashlib
import json
import re
import socket
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import pymysql


TRACE_ID_RE = re.compile(r"traceId=([^\s]+)")
TRACK_ID_RE = re.compile(r"trackId=([^,\s}]+)")
OAID_RE = re.compile(r"oaid=([^,\s}]+)")


@dataclass
class LogEntry:
    trace_id: str
    oaid: str


def md5_hex(value: str) -> str:
    return hashlib.md5(value.encode("utf-8")).hexdigest()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Backfill honor oaid/oaidMd5 into TiDB cold data and Redis hot data from logs."
    )
    parser.add_argument(
        "--log",
        action="append",
        required=True,
        help="Log file path or glob. Repeatable, e.g. --log 'logs/tencent-baidu-tracking.2026-07-2[56].log'",
    )
    parser.add_argument("--mysql-host", required=True)
    parser.add_argument("--mysql-port", type=int, default=3390)
    parser.add_argument("--mysql-user", required=True)
    parser.add_argument("--mysql-password", required=True)
    parser.add_argument("--mysql-db", default="adx_tencent")
    parser.add_argument("--redis-host", required=True)
    parser.add_argument("--redis-port", type=int, default=6379)
    parser.add_argument("--redis-user", default="")
    parser.add_argument("--redis-password", required=True)
    parser.add_argument("--redis-db", type=int, default=5)
    parser.add_argument("--redis-prefix", default="adx:honor:")
    parser.add_argument("--limit", type=int, default=0, help="Only process first N extracted mappings")
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def expand_logs(patterns: Iterable[str]) -> List[Path]:
    files: List[Path] = []
    for pattern in patterns:
        matched = sorted(Path(p) for p in glob.glob(pattern))
        if matched:
            files.extend(matched)
        else:
            path = Path(pattern)
            if path.exists():
                files.append(path)
    deduped: List[Path] = []
    seen = set()
    for path in files:
        key = str(path.resolve())
        if key not in seen:
            seen.add(key)
            deduped.append(path)
    return deduped


def scan_logs(paths: Iterable[Path]) -> Dict[str, LogEntry]:
    result: Dict[str, LogEntry] = {}
    for path in paths:
        with path.open("r", encoding="utf-8", errors="ignore") as fh:
            for line in fh:
                if "[荣耀][曝光][入参]" not in line and "[荣耀][点击][入参]" not in line:
                    continue
                trace_match = TRACE_ID_RE.search(line)
                oaid_match = OAID_RE.search(line)
                if not trace_match or not oaid_match:
                    track_match = TRACK_ID_RE.search(line)
                    if not track_match or not oaid_match:
                        continue
                    trace_id = track_match.group(1)
                else:
                    trace_id = trace_match.group(1)
                oaid = oaid_match.group(1).strip()
                if not oaid:
                    continue
                result[trace_id] = LogEntry(trace_id=trace_id, oaid=oaid)
    return result


class RespRedisClient:
    def __init__(self, host: str, port: int, user: str, password: str, db: int):
        self.sock = socket.create_connection((host, port), timeout=10)
        self.file = self.sock.makefile("rb")
        self._auth(user, password)
        self._select(db)

    def close(self) -> None:
        try:
            self.file.close()
        finally:
            self.sock.close()

    def _write_command(self, *parts: str) -> None:
        payload = [f"*{len(parts)}\r\n".encode("utf-8")]
        for part in parts:
            data = part.encode("utf-8")
            payload.append(f"${len(data)}\r\n".encode("utf-8"))
            payload.append(data + b"\r\n")
        self.sock.sendall(b"".join(payload))

    def _read_response(self):
        prefix = self.file.read(1)
        if not prefix:
            raise RuntimeError("redis connection closed")
        if prefix == b"+":
            return self.file.readline().decode("utf-8").rstrip("\r\n")
        if prefix == b"-":
            raise RuntimeError(self.file.readline().decode("utf-8").rstrip("\r\n"))
        if prefix == b":":
            return int(self.file.readline().decode("utf-8").rstrip("\r\n"))
        if prefix == b"$":
            length = int(self.file.readline().decode("utf-8").rstrip("\r\n"))
            if length == -1:
                return None
            data = self.file.read(length)
            self.file.read(2)
            return data.decode("utf-8")
        if prefix == b"*":
            length = int(self.file.readline().decode("utf-8").rstrip("\r\n"))
            if length == -1:
                return None
            return [self._read_response() for _ in range(length)]
        raise RuntimeError(f"unsupported redis response prefix: {prefix!r}")

    def _auth(self, user: str, password: str) -> None:
        if user:
            self._write_command("AUTH", user, password)
        else:
            self._write_command("AUTH", password)
        self._read_response()

    def _select(self, db: int) -> None:
        self._write_command("SELECT", str(db))
        self._read_response()

    def get(self, key: str) -> Optional[str]:
        self._write_command("GET", key)
        return self._read_response()

    def pttl(self, key: str) -> int:
        self._write_command("PTTL", key)
        return int(self._read_response())

    def set_px(self, key: str, value: str, px: int) -> None:
        self._write_command("SET", key, value, "PX", str(px))
        self._read_response()


def fetch_rows(conn, trace_ids: List[str]) -> Dict[str, dict]:
    rows: Dict[str, dict] = {}
    chunk_size = 500
    sql = """
        SELECT qk, media_trace_id, media_params
        FROM honor_ad_bid_events
        WHERE media_trace_id IN ({})
    """
    with conn.cursor() as cur:
        for i in range(0, len(trace_ids), chunk_size):
            chunk = trace_ids[i:i + chunk_size]
            placeholders = ",".join(["%s"] * len(chunk))
            cur.execute(sql.format(placeholders), chunk)
            for qk, media_trace_id, media_params in cur.fetchall():
                rows[media_trace_id] = {
                    "qk": qk,
                    "trace_id": media_trace_id,
                    "media_params": json.loads(media_params) if media_params else {},
                }
    return rows


def update_mysql(conn, trace_id: str, params: dict, dry_run: bool) -> bool:
    sql = """
        UPDATE honor_ad_bid_events
        SET media_params = %s
        WHERE media_trace_id = %s
    """
    if dry_run:
        return True
    with conn.cursor() as cur:
        cur.execute(sql, (json.dumps(params, ensure_ascii=False), trace_id))
    return True


def patch_redis_payload(payload: str, oaid: str, oaid_md5: str) -> Tuple[str, bool]:
    data = json.loads(payload)
    media_params = data.get("mediaParams") or {}
    changed = False
    if media_params.get("oaid") != oaid:
        media_params["oaid"] = oaid
        changed = True
    if media_params.get("oaidMd5") != oaid_md5:
        media_params["oaidMd5"] = oaid_md5
        changed = True
    data["mediaParams"] = media_params
    return json.dumps(data, ensure_ascii=False), changed


def main() -> int:
    args = parse_args()
    log_files = expand_logs(args.log)
    if not log_files:
        print("no log files matched")
        return 1

    mappings = scan_logs(log_files)
    if not mappings:
        print("no honor traceId/oaid mappings found")
        return 1
    items = list(mappings.items())
    if args.limit > 0:
        items = items[:args.limit]
    trace_ids = [trace_id for trace_id, _ in items]

    conn = pymysql.connect(
        host=args.mysql_host,
        port=args.mysql_port,
        user=args.mysql_user,
        password=args.mysql_password,
        database=args.mysql_db,
        charset="utf8mb4",
        autocommit=False,
    )
    redis_client = RespRedisClient(
        host=args.redis_host,
        port=args.redis_port,
        user=args.redis_user,
        password=args.redis_password,
        db=args.redis_db,
    )

    try:
        db_rows = fetch_rows(conn, trace_ids)
        stats = {
            "log_mappings": len(items),
            "db_found": 0,
            "mysql_updated": 0,
            "redis_keys_checked": 0,
            "redis_keys_updated": 0,
            "already_ok": 0,
            "db_missing": 0,
            "redis_missing": 0,
        }

        for trace_id, entry in items:
            row = db_rows.get(trace_id)
            if not row:
                stats["db_missing"] += 1
                continue
            stats["db_found"] += 1
            oaid_md5 = md5_hex(entry.oaid)

            media_params = row["media_params"] or {}
            mysql_changed = False
            if media_params.get("oaid") != entry.oaid:
                media_params["oaid"] = entry.oaid
                mysql_changed = True
            if media_params.get("oaidMd5") != oaid_md5:
                media_params["oaidMd5"] = oaid_md5
                mysql_changed = True
            if mysql_changed:
                update_mysql(conn, trace_id, media_params, args.dry_run)
                stats["mysql_updated"] += 1
            else:
                stats["already_ok"] += 1

            redis_keys = [
                f"{args.redis_prefix}bid:{row['qk']}",
                f"{args.redis_prefix}media:honor:{trace_id}",
            ]
            for key in redis_keys:
                stats["redis_keys_checked"] += 1
                payload = redis_client.get(key)
                if payload is None:
                    stats["redis_missing"] += 1
                    continue
                ttl_ms = redis_client.pttl(key)
                if ttl_ms <= 0:
                    stats["redis_missing"] += 1
                    continue
                new_payload, changed = patch_redis_payload(payload, entry.oaid, oaid_md5)
                if changed:
                    if not args.dry_run:
                        redis_client.set_px(key, new_payload, ttl_ms)
                    stats["redis_keys_updated"] += 1

        if args.dry_run:
            conn.rollback()
        else:
            conn.commit()
        print(json.dumps(stats, ensure_ascii=False, indent=2))
        return 0
    finally:
        redis_client.close()
        conn.close()


if __name__ == "__main__":
    raise SystemExit(main())