defect-analysis/import_to_database.py

"""把缺陷 CSV 导入 SQLite 数据库。"""

import argparse

import pandas as pd

from defect_analysis.data_quality import build_data_quality_report
from defect_analysis.database import create_import_batch, init_database, insert_defects
from defect_analysis.schemas import normalize_defect_schema


def import_csv_to_database(csv_path, db_path, source_name=None):
    df = pd.read_csv(csv_path, parse_dates=["timestamp"])
    df = normalize_defect_schema(df)
    quality_report = build_data_quality_report(df)

    init_database(db_path)
    import_id = create_import_batch(
        db_path,
        source_name=source_name or str(csv_path),
        row_count=len(df),
        quality_score=quality_report["score"],
    )
    inserted = insert_defects(db_path, df, import_id=import_id)
    return {
        "import_id": import_id,
        "rows": len(df),
        "inserted": inserted,
        "quality_score": quality_report["score"],
        "issues": quality_report["issues"],
    }


def main():
    parser = argparse.ArgumentParser(description="导入缺陷 CSV 到 SQLite 数据库")
    parser.add_argument("--csv", default="defect_data.csv", help="缺陷 CSV 文件路径")
    parser.add_argument("--db", default="defect_analysis.db", help="SQLite 数据库路径")
    args = parser.parse_args()

    result = import_csv_to_database(args.csv, args.db)
    print(
        f"导入批次 {result['import_id']} 完成: "
        f"读取 {result['rows']} 行，新增 {result['inserted']} 行，数据质量分 {result['quality_score']:.1f}"
    )
    if result["issues"] != ["数据质量良好"]:
        print("数据质量提示: " + "；".join(result["issues"]))


if __name__ == "__main__":
    main()