SBrainCO/db/pipeline/import_bill_records.py

#!/usr/bin/env python3
"""
账单查询原始数据导入脚本
从Excel文件导入bill_records和bill_columns表。

用法:
  python3 import_bill_records.py --month 2026-04-01 --file <文件或目录> --db bill_query_test
  python3 import_bill_records.py --month 2026-04-01 --file <目录> --db bill_query_test --skip-columns
"""
import argparse
import hashlib
import os
import sys

import pandas as pd
import psycopg2
from psycopg2.extras import execute_values


def file_sha256(filepath):
    h = hashlib.sha256()
    with open(filepath, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            h.update(chunk)
    return h.hexdigest()


def to_text(val):
    if val is None:
        return None
    if isinstance(val, float) and pd.isna(val):
        return None
    s = str(val).strip()
    if s.lower() in ('nan', 'none', ''):
        return None
    return s


def import_bill_columns(conn, filepaths):
    """从第一个Excel文件提取列定义，导入bill_columns表"""
    cur = conn.cursor()
    cur.execute("SELECT count(*) FROM public.bill_columns")
    existing = cur.fetchone()[0]
    if existing > 0:
        print(f"bill_columns已有{existing}行，跳过列定义导入")
        return

    if not filepaths:
        print("无文件可用，跳过列定义导入")
        return

    fpath = filepaths[0]
    df = pd.read_excel(fpath, header=None, nrows=2, engine='openpyxl')
    main_headers = df.iloc[0].tolist()
    sub_headers = df.iloc[1].tolist()

    columns = []
    for i, (mh, sh) in enumerate(zip(main_headers, sub_headers)):
        col_name = f"c{i+1:03d}"
        group_h = to_text(mh) or ''
        sub_h = to_text(sh) or ''
        # If group and sub are the same, only keep group
        if group_h and sub_h and group_h == sub_h:
            sub_h = ''
        columns.append((col_name, i + 1, group_h, sub_h))

    execute_values(
        cur,
        "INSERT INTO public.bill_columns (column_name, excel_position, group_header, sub_header) VALUES %s",
        columns,
        page_size=500
    )
    conn.commit()
    print(f"bill_columns导入完成: {len(columns)} 列")
    cur.close()


def import_bill_records(conn, filepath, report_month):
    """导入单个Excel文件的账单数据到bill_records"""
    fname = os.path.basename(filepath)
    sha = file_sha256(filepath)

    cur = conn.cursor()

    # 幂等检查
    cur.execute(
        "SELECT 1 FROM public.bill_records WHERE source_file = %s LIMIT 1",
        (fname,)
    )
    if cur.fetchone():
        print(f"  {fname} 已导入过，跳过")
        cur.close()
        return 0

    # 读取Excel：row0=主表头, row1=子表头, 数据从row2开始
    df = pd.read_excel(filepath, header=None, engine='openpyxl')
    # 跳过前2行（表头）
    data_df = df.iloc[2:]
    # 过滤空行（第一列为空）
    data_df = data_df.dropna(subset=[data_df.columns[0]])
    data_rows = data_df.values.tolist()

    print(f"  数据行数: {len(data_rows)}, 列数: {data_df.shape[1]}")

    if not data_rows:
        cur.close()
        return 0

    # 获取DB列名 (c001 ~ c196)
    num_cols = data_df.shape[1]
    col_names = [f"c{i+1:03d}" for i in range(num_cols)]
    all_cols = ['source_file', 'source_row'] + col_names
    col_str = ','.join(all_cols)

    batch = []
    for idx, row in enumerate(data_rows):
        source_row = idx + 3  # 1-indexed, data starts at Excel row 3
        vals = [fname, source_row]
        for i in range(num_cols):
            v = row[i] if i < len(row) else None
            vals.append(to_text(v))
        batch.append(tuple(vals))
        if len(batch) >= 500:
            execute_values(
                cur,
                f"INSERT INTO public.bill_records ({col_str}) VALUES %s",
                batch,
                page_size=500
            )
            conn.commit()
            if (idx + 1) % 10000 == 0:
                print(f"  已导入 {idx + 1}/{len(data_rows)} 行")
            batch = []

    if batch:
        execute_values(
            cur,
            f"INSERT INTO public.bill_records ({col_str}) VALUES %s",
            batch,
            page_size=500
        )
        conn.commit()

    print(f"  已导入 {len(data_rows)}/{len(data_rows)} 行")
    cur.close()
    return len(data_rows)


def main():
    parser = argparse.ArgumentParser(description='账单查询原始数据导入')
    parser.add_argument('--month', required=True, help='报告月份 (YYYY-MM-01)')
    parser.add_argument('--file', required=True, help='Excel文件路径或目录')
    parser.add_argument('--db', default='bill_query', help='目标数据库 (默认bill_query)')
    parser.add_argument('--skip-columns', action='store_true', help='跳过bill_columns导入')

    args = parser.parse_args()

    conn = psycopg2.connect(host='localhost', port=5432, dbname=args.db, user='freedak')
    conn.autocommit = False

    # 收集文件
    files = []
    if os.path.isdir(args.file):
        for f in sorted(os.listdir(args.file)):
            if f.endswith('.xlsx') and not f.startswith('.'):
                files.append(os.path.join(args.file, f))
    else:
        files.append(args.file)

    print(f"\n=== 账单查询原始数据导入 ===")
    print(f"文件数: {len(files)}")
    print(f"数据库: {args.db}")

    try:
        # 导入列定义
        if not args.skip_columns:
            import_bill_columns(conn, files)

        # 导入数据
        total = 0
        for idx, fpath in enumerate(files):
            fname = os.path.basename(fpath)
            print(f"\n[{idx+1}/{len(files)}] {fname}")
            total += import_bill_records(conn, fpath, args.month)

        print(f"\n账单数据导入完成: 共 {total} 行")
        print("=== 导入完成 ===")

    except Exception as e:
        conn.rollback()
        print(f"错误: {e}", file=sys.stderr)
        raise
    finally:
        conn.close()


if __name__ == '__main__':
    main()