cfc/scripts/backfill-oss-to-db.py

#!/usr/bin/env python3
"""
将京东云 OSS 上的所有文件回写进 file_record 表。

策略：
1. 遍历 OSS bucket 所有对象（分页）
2. 并发下载每个对象内容，计算 SHA-256
3. 写入 file_record 表（hash 唯一，已存在则跳过）

用法:
    python3 scripts/backfill-oss-to-db.py
    python3 scripts/backfill-oss-to-db.py --dry-run        # 仅统计不上传
    python3 scripts/backfill-oss-to-db.py --workers 20     # 并发数
"""

import boto3
import hashlib
import pymysql
from botocore.config import Config
from concurrent.futures import ThreadPoolExecutor, as_completed
import argparse
import logging
import threading
from datetime import datetime

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("backfill")

ENDPOINT = "s3.cn-north-1.jdcloud-oss.com"
ACCESS_KEY = "JDC_F351EC3CB1F8593204ABCDA2CBF2"
SECRET_KEY = "970C0C6985A6BF29D0C79AE3693B9899"
BUCKET = "cfc"
REGION = "cn-north-1"

DB_HOST = "mysql-internet-cn-north-1-23feae22680e4dfa.rds.jdcloud.com"
DB_PORT = 3306
DB_USER = "cfc"
DB_PASS = "cfc@1314"
DB_NAME = "zxyj"

PUBLIC_URL_PREFIX = f"https://{BUCKET}.{ENDPOINT}/"

# 测试目录，跳过
SKIP_PREFIXES = ("test-connectivity",)

db_lock = threading.Lock()


def get_s3():
    return boto3.client(
        "s3",
        endpoint_url=f"https://{ENDPOINT}",
        aws_access_key_id=ACCESS_KEY,
        aws_secret_access_key=SECRET_KEY,
        config=Config(s3={"addressing_style": "path"}, signature_version="s3v4"),
        region_name=REGION,
    )


def list_all_objects(s3):
    """列出所有对象，返回 [{Key, Size, ContentType, LastModified}]"""
    objects = []
    marker = None
    while True:
        if marker:
            resp = s3.list_objects_v2(Bucket=BUCKET, MaxKeys=1000, StartAfter=marker)
        else:
            resp = s3.list_objects_v2(Bucket=BUCKET, MaxKeys=1000)
        if "Contents" not in resp:
            break
        for obj in resp["Contents"]:
            key = obj["Key"]
            if key.startswith(SKIP_PREFIXES):
                continue
            objects.append({"Key": key, "Size": obj["Size"], "LastModified": obj["LastModified"]})
            marker = key
        if not resp.get("IsTruncated"):
            break
    return objects


def guess_content_type(key):
    """根据 key 扩展名猜测 Content-Type"""
    ext = key[key.rfind(".") + 1:].lower() if "." in key else ""
    mapping = {
        "jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png",
        "gif": "image/gif", "webp": "image/webp", "bmp": "image/bmp",
        "svg": "image/svg+xml", "pdf": "application/pdf",
        "mp4": "video/mp4", "avi": "video/x-msvideo", "mov": "video/quicktime",
        "mp3": "audio/mpeg", "wav": "audio/wav", "m4a": "audio/mp4",
        "txt": "text/plain", "json": "application/json", "html": "text/html",
        "doc": "application/msword", "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "xls": "application/vnd.ms-excel", "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        "zip": "application/zip",
    }
    return mapping.get(ext, "application/octet-stream")


def guess_file_type(content_type):
    if content_type.startswith("image/"):
        return "image"
    if content_type == "application/pdf" or content_type.endswith("pdf"):
        return "pdf"
    if content_type.startswith("video/"):
        return "video"
    if content_type.startswith("audio/"):
        return "audio"
    return "other"


def guess_category(key):
    """根据 key 路径推断分类"""
    parts = key.split("/")
    if len(parts) >= 3 and parts[1] == "health-reports":
        return "health-reports"
    if len(parts) >= 2:
        return parts[0]
    return "general"


def process_object(s3, obj, db):
    """下载对象、计算 hash、写库。返回 ('ok'|'skip'|'fail', key)"""
    key = obj["Key"]
    url = PUBLIC_URL_PREFIX + key
    content_type = guess_content_type(key)
    file_type = guess_file_type(content_type)
    category = guess_category(key)
    original_name = key.split("/")[-1]

    try:
        # 下载并计算 SHA-256
        resp = s3.get_object(Bucket=BUCKET, Key=key)
        sha = hashlib.sha256()
        size = 0
        body = resp["Body"]
        while True:
            chunk = body.read(65536)
            if not chunk:
                break
            sha.update(chunk)
            size += len(chunk)
        file_hash = sha.hexdigest()

        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        with db_lock:
            with db.cursor() as cursor:
                cursor.execute("SELECT id FROM file_record WHERE hash = %s", (file_hash,))
                if cursor.fetchone():
                    return "skip", key
                cursor.execute(
                    "INSERT INTO file_record (hash, url, file_size, content_type, "
                    "original_filename, category, file_type, created_at, updated_at) "
                    "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
                    (file_hash, url, size, content_type, original_name,
                     category, file_type, now, now),
                )
            db.commit()
        return "ok", key
    except Exception as e:
        log.error("[FAIL] %s: %s", key, e)
        return "fail", key


def main():
    parser = argparse.ArgumentParser(description="回写 OSS 文件到 file_record 表")
    parser.add_argument("--dry-run", action="store_true", help="仅统计")
    parser.add_argument("--workers", type=int, default=20, help="并发数")
    parser.add_argument("--limit", type=int, default=0, help="处理前 N 个（测试用，0=全部）")
    args = parser.parse_args()

    s3 = get_s3()
    log.info("连接 OSS 成功，正在列举对象...")
    objects = list_all_objects(s3)
    log.info("OSS 对象总数: %d", len(objects))

    if args.limit > 0:
        objects = objects[:args.limit]
        log.info("限制处理前 %d 个", args.limit)

    if args.dry_run:
        log.info("DRY-RUN 模式，仅统计不写库")
        by_type = {}
        for o in objects:
            ct = guess_file_type(guess_content_type(o["Key"]))
            by_type[ct] = by_type.get(ct, 0) + 1
        total_size = sum(o["Size"] for o in objects)
        log.info("文件类型分布: %s", by_type)
        log.info("文件总大小: %.1f MB", total_size / 1024 / 1024)
        return

    db = pymysql.connect(
        host=DB_HOST, port=DB_PORT, user=DB_USER, password=DB_PASS,
        database=DB_NAME, charset="utf8mb4",
    )
    log.info("MySQL 连接成功")

    ok = skip = fail = 0
    try:
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            futures = [pool.submit(process_object, s3, obj, db) for obj in objects]
            for i, fut in enumerate(as_completed(futures), 1):
                status, key = fut.result()
                if status == "ok":
                    ok += 1
                elif status == "skip":
                    skip += 1
                else:
                    fail += 1
                if i % 200 == 0:
                    log.info("进度: %d/%d (ok=%d skip=%d fail=%d)",
                             i, len(objects), ok, skip, fail)
    finally:
        db.close()

    log.info("=" * 50)
    log.info("回写完成: 新增=%d, 跳过=%d, 失败=%d", ok, skip, fail)


if __name__ == "__main__":
    main()