ZhiYuanTong/scripts/import_gaokao_compass_2025.py

import argparse
import csv
import hashlib
import json
import os
import subprocess
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


YEAR = 2025
DATASET = "choucsan/Gaokao-Compass-11M"
REVISION = "main"
TREE_URL = f"https://huggingface.co/api/datasets/{DATASET}/tree/{REVISION}/data/{YEAR}?recursive=true&expand=false&limit=1000"
LANDING_URL = f"https://huggingface.co/datasets/{DATASET}"
USER_AGENT = "Mozilla/5.0 (compatible; ZhiyuantongDataAudit/1.0)"

PROVINCES = {
    "anhui": "安徽省", "beijing": "北京市", "chongqing": "重庆市", "fujian": "福建省",
    "gansu": "甘肃省", "guangdong": "广东省", "guangxi": "广西壮族自治区",
    "guizhou": "贵州省", "hainan": "海南省", "hebei": "河北省",
    "heilongjiang": "黑龙江省", "henan": "河南省", "hubei": "湖北省", "hunan": "湖南省",
    "jiangsu": "江苏省", "jiangxi": "江西省", "jilin": "吉林省", "liaoning": "辽宁省",
    "neimenggu": "内蒙古自治区", "ningxia": "宁夏回族自治区", "qinghai": "青海省",
    "shaanxi": "陕西省", "shandong": "山东省", "shanghai": "上海市", "shanxi": "山西省",
    "sichuan": "四川省", "tianjin": "天津市", "xinjiang": "新疆维吾尔自治区",
    "xizang": "西藏自治区", "yunnan": "云南省", "zhejiang": "浙江省",
}

EXPECTED_HEADERS = {
    "enrollment-plan.csv": [
        "province", "year", "batch", "category", "university_code", "university_name",
        "major_group", "major_code", "major_name", "major_note", "subject_req", "plan_count",
        "duration", "tuition",
    ],
    "school-admission.csv": [
        "province", "year", "category", "batch", "university_code", "university_name",
        "major_group", "subject_req", "min_score", "min_rank", "control_score", "score_diff",
        "admit_count", "school_province", "school_nature", "is_985", "is_211",
    ],
    "major-admission.csv": [
        "province", "year", "category", "batch", "university_code", "university_name",
        "major_group", "major_code", "major_name", "major_note", "subject_req", "min_score",
        "min_rank", "max_score", "avg_score", "admit_count", "school_province", "school_nature",
        "is_985", "is_211",
    ],
    "score-range.csv": [
        "province", "year", "category", "batch", "control_score", "score", "score_range",
        "segment_count", "cumulative_count", "rank_range",
    ],
}

HEADER_VARIANTS = {
    ("fujian", "enrollment-plan.csv"): [
        "province", "year", "batch", "category", "university_code", "university_name",
        "major_group", "major_group", "major_code", "major_name", "major_note", "subject_req",
        "plan_count", "duration", "tuition",
    ],
    ("shanghai", "enrollment-plan.csv"): [
        "province", "year", "batch", "category", "university_code", "university_name",
        "major_group", "major_group", "major_code", "major_name", "major_note", "subject_req",
        "plan_count", "duration", "tuition",
    ],
    ("guangdong", "enrollment-plan.csv"): [
        "province", "year", "batch", "category", "university_code", "university_name",
        "major_group", "major_code", "major_name", "major_note", "major_note", "subject_req",
        "plan_count", "duration", "tuition",
    ],
}

STAGING = {
    "enrollment-plan.csv": (
        "stg_gc_enrollment_plans",
        "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw,"
        "group_code,major_code,major_name,major_note,subject_requirements,plan_count_text,study_duration,tuition_text",
    ),
    "school-admission.csv": (
        "stg_gc_school_admissions",
        "@raw_province,@raw_year,subject_category,batch_name,source_school_code,school_name_raw,"
        "group_code,subject_requirements,min_score_text,min_rank_text,control_score_text,score_diff_text,"
        "admit_count_text,school_province,school_nature,is_985_text,is_211_text",
    ),
    "major-admission.csv": (
        "stg_gc_major_admissions",
        "@raw_province,@raw_year,subject_category,batch_name,source_school_code,school_name_raw,"
        "group_code,major_code,major_name,major_note,subject_requirements,min_score_text,min_rank_text,"
        "max_score_text,avg_score_text,admit_count_text,school_province,school_nature,is_985_text,is_211_text",
    ),
    "score-range.csv": (
        "stg_gc_score_ranges",
        "@raw_province,@raw_year,subject_category,batch_name,control_score_text,score_text,score_range_text,"
        "segment_count_text,cumulative_count_text,rank_range_text",
    ),
}


