feat: 新增db/pipeline数据导入管线，适配5月数据格式
- 整合所有导入脚本到db/pipeline/目录 - config.py: 按月份动态配置数据源路径，支持4月/5月不同目录结构 - import_bill_records_may.py: 5月专用账单导入(51列全渠道订单明细→bill_records + 50列品项销售明细→dish_sales_details) - import_salary_attendance.py: 添加--salary-file/--attendance-file参数支持动态文件路径 - import_may_data.sh: 5月一键导入脚本(12步全流程) - run_all.py: 一键全量导入+物化视图刷新编排 - refresh_materialized_views.py: 按依赖顺序刷新38个物化视图 - verify.py/verify_import.py: 数据一致性验证 - README.md: 管线文档
This commit is contained in:
freedakgmail
@@ -0,0 +1,198 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user