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 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}]+)")


class LogEntry:
    def __init__(self, trace_id: str, oaid: str):
        self.trace_id = trace_id
        self.oaid = oaid


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(
        "--progress-every",
        type=int,
        default=100000,
        help="Print progress every N scanned log lines / processed 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], progress_every: int) -> Dict[str, LogEntry]:
    result: Dict[str, LogEntry] = {}
    for path in paths:
        print(f"[scan] start file={path}")
        line_count = 0
        match_count = 0
        with path.open("r", encoding="utf-8", errors="ignore") as fh:
            for line in fh:
                line_count += 1
                if progress_every > 0 and line_count % progress_every == 0:
                    print(
                        f"[scan] file={path} lines={line_count} mappings={len(result)} fileMatches={match_count}"
                    )
                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)
                match_count += 1
        print(f"[scan] done file={path} lines={line_count} mappings={len(result)} fileMatches={match_count}")
    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], progress_every: int) -> 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 = 'honor' AND media_trace_id IN ({})
    """
    total_chunks = (len(trace_ids) + chunk_size - 1) // chunk_size
    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))
            chunk_no = i // chunk_size + 1
            if chunk_no == 1 or chunk_no == total_chunks or (progress_every > 0 and chunk_no % max(progress_every // chunk_size, 1) == 0):
                print(
                    f"[db] fetching chunk={chunk_no}/{total_chunks} chunkSize={len(chunk)} rowsLoaded={len(rows)}"
                )
            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 {},
                }
    print(f"[db] fetched rows={len(rows)}")
    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 = 'honor' AND 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

    print(
        f"[start] logs={len(log_files)} dry_run={args.dry_run} limit={args.limit or 'all'} progress_every={args.progress_every}"
    )
    for path in log_files:
        print(f"[start] log={path}")

    mappings = scan_logs(log_files, args.progress_every)
    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]
    print(f"[scan] extracted_mappings={len(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, args.progress_every)
        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 index, (trace_id, entry) in enumerate(items, start=1):
            if args.progress_every > 0 and index % args.progress_every == 0:
                print(
                    f"[apply] processed={index}/{len(items)} mysql_updated={stats['mysql_updated']} "
                    f"redis_updated={stats['redis_keys_updated']} db_missing={stats['db_missing']}"
                )
            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(f"[done] processed={len(items)}")
        print(json.dumps(stats, ensure_ascii=False, indent=2))
        return 0
    finally:
        redis_client.close()
        conn.close()


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