def sql_string(value):
    if value is None:
        return "NULL"
    text = str(value).replace("\\", "\\\\").replace("'", "\\'")
    text = text.replace("\x00", "").replace("\r", " ").replace("\n", " ")
    return f"'{text}'"


def request_bytes(url, timeout=180):
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    last_error = None
    for attempt in range(3):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                return response.read()
        except Exception as error:  # noqa: BLE001 - 重试后由调用方报告完整URL
            last_error = error
            if attempt < 2:
                time.sleep(2 ** attempt)
    raise RuntimeError(f"下载失败: {url}: {last_error}")


def sha256(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def download_one(item, raw_root, refresh):
    relative = Path(item["path"]).relative_to(f"data/{YEAR}")
    target = raw_root / relative
    target.parent.mkdir(parents=True, exist_ok=True)
    expected_size = int(item["size"])
    if not refresh and target.exists() and target.stat().st_size == expected_size:
        return target
    url = f"https://huggingface.co/datasets/{DATASET}/resolve/{REVISION}/{item['path']}"
    content = request_bytes(url)
    if len(content) != expected_size:
        raise RuntimeError(f"文件长度异常: {item['path']}，期望{expected_size}，实际{len(content)}")
    partial = target.with_suffix(target.suffix + ".part")
    partial.write_bytes(content)
    os.replace(partial, target)
    return target


def csv_summary(path):
    expected = HEADER_VARIANTS.get((path.parent.name, path.name), EXPECTED_HEADERS[path.name])
    row_count = 0
    invalid_width = 0
    with path.open("r", encoding="utf-8-sig", newline="") as stream:
        reader = csv.reader(stream)
        header = next(reader, None)
        if header != expected:
            raise RuntimeError(f"字段不匹配: {path}\n期望: {expected}\n实际: {header}")
        for row in reader:
            row_count += 1
            if len(row) != len(expected):
                invalid_width += 1
    if invalid_width:
        raise RuntimeError(f"CSV列数异常: {path}，异常行数{invalid_width}")
    return row_count


def load_mapping(path):
    slug = path.parent.name
    if path.name == "enrollment-plan.csv" and slug in {"fujian", "shanghai"}:
        return (
            "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw,"
            "@composite_group_code,group_code,major_code,major_name,major_note,subject_requirements,"
            "plan_count_text,study_duration,tuition_text",
            "",
        )
    if path.name == "enrollment-plan.csv" and slug == "guangdong":
        return (
            "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw,"
            "group_code,major_code,major_name,@major_note_one,@major_note_two,subject_requirements,"
            "plan_count_text,study_duration,tuition_text",
            ", major_note=NULLIF(CONCAT_WS('；', NULLIF(TRIM(@major_note_one), ''), NULLIF(TRIM(@major_note_two), '')), '')",
        )
    return STAGING[path.name][1], ""


def mysql_command(args, extra=None):
    if not os.environ.get("MYSQL_PWD"):
        raise RuntimeError("必须通过MYSQL_PWD环境变量提供数据库密码")
    command = [
        args.mysql_bin, "--protocol=TCP", "-h", args.host, "-P", str(args.port),
        "-u", args.user, "--default-character-set=utf8mb4",
    ]
    if extra:
        command.extend(extra)
    return command


def mysql_run(args, sql, local_infile=False, capture=True):
    extra = ["--local-infile=1"] if local_infile else []
    process = subprocess.run(
        mysql_command(args, extra), input=sql, text=True, encoding="utf-8",
        capture_output=capture,
    )
    if process.returncode:
        raise RuntimeError(process.stderr.strip() or process.stdout.strip())
    return process.stdout


def mysql_run_file(args, path):
    mysql_run(args, path.read_text(encoding="utf-8"))


def build_load_sql(csv_files, summaries, database):
    statements = [
        "SET NAMES utf8mb4;", f"USE `{database}`;",
        "SET SESSION sql_log_bin=0;",
        "SET SESSION sql_mode='STRICT_TRANS_TABLES,ERROR_FOR_DIVISION_BY_ZERO,NO_ENGINE_SUBSTITUTION';",
        "TRUNCATE TABLE stg_gc_enrollment_plans;",
        "TRUNCATE TABLE stg_gc_school_admissions;",
        "TRUNCATE TABLE stg_gc_major_admissions;",
        "TRUNCATE TABLE stg_gc_score_ranges;",
    ]
    for path in csv_files:
        slug = path.parent.name
        province = PROVINCES[slug]
        filename = path.name
        source_url = f"https://huggingface.co/datasets/{DATASET}/resolve/{REVISION}/data/{YEAR}/{slug}/{filename}"
        digest = sha256(path)
        title_map = {
            "enrollment-plan.csv": "招生计划",
            "school-admission.csv": "院校投档线",
            "major-admission.csv": "专业录取线",
            "score-range.csv": "一分一段",
        }
        notes = (
            f"GaokaoCompass开放数据集{YEAR}年{province}{title_map[filename]}CSV；"
            "来源方说明其由各省教育考试院、阳光高考等公开渠道清洗整理。"
            "本系统将该文件作为待审核快照，不替代省级考试院原始公告。"
        )
        statements.append(f"""
        INSERT INTO data_sources (
          source_type, publisher, title, landing_page_url, file_url, published_on,
          retrieved_at, sha256, raw_file_path, notes
        ) VALUES (
          'open_dataset_csv', 'GaokaoCompass',
          {sql_string(f'{province}{YEAR}年{title_map[filename]}')},
          {sql_string(LANDING_URL)}, {sql_string(source_url)}, NULL, NOW(),
          {sql_string(digest)}, {sql_string(str(path.resolve()))}, {sql_string(notes)}
        ) AS incoming ON DUPLICATE KEY UPDATE
          retrieved_at=incoming.retrieved_at, raw_file_path=incoming.raw_file_path,
          file_url=incoming.file_url, notes=incoming.notes;
        SET @source_id=(SELECT source_id FROM data_sources WHERE sha256={sql_string(digest)});
        SET @row_no=0;
        """)
        table = STAGING[filename][0]
        columns, extra_set = load_mapping(path)
        local_path = path.resolve().as_posix()
        statements.append(f"""
        LOAD DATA LOCAL INFILE {sql_string(local_path)} INTO TABLE {table}
        CHARACTER SET utf8mb4
        FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '"' ESCAPED BY '"'
        LINES TERMINATED BY '\n'
        IGNORE 1 LINES
        ({columns})
        SET source_id=@source_id,
            source_row_number=(@row_no:=@row_no+1),
            province_name={sql_string(province)},
            admission_year=CAST(TRIM(@raw_year) AS UNSIGNED){extra_set};
        """)
    statements.append("SELECT 'staging_loaded' AS status;")
    return "\n".join(statements)


def main():
    parser = argparse.ArgumentParser(description="下载并覆盖导入GaokaoCompass 2025全国招录数据")
    parser.add_argument("--mysql-bin", required=True)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=3306)
    parser.add_argument("--user", default="root")
    parser.add_argument("--database", default="zhiyuantong")
    parser.add_argument("--root", default=str(Path(__file__).resolve().parents[1]))
    parser.add_argument("--refresh", action="store_true", help="强制重新下载文件")
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--province", choices=sorted(PROVINCES), help="只覆盖指定省份目录（英文目录名）")
    args = parser.parse_args()

    root = Path(args.root).resolve()
    raw_root = root / "data" / "raw" / "gaokao_compass" / str(YEAR)
    raw_root.mkdir(parents=True, exist_ok=True)
    tree = json.loads(request_bytes(TREE_URL, timeout=60).decode("utf-8"))
    items = [
        item for item in tree
        if item.get("type") == "file"
        and Path(item["path"]).name in set(EXPECTED_HEADERS) | {"meta.json"}
    ]
    expected_provinces = {Path(item["path"]).parts[-2] for item in items if Path(item["path"]).name == "enrollment-plan.csv"}
    if expected_provinces != set(PROVINCES):
        raise RuntimeError(f"招生计划省份不完整: {sorted(set(PROVINCES) - expected_provinces)}")

    print(f"开始下载/校验{len(items)}个文件，共{sum(int(item['size']) for item in items):,}字节")
    downloaded = []
    with ThreadPoolExecutor(max_workers=max(1, min(args.workers, 12))) as executor:
        futures = {executor.submit(download_one, item, raw_root, args.refresh): item for item in items}
        for index, future in enumerate(as_completed(futures), 1):
            downloaded.append(future.result())
            if index % 15 == 0 or index == len(items):
                print(f"已完成 {index}/{len(items)} 个文件")

    csv_files = sorted(path for path in downloaded if path.suffix.lower() == ".csv")
    summaries = {str(path): csv_summary(path) for path in csv_files}
    manifest = {
        "dataset": DATASET, "revision": REVISION, "year": YEAR,
        "landing_url": LANDING_URL, "downloaded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "files": [
            {
                "path": str(path.relative_to(root)).replace("\\", "/"),
                "size": path.stat().st_size, "sha256": sha256(path),
                "rows": summaries.get(str(path)),
            }
            for path in sorted(downloaded)
        ],
    }
    (raw_root / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"CSV结构校验通过：{len(csv_files)}个文件，共{sum(summaries.values()):,}行")

    migration = root / "database" / "migrations" / "007_add_gaokao_compass_2025_tables.sql"
    replacement = root / "database" / "import_gaokao_compass_2025.sql"
    migration_state = mysql_run(args, f"""
        SELECT COUNT(*) FROM information_schema.columns
        WHERE table_schema={sql_string(args.database)}
          AND table_name='admission_groups' AND column_name='batch_name';
    """).strip().splitlines()[-1]
    if migration_state == "0":
        mysql_run_file(args, migration)
    else:
        print("数据库007迁移已存在，跳过重复执行")
    mysql_run(args, "SET GLOBAL local_infile=ON;")
    try:
        province_slugs = [args.province] if args.province else sorted(PROVINCES)
        for index, slug in enumerate(province_slugs, 1):
            province_files = [path for path in csv_files if path.parent.name == slug]
            if not any(path.name == "enrollment-plan.csv" for path in province_files):
                raise RuntimeError(f"{PROVINCES[slug]}缺少招生计划文件，不能执行省级覆盖")
            load_output = mysql_run(
                args, build_load_sql(province_files, summaries, args.database), local_infile=True
            )
            if "staging_loaded" not in load_output:
                raise RuntimeError(f"{PROVINCES[slug]}暂存表装载未返回完成标志")
            stage_counts = mysql_run(args, f"""
                USE `{args.database}`;
                SELECT 'enrollment-plan.csv', COUNT(*) FROM stg_gc_enrollment_plans
                UNION ALL SELECT 'school-admission.csv', COUNT(*) FROM stg_gc_school_admissions
                UNION ALL SELECT 'major-admission.csv', COUNT(*) FROM stg_gc_major_admissions
                UNION ALL SELECT 'score-range.csv', COUNT(*) FROM stg_gc_score_ranges;
            """)
            actual = {}
            for line in stage_counts.splitlines():
                parts = line.split("\t")
                if len(parts) == 2 and parts[1].isdigit():
                    actual[parts[0]] = int(parts[1])
            expected_by_type = {
                filename: sum(
                    summaries[str(path)] for path in province_files if path.name == filename
                )
                for filename in EXPECTED_HEADERS
            }
            if actual != expected_by_type:
                raise RuntimeError(
                    f"{PROVINCES[slug]}暂存表行数校验失败，期望{expected_by_type}，实际{actual}"
                )
            distinct_provinces = mysql_run(args, f"""
                USE `{args.database}`;
                SELECT COUNT(DISTINCT province_name), MIN(province_name)
                FROM stg_gc_enrollment_plans;
            """).strip().splitlines()[-1].split("\t")
            if distinct_provinces != ["1", PROVINCES[slug]]:
                raise RuntimeError(f"暂存表省份边界校验失败：{distinct_provinces}")
            mysql_run_file(args, replacement)
            print(
                f"[{index}/{len(province_slugs)}] {PROVINCES[slug]}覆盖完成："
                + ", ".join(f"{name}={count:,}" for name, count in actual.items()),
                flush=True,
            )
    finally:
        mysql_run(args, "SET GLOBAL local_infile=OFF;")

    verification = mysql_run(args, f"""
        USE `{args.database}`;
        SELECT '专业招生计划', COUNT(*) FROM admission_major_plans plans
          JOIN admission_groups groups_ ON groups_.admission_group_id=plans.admission_group_id
          WHERE groups_.admission_year=2025
        UNION ALL SELECT '院校专业组', COUNT(*) FROM admission_groups WHERE admission_year=2025
        UNION ALL SELECT '院校投档线', COUNT(*) FROM admission_cutoff_lines WHERE admission_year=2025
        UNION ALL SELECT '专业录取线', COUNT(*) FROM admission_major_lines WHERE admission_year=2025
        UNION ALL SELECT '一分一段', COUNT(*) FROM score_rank_points WHERE admission_year=2025
        UNION ALL SELECT '省级控制线', COUNT(*) FROM control_score_lines WHERE admission_year=2025;
        SELECT '覆盖省份', COUNT(DISTINCT province_name) FROM admission_major_plans plans
          JOIN admission_groups groups_ ON groups_.admission_group_id=plans.admission_group_id
          WHERE groups_.admission_year=2025;
    """)
    print("覆盖导入完成：")
    print(verification.strip())


if __name__ == "__main__":
    main()