""" 课时核对表 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()