stat_center/script/sync_jump_tab_avatar.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Mirror jump_tab.avatar images from http://yuan.bbaazz.me to AWS S3 (bucket: meirenyu).

Behaviour:
  - reads bili_app.jump_tab from the source MongoDB
  - for every doc whose `avatar` is on yuan.bbaazz.me, downloads the file and
    uploads it to S3 using the same URL path as the object key
  - DOES NOT modify MongoDB
  - skips objects that already exist in the bucket (HEAD check)

Run:
    python3 script/sync_jump_tab_avatar.py            # dry-run, prints actions only
    python3 script/sync_jump_tab_avatar.py --apply    # actually download + upload
"""

from __future__ import annotations

import argparse
import mimetypes
import sys
from urllib.parse import urlparse

import boto3
import requests
from botocore.exceptions import ClientError
from pymongo import MongoClient

# ---------- config ---------------------------------------------------------

MONGO_URI = (
    "mongodb://bili_app:SbKE9KN55fptDXeNaCDX@"
    "172.104.52.130:27017,172.104.52.26:27017,172.104.52.242:27017/"
    "bili_app?replicaSet=fish0&maxPoolSize=200&authSource=admin"
)
DB_NAME = "bili_app"
COLL = "jump_tab"
AVATAR_FIELD = "avatar"

SOURCE_HOST = "yuan.bbaazz.me"
SOURCE_BASE = f"http://{SOURCE_HOST}"

AWS_REGION = "ap-east-1"
AWS_ACCESS_KEY = "AKIA4FIZIQT3RUPZY3AR"
AWS_SECRET_KEY = "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"
S3_BUCKET = "meirenyu"

HTTP_TIMEOUT = 30
# ---------------------------------------------------------------------------


def avatar_to_key(avatar: str) -> str:
    """Return the S3 key for an avatar value.

    Handles both relative paths ("v3/image/foo.png") and absolute URLs
    ("http://yuan.bbaazz.me/v3/image/foo.png").
    """
    if not isinstance(avatar, str) or not avatar.strip():
        return ""
    s = avatar.strip()
    if s.startswith("http://") or s.startswith("https://"):
        return urlparse(s).path.lstrip("/")
    return s.lstrip("/")


def avatar_to_source_url(avatar: str) -> str:
    """Return the URL to download from yuan.bbaazz.me."""
    s = avatar.strip()
    if s.startswith("http://") or s.startswith("https://"):
        return s
    return f"{SOURCE_BASE}/{s.lstrip('/')}"


def guess_content_type(url: str, fallback: str | None) -> str:
    ct, _ = mimetypes.guess_type(url)
    return ct or fallback or "application/octet-stream"


def s3_object_exists(s3, key: str) -> bool:
    try:
        s3.head_object(Bucket=S3_BUCKET, Key=key)
        return True
    except ClientError as e:
        if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey", "NotFound"):
            return False
        raise


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--apply", action="store_true",
                        help="actually perform downloads and uploads")
    parser.add_argument("--limit", type=int, default=0,
                        help="process at most N matching docs (0 = no limit)")
    parser.add_argument("--inspect", action="store_true",
                        help="print sample docs and host distribution, then exit")
    args = parser.parse_args()

    mc = MongoClient(MONGO_URI)
    coll = mc[DB_NAME][COLL]

    if args.inspect:
        total = coll.estimated_document_count()
        print(f"collection={DB_NAME}.{COLL}  estimated_count={total}")
        print("--- 3 sample docs ---")
        for d in coll.find({}, limit=3):
            print(d)
        print("--- avatar field stats ---")
        with_avatar = coll.count_documents({AVATAR_FIELD: {"$exists": True}})
        print(f"docs with `{AVATAR_FIELD}` field: {with_avatar}")
        # bucket hosts
        from collections import Counter
        hosts = Counter()
        for d in coll.find({AVATAR_FIELD: {"$exists": True}},
                           {AVATAR_FIELD: 1}, limit=2000):
            v = d.get(AVATAR_FIELD)
            if isinstance(v, str) and v:
                try:
                    h = urlparse(v).hostname or "(no-host)"
                except ValueError:
                    h = "(parse-error)"
                hosts[h] += 1
            else:
                hosts[f"({type(v).__name__})"] += 1
        for h, n in hosts.most_common(20):
            print(f"  {n:>6}  {h}")
        mc.close()
        return 0

    s3 = boto3.client(
        "s3",
        region_name=AWS_REGION,
        aws_access_key_id=AWS_ACCESS_KEY,
        aws_secret_access_key=AWS_SECRET_KEY,
    )

    # Every doc with a non-empty avatar is a candidate (per user: store as relative path).
    cursor = coll.find(
        {AVATAR_FIELD: {"$exists": True, "$ne": ""}},
        {AVATAR_FIELD: 1},
        no_cursor_timeout=True,
    )

    seen, hit, uploaded, skipped, failed = 0, 0, 0, 0, 0
    sess = requests.Session()

    try:
        for doc in cursor:
            seen += 1
            avatar = doc.get(AVATAR_FIELD, "")
            key = avatar_to_key(avatar)
            if not key:
                continue
            hit += 1
            src = avatar_to_source_url(avatar)

            if args.apply and s3_object_exists(s3, key):
                print(f"[exists] s3://{S3_BUCKET}/{key}")
                skipped += 1
                if args.limit and hit >= args.limit:
                    break
                continue

            if not args.apply:
                print(f"[dry-run] {src}  ->  s3://{S3_BUCKET}/{key}")
                if args.limit and hit >= args.limit:
                    break
                continue

            try:
                r = sess.get(src, timeout=HTTP_TIMEOUT, stream=True)
                r.raise_for_status()
                ctype = guess_content_type(src, r.headers.get("Content-Type"))
                body = r.content
                s3.put_object(
                    Bucket=S3_BUCKET,
                    Key=key,
                    Body=body,
                    ContentType=ctype,
                )
                uploaded += 1
                print(f"[ok] {src}  ->  s3://{S3_BUCKET}/{key}  "
                      f"({len(body)} bytes, {ctype})")
            except Exception as e:
                failed += 1
                print(f"[fail] {src}: {e}", file=sys.stderr)

            if args.limit and hit >= args.limit:
                break
    finally:
        cursor.close()
        mc.close()

    print()
    print(f"scanned={seen}  matched={hit}  uploaded={uploaded}  "
          f"skipped={skipped}  failed={failed}  "
          f"mode={'APPLY' if args.apply else 'DRY-RUN'}")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())