579 lines
21 KiB
Python
579 lines
21 KiB
Python
"""
|
||
课时核对表 Excel → 数据库 ETL（批量校验 + 批量写入）
|
||
|
||
用法:
|
||
python -m etl.import_keshibiao --file keshibiao/课时核对表（2026.1月份）.xlsx --update
|
||
python -m etl.import_keshibiao --dir keshibiao --update --month 1
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import logging
|
||
import sys
|
||
import time
|
||
from collections import defaultdict
|
||
from decimal import Decimal
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
ROOT = Path(__file__).resolve().parent.parent
|
||
if str(ROOT) not in sys.path:
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
from etl.keshi_parser import (
|
||
ParsedWorkbook,
|
||
ParsedRow,
|
||
parse_workbook,
|
||
batch_validate_rows,
|
||
is_real_workbook_file,
|
||
)
|
||
|
||
logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s')
|
||
logger = logging.getLogger('keshi_etl')
|
||
|
||
DEFAULT_DIR = ROOT / 'keshibiao'
|
||
BULK_CHUNK = 500
|
||
COURSE_DISPLAY = {
|
||
'scratch': 'Scratch',
|
||
'python': 'Python',
|
||
'c++': 'C++',
|
||
'wedo': 'WeDo',
|
||
'大班': '大班',
|
||
'中班': '中班',
|
||
'小班': '小班',
|
||
'lab': 'Lab',
|
||
'stem': 'STEM',
|
||
'free': '免费课',
|
||
'trial': '试学',
|
||
}
|
||
|
||
|
||
def _d(val) -> Optional[Decimal]:
|
||
if val is None:
|
||
return None
|
||
return Decimal(str(round(float(val), 4)))
|
||
|
||
|
||
def _graduated_from_notes(notes: str) -> bool:
|
||
if not notes:
|
||
return False
|
||
return any(k in notes for k in ('结课', '毕业', '退费结课', '不再上课'))
|
||
|
||
|
||
def _fill_sparse_snapshot(data: dict, row: ParsedRow) -> dict:
|
||
"""试学缴费等仅含金额的行：补全数值字段，避免 Doris strict 模式过滤 INSERT"""
|
||
if row.course_code != 'trial' and data.get('end_balance') is not None:
|
||
return data
|
||
if row.course_code == 'trial' or (
|
||
data.get('new_signup_amount') is not None
|
||
and data.get('prev_lessons') is None
|
||
and data.get('end_lessons') is None
|
||
):
|
||
amt = data.get('new_signup_amount') or Decimal(0)
|
||
for key in (
|
||
'prev_lessons', 'prev_gift_lessons', 'prev_total_lessons',
|
||
'end_lessons', 'end_gift_lessons', 'end_total_lessons',
|
||
'consumed_lessons', 'consumed_amount',
|
||
'refund_lessons', 'refund_gift_lessons', 'refund_amount', 'account_fee',
|
||
):
|
||
data.setdefault(key, Decimal(0))
|
||
data.setdefault('prev_balance', Decimal(0))
|
||
data.setdefault('end_balance', amt)
|
||
return data
|
||
|
||
|
||
def build_snapshot_fields(row: ParsedRow, year: int, month: int, source_file: str) -> dict:
|
||
f = row.fields
|
||
end_lessons = f.get('end_lessons')
|
||
end_total = f.get('end_total_lessons')
|
||
end_gift = None
|
||
if end_lessons is not None and end_total is not None:
|
||
end_gift = end_total - end_lessons
|
||
elif f.get('prev_gift_lessons') is not None:
|
||
end_gift = f.get('prev_gift_lessons')
|
||
|
||
return {
|
||
'year': year,
|
||
'month': month,
|
||
'seq_no': row.seq_no,
|
||
'prev_lessons': _d(f.get('prev_lessons')),
|
||
'prev_gift_lessons': _d(f.get('prev_gift_lessons')),
|
||
'prev_total_lessons': _d(f.get('prev_total_lessons')),
|
||
'prev_balance': _d(f.get('prev_balance')),
|
||
'new_signup_amount': _d(f.get('new_signup_amount')),
|
||
'new_signup_lessons': _d(f.get('new_signup_lessons')),
|
||
'new_signup_gift_lessons': _d(f.get('new_signup_gift_lessons')),
|
||
'renewal_amount': _d(f.get('renewal_amount')),
|
||
'renewal_lessons': _d(f.get('renewal_lessons')),
|
||
'renewal_gift_lessons': _d(f.get('renewal_gift_lessons')),
|
||
'consumed_lessons': _d(f.get('consumed_lessons')),
|
||
'consumed_amount': _d(f.get('consumed_amount')),
|
||
'refund_lessons': _d(f.get('refund_lessons')),
|
||
'account_fee': _d(f.get('account_fee')),
|
||
'refund_gift_lessons': _d(f.get('refund_gift_lessons')),
|
||
'refund_amount': _d(f.get('refund_amount')),
|
||
'end_lessons': _d(end_lessons),
|
||
'end_gift_lessons': _d(end_gift),
|
||
'end_total_lessons': _d(end_total),
|
||
'end_balance': _d(f.get('end_balance')),
|
||
'unit_price': _d(f.get('unit_price')),
|
||
'notes': row.notes or None,
|
||
'source_file': source_file,
|
||
}
|
||
|
||
|
||
class ImportCaches:
|
||
"""预加载学员/课程/账户/当月快照，减少逐行查询"""
|
||
|
||
def __init__(self, db, Student, Course, StudentAccount, MonthlySnapshot, year, month):
|
||
self.db = db
|
||
self.Student = Student
|
||
self.Course = Course
|
||
self.StudentAccount = StudentAccount
|
||
self.MonthlySnapshot = MonthlySnapshot
|
||
self.year = year
|
||
self.month = month
|
||
|
||
self.courses_by_code: dict[str, Course] = {}
|
||
for c in Course.query.all():
|
||
for key in (c.course_code, c.level):
|
||
if key:
|
||
self.courses_by_code[str(key).lower()] = c
|
||
|
||
self.students_by_name: dict[str, Student] = {
|
||
s.name: s for s in Student.query.all()
|
||
}
|
||
self.accounts: dict[tuple[int, int], StudentAccount] = {}
|
||
for a in StudentAccount.query.all():
|
||
self.accounts[(a.student_id, a.course_id)] = a
|
||
|
||
self.snapshots: dict[int, MonthlySnapshot] = {
|
||
s.account_id: s
|
||
for s in MonthlySnapshot.query.filter_by(year=year, month=month).all()
|
||
}
|
||
|
||
self._pending_students: list = []
|
||
self._pending_courses: list = []
|
||
self._pending_accounts: list = []
|
||
|
||
def _course_key(self, code: str) -> str:
|
||
return (code or 'unknown').lower()
|
||
|
||
def ensure_courses(self, codes: set[str]) -> None:
|
||
for code in codes:
|
||
key = self._course_key(code)
|
||
if key in self.courses_by_code:
|
||
continue
|
||
name = COURSE_DISPLAY.get(key, key.upper())
|
||
course = self.Course(
|
||
name=name, course_code=key, level=key,
|
||
type='编程', status=1,
|
||
)
|
||
self._pending_courses.append(course)
|
||
self.courses_by_code[key] = course
|
||
|
||
def ensure_students(self, rows: list[ParsedRow]) -> None:
|
||
seen = set(self.students_by_name)
|
||
for row in rows:
|
||
if row.student_name in seen:
|
||
continue
|
||
seen.add(row.student_name)
|
||
student = self.Student(
|
||
name=row.student_name,
|
||
display_name=row.display_name or row.student_name,
|
||
nickname=row.nickname or None,
|
||
status=1,
|
||
source='课时表导入',
|
||
)
|
||
self._pending_students.append(student)
|
||
self.students_by_name[row.student_name] = student
|
||
|
||
def flush_pending(self) -> None:
|
||
need_flush = False
|
||
if self._pending_courses:
|
||
self.db.session.add_all(self._pending_courses)
|
||
self._pending_courses.clear()
|
||
need_flush = True
|
||
if self._pending_students:
|
||
self.db.session.add_all(self._pending_students)
|
||
self._pending_students.clear()
|
||
need_flush = True
|
||
if need_flush:
|
||
self.db.session.flush()
|
||
if self._pending_accounts:
|
||
self.db.session.add_all(self._pending_accounts)
|
||
self._pending_accounts.clear()
|
||
self.db.session.flush()
|
||
|
||
def ensure_accounts(self, pairs: set[tuple[str, str]]) -> None:
|
||
for student_name, course_code in pairs:
|
||
student = self.students_by_name.get(student_name)
|
||
course = self.courses_by_code.get(self._course_key(course_code))
|
||
if not student or not course:
|
||
continue
|
||
key = (student.id, course.id)
|
||
if key in self.accounts:
|
||
continue
|
||
orig = float(course.price_per_hour or 0) if course else 0
|
||
acc = self.StudentAccount(
|
||
student_id=student.id,
|
||
course_id=course.id,
|
||
normal_hours=0,
|
||
gifted_hours=0,
|
||
account_status='active',
|
||
original_price_per_lesson=orig if orig > 0 else None,
|
||
)
|
||
self._pending_accounts.append(acc)
|
||
self.accounts[key] = acc
|
||
|
||
def get_student(self, row: ParsedRow):
|
||
s = self.students_by_name[row.student_name]
|
||
if row.display_name and not s.display_name:
|
||
s.display_name = row.display_name
|
||
if row.nickname and not s.nickname:
|
||
s.nickname = row.nickname
|
||
return s
|
||
|
||
def get_course(self, course_code: str):
|
||
return self.courses_by_code[self._course_key(course_code)]
|
||
|
||
def get_account(self, student, course, unit_price=None):
|
||
acc = self.accounts[(student.id, course.id)]
|
||
if unit_price is not None:
|
||
acc.unit_price = _d(unit_price) or acc.unit_price
|
||
return acc
|
||
|
||
def get_snapshot(self, account_id: int):
|
||
return self.snapshots.get(account_id)
|
||
|
||
|
||
def _chunked(items: list, size: int):
|
||
for i in range(0, len(items), size):
|
||
yield items[i:i + size]
|
||
|
||
|
||
def import_workbook(
|
||
app,
|
||
path: Path,
|
||
*,
|
||
update: bool = False,
|
||
default_lessons: float = 2.0,
|
||
verbose: bool = False,
|
||
repair_chains: bool = True,
|
||
) -> dict:
|
||
from models import db, Student, Course, StudentAccount, MonthlySnapshot, ClassRecord
|
||
import business_rules as rules
|
||
|
||
t0 = time.perf_counter()
|
||
stats = {
|
||
'file': path.name,
|
||
'rows': 0,
|
||
'parsed_rows': 0,
|
||
'snapshots_new': 0,
|
||
'snapshots_updated': 0,
|
||
'snapshots_skipped': 0,
|
||
'class_records': 0,
|
||
'warnings': [],
|
||
'validation_warning_count': 0,
|
||
'sheets_parsed': [],
|
||
'sheets_skipped': [],
|
||
'elapsed_sec': 0,
|
||
'account_ids_touched': [],
|
||
}
|
||
|
||
parsed: ParsedWorkbook = parse_workbook(path, default_lessons=default_lessons)
|
||
stats['year'] = parsed.year
|
||
stats['month'] = parsed.month
|
||
stats['sheets_parsed'] = parsed.sheets_parsed
|
||
stats['sheets_skipped'] = parsed.sheets_skipped
|
||
stats['parsed_rows'] = len(parsed.rows)
|
||
|
||
valid_rows: list[ParsedRow] = []
|
||
for row in parsed.rows:
|
||
if not row.student_name or not row.course_code:
|
||
stats['warnings'].append(f'跳过缺姓名/课程行: {row.student_name}')
|
||
continue
|
||
valid_rows.append(row)
|
||
|
||
val_warnings = batch_validate_rows(valid_rows)
|
||
stats['validation_warning_count'] = len(val_warnings)
|
||
if val_warnings:
|
||
stats['warnings'].extend(val_warnings[:30])
|
||
logger.info(
|
||
'批量校验完成: %s 条余额/课时警告（已合并，不逐条刷屏；--verbose 可看样例）',
|
||
len(val_warnings),
|
||
)
|
||
if verbose:
|
||
for w in val_warnings[:20]:
|
||
logger.warning(w)
|
||
|
||
with app.app_context():
|
||
from db_bootstrap import migrate_keshi_ledger_columns
|
||
migrate_keshi_ledger_columns()
|
||
|
||
caches = ImportCaches(
|
||
db, Student, Course, StudentAccount, MonthlySnapshot,
|
||
parsed.year, parsed.month,
|
||
)
|
||
|
||
course_codes = {r.course_code for r in valid_rows}
|
||
account_pairs = {(r.student_name, r.course_code) for r in valid_rows}
|
||
|
||
caches.ensure_courses(course_codes)
|
||
caches.ensure_students(valid_rows)
|
||
caches.flush_pending()
|
||
caches.ensure_accounts(account_pairs)
|
||
caches.flush_pending()
|
||
|
||
new_snapshots: list = []
|
||
class_mappings: list[dict] = []
|
||
account_ids_touched: set[int] = set()
|
||
|
||
for row in valid_rows:
|
||
stats['rows'] += 1
|
||
course = caches.get_course(row.course_code)
|
||
student = caches.get_student(row)
|
||
unit_price = row.fields.get('unit_price')
|
||
account = caches.get_account(student, course, unit_price)
|
||
account_ids_touched.add(account.id)
|
||
|
||
snap_data = _fill_sparse_snapshot(
|
||
build_snapshot_fields(row, parsed.year, parsed.month, path.name),
|
||
row,
|
||
)
|
||
snap = caches.get_snapshot(account.id)
|
||
|
||
if snap:
|
||
if not update:
|
||
stats['snapshots_skipped'] += 1
|
||
else:
|
||
for k, v in snap_data.items():
|
||
setattr(snap, k, v)
|
||
stats['snapshots_updated'] += 1
|
||
else:
|
||
new_snapshots.append(
|
||
MonthlySnapshot(account_id=account.id, **snap_data)
|
||
)
|
||
stats['snapshots_new'] += 1
|
||
caches.snapshots[account.id] = new_snapshots[-1]
|
||
|
||
if snap_data.get('end_lessons') is not None:
|
||
account.normal_hours = snap_data['end_lessons']
|
||
if snap_data.get('end_gift_lessons') is not None:
|
||
account.gifted_hours = snap_data['end_gift_lessons']
|
||
elif snap_data.get('end_total_lessons') is not None and snap_data.get('end_lessons') is not None:
|
||
account.gifted_hours = snap_data['end_total_lessons'] - snap_data['end_lessons']
|
||
if snap_data.get('unit_price') is not None:
|
||
account.unit_price = snap_data['unit_price']
|
||
if _graduated_from_notes(row.notes):
|
||
account.account_status = 'graduated'
|
||
student.status = 0
|
||
|
||
for class_date, lessons, raw in row.class_sessions:
|
||
class_mappings.append({
|
||
'account_id': account.id,
|
||
'year': int(parsed.year),
|
||
'month': int(parsed.month),
|
||
'class_date': class_date,
|
||
'lessons_consumed': float(lessons) if lessons is not None else 2.0,
|
||
'raw_text': (raw or '')[:50] or None,
|
||
})
|
||
|
||
if update and account_ids_touched:
|
||
ClassRecord.query.filter(
|
||
ClassRecord.account_id.in_(account_ids_touched),
|
||
ClassRecord.year == parsed.year,
|
||
ClassRecord.month == parsed.month,
|
||
).delete(synchronize_session=False)
|
||
|
||
if new_snapshots:
|
||
db.session.add_all(new_snapshots)
|
||
|
||
if class_mappings:
|
||
# Doris 对 executemany + Decimal/datetime 不友好，用 add_all 分批提交
|
||
for chunk in _chunked(class_mappings, BULK_CHUNK):
|
||
db.session.add_all([ClassRecord(**m) for m in chunk])
|
||
stats['class_records'] = len(class_mappings)
|
||
|
||
db.session.commit()
|
||
stats['account_ids_touched'] = list(account_ids_touched)
|
||
|
||
if repair_chains and account_ids_touched:
|
||
repaired = rules.repair_snapshot_chains_for_accounts(account_ids_touched)
|
||
stats['chain_repaired'] = repaired
|
||
logger.info('快照链修复: %s 个衔接点（仅本次涉及 %s 个账户）', repaired, len(account_ids_touched))
|
||
|
||
stats['elapsed_sec'] = round(time.perf_counter() - t0, 2)
|
||
return stats
|
||
|
||
|
||
def validate_month_chain(app, tolerance: float = 0.05) -> list[str]:
|
||
from models import MonthlySnapshot
|
||
|
||
warnings = []
|
||
with app.app_context():
|
||
snaps = MonthlySnapshot.query.order_by(
|
||
MonthlySnapshot.account_id,
|
||
MonthlySnapshot.year,
|
||
MonthlySnapshot.month,
|
||
).all()
|
||
by_account: dict[int, list] = defaultdict(list)
|
||
for s in snaps:
|
||
by_account[s.account_id].append(s)
|
||
|
||
for account_id, items in by_account.items():
|
||
items.sort(key=lambda x: (x.year, x.month))
|
||
for i in range(len(items) - 1):
|
||
cur, nxt = items[i], items[i + 1]
|
||
expected_next = cur.month + 1
|
||
expected_year = cur.year
|
||
if expected_next > 12:
|
||
expected_next = 1
|
||
expected_year += 1
|
||
if (nxt.year, nxt.month) != (expected_year, expected_next):
|
||
continue
|
||
|
||
pairs = [
|
||
('正课', cur.end_lessons, nxt.prev_lessons),
|
||
('赠课', cur.end_gift_lessons, nxt.prev_gift_lessons),
|
||
('余额', cur.end_balance, nxt.prev_balance),
|
||
]
|
||
for label, end_v, prev_v in pairs:
|
||
if end_v is None or prev_v is None:
|
||
continue
|
||
if abs(float(end_v) - float(prev_v)) > tolerance:
|
||
warnings.append(
|
||
f'账户{account_id} {cur.period_label}→{nxt.period_label} '
|
||
f'{label}不衔接: 月末{float(end_v)} vs 下月初{float(prev_v)}'
|
||
)
|
||
return warnings
|
||
|
||
|
||
def _filter_files(files: list[Path], month: Optional[int]) -> list[Path]:
|
||
if month is None:
|
||
return files
|
||
out = []
|
||
for f in files:
|
||
name = f.name
|
||
if (
|
||
f'_{month}月' in name
|
||
or f'.{month}月' in name
|
||
or f'.{month}月份' in name
|
||
or f'（{month}月' in name
|
||
or f'({month}月' in name
|
||
):
|
||
out.append(f)
|
||
return out
|
||
|
||
|
||
def import_directory(
|
||
app,
|
||
directory: Path,
|
||
*,
|
||
update: bool = False,
|
||
default_lessons: float = 2.0,
|
||
month: Optional[int] = None,
|
||
verbose: bool = False,
|
||
repair_chains: bool = True,
|
||
validate_chain: bool = False,
|
||
) -> list[dict]:
|
||
directory = Path(directory)
|
||
if not directory.is_dir():
|
||
raise FileNotFoundError(f'目录不存在: {directory}')
|
||
|
||
files = sorted(directory.glob('*.xlsx'))
|
||
files = [
|
||
f for f in files
|
||
if not f.name.startswith('~$') and is_real_workbook_file(f)
|
||
]
|
||
files = _filter_files(files, month)
|
||
if not files:
|
||
logger.warning('目录内无匹配 xlsx: %s month=%s', directory, month)
|
||
return []
|
||
|
||
results = []
|
||
all_touched: set[int] = set()
|
||
t0 = time.perf_counter()
|
||
|
||
for path in files:
|
||
logger.info('导入 %s', path.name)
|
||
try:
|
||
stats = import_workbook(
|
||
app, path, update=update, default_lessons=default_lessons,
|
||
verbose=verbose, repair_chains=False,
|
||
)
|
||
all_touched.update(stats.get('account_ids_touched') or [])
|
||
results.append(stats)
|
||
logger.info(
|
||
' 完成 [%ss]: 解析=%s 入库=%s 表=%s 新增=%s 更新=%s 上课=%s 校验警告=%s',
|
||
stats['elapsed_sec'],
|
||
stats.get('parsed_rows'), stats['rows'],
|
||
stats.get('sheets_parsed'),
|
||
stats['snapshots_new'],
|
||
stats['snapshots_updated'],
|
||
stats['class_records'],
|
||
stats.get('validation_warning_count', 0),
|
||
)
|
||
except Exception as e:
|
||
logger.exception('导入失败 %s: %s', path.name, e)
|
||
results.append({'file': path.name, 'error': str(e)})
|
||
|
||
if repair_chains and all_touched:
|
||
import business_rules as rules
|
||
with app.app_context():
|
||
n = rules.repair_snapshot_chains_for_accounts(all_touched)
|
||
logger.info('目录导入后快照链修复: %s 点 / %s 账户', n, len(all_touched))
|
||
|
||
if validate_chain:
|
||
chain_warnings = validate_month_chain(app)
|
||
logger.info('全库链式校验: %s 条', len(chain_warnings))
|
||
if results:
|
||
results[-1]['chain_warnings'] = chain_warnings
|
||
|
||
logger.info('目录导入总耗时 %.1fs，文件数 %s', time.perf_counter() - t0, len(files))
|
||
return results
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='课时核对表 Excel 导入（批量模式）')
|
||
parser.add_argument('--dir', type=str, default=str(DEFAULT_DIR), help='xlsx 目录')
|
||
parser.add_argument('--file', type=str, help='单个 xlsx 文件')
|
||
parser.add_argument('--month', type=int, help='仅导入文件名含该月份的表（与 --dir 联用）')
|
||
parser.add_argument('--update', action='store_true', help='已存在快照则更新')
|
||
parser.add_argument('--lessons-per-class', type=float, default=2.0, help='默认每次上课课时')
|
||
parser.add_argument('--verbose', action='store_true', help='打印校验警告样例')
|
||
parser.add_argument('--no-repair', action='store_true', help='跳过快照链修复（最快）')
|
||
parser.add_argument('--validate-chain', action='store_true', help='导入后全库链式校验（慢）')
|
||
args = parser.parse_args()
|
||
|
||
from app import create_app
|
||
app = create_app()
|
||
repair = not args.no_repair
|
||
|
||
if args.file:
|
||
stats = import_workbook(
|
||
app, Path(args.file),
|
||
update=args.update,
|
||
default_lessons=args.lessons_per_class,
|
||
verbose=args.verbose,
|
||
repair_chains=repair,
|
||
)
|
||
print(stats)
|
||
else:
|
||
results = import_directory(
|
||
app, Path(args.dir),
|
||
update=args.update,
|
||
default_lessons=args.lessons_per_class,
|
||
month=args.month,
|
||
verbose=args.verbose,
|
||
repair_chains=repair,
|
||
validate_chain=args.validate_chain,
|
||
)
|
||
print(f'共处理 {len(results)} 个文件')
|
||
for r in results:
|
||
print(r)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|