all-project/学校学生操行历史数据迁移/04_import_to_production.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
学生操行分历史数据迁移 - 脚本4：正式导入处理（简化版本）

功能说明：
- 使用高效的INSERT INTO SELECT语句直接导入数据
- 自动生成雪花ID作为主键
- 支持重新处理模式：重新导入所有已关联的数据
- 简洁高效，避免复杂的重复检查和批量处理逻辑

执行方式：
python 04_import_to_production.py              # 正常模式：仅处理待导入数据
python 04_import_to_production.py --reprocess  # 重新处理模式：重新处理所有已关联数据

注意事项：
- 依赖脚本3已执行完成
- 使用数据库的INSERT IGNORE处理重复数据
- 性能大幅提升，代码大幅简化
"""

from __future__ import annotations

import argparse
import logging
import sys
import time
from datetime import datetime
from typing import Any, Dict, List

import pymysql

from db_config import DB_CONFIG, TABLE_CONFIG, CONSTANTS

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler("logs/04_import_production.log", mode='a', encoding='utf-8')
    ]
)
logger = logging.getLogger("import_production")

class ProductionImporter:
    """正式表导入处理器（简化版本）"""

    def __init__(self):
        self.db_config = DB_CONFIG
        self.temp_parsed_table = TABLE_CONFIG["temp_parsed"]
        self.production_table = TABLE_CONFIG["production"]
        self.tenant_id = CONSTANTS["DEFAULT_TENANT_ID"]
        self.create_by = CONSTANTS["DEFAULT_CREATE_BY"]
        self.update_by = CONSTANTS["DEFAULT_UPDATE_BY"]
        self.connection = None
        self.sequence_counter = 0  # 雪花ID序列号计数器
        self.last_timestamp = -1    # 上次时间戳

    def connect_db(self) -> None:
        """连接数据库"""
        try:
            self.connection = pymysql.connect(**self.db_config)
            logger.info("数据库连接成功")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise

    def disconnect_db(self) -> None:
        """断开数据库连接"""
        if self.connection:
            self.connection.close()
            logger.info("数据库连接已断开")

    def generate_snowflake_id(self) -> int:
        """生成雪花ID"""
        epoch = CONSTANTS["SNOWFLAKE_EPOCH"]
        worker_id = CONSTANTS["WORKER_ID"]
        datacenter_id = CONSTANTS["DATACENTER_ID"]

        timestamp = int(time.time() * 1000)

        if timestamp == self.last_timestamp:
            self.sequence_counter = (self.sequence_counter + 1) % 4096
        else:
            self.sequence_counter = 0
            self.last_timestamp = timestamp

        timestamp_part = (timestamp - epoch) << 22
        datacenter_part = datacenter_id << 17
        worker_part = worker_id << 12
        sequence_part = self.sequence_counter

        return timestamp_part | datacenter_part | worker_part | sequence_part

    def get_pending_ids(self, reprocess: bool = False) -> List[int]:
        """获取待处理的记录ID"""
        try:
            with self.connection.cursor() as cursor:
                if reprocess:
                    # 重新处理模式：获取所有已关联的数据
                    sql = f"""
                    SELECT id FROM {self.temp_parsed_table}
                    WHERE processing_status IN ('linked', 'imported', 'duplicate')
                        AND student_id IS NOT NULL
                    """
                else:
                    # 正常模式：只获取待导入的数据
                    sql = f"""
                    SELECT id FROM {self.temp_parsed_table}
                    WHERE processing_status = 'linked'
                        AND student_id IS NOT NULL
                    """

                cursor.execute(sql)
                return [row[0] for row in cursor.fetchall()]

        except Exception as e:
            logger.error(f"获取待处理记录ID失败: {e}")
            return []

    def import_data_directly(self, record_ids: List[int]) -> Dict[str, int]:
        """使用INSERT INTO SELECT直接导入数据"""
        if not record_ids:
            return {'imported': 0, 'duplicates': 0, 'errors': 0}

        try:
            with self.connection.cursor() as cursor:
                # 先预生成一批雪花ID
                snowflake_ids = []
                for _ in range(len(record_ids)):
                    snowflake_ids.append(self.generate_snowflake_id())

                current_time = datetime.now()

                # 创建临时表存储ID映射
                temp_id_table = "temp_snowflake_ids"
                cursor.execute(f"DROP TEMPORARY TABLE IF EXISTS {temp_id_table}")

                cursor.execute(f"""
                CREATE TEMPORARY TABLE {temp_id_table} (
                    record_id BIGINT PRIMARY KEY,
                    snowflake_id BIGINT UNIQUE
                )
                """)

                # 批量插入ID映射
                id_mapping = list(zip(record_ids, snowflake_ids))
                cursor.executemany(f"""
                INSERT INTO {temp_id_table} (record_id, snowflake_id) VALUES (%s, %s)
                """, id_mapping)

                # 使用INSERT IGNORE直接导入数据
                insert_sql = f"""
                INSERT IGNORE INTO {self.production_table} (
                    id, tenant_id, create_dept, create_by, create_time,
                    update_by, update_time, deleted,
                    academic_year, semester, month,
                    student_id, class_id, class_teacher_id,
                    score, remark
                )
                SELECT
                    t.snowflake_id,
                    '{self.tenant_id}',
                    NULL,
                    {self.create_by},
                    '{current_time.strftime('%Y-%m-%d %H:%M:%S')}',
                    {self.update_by},
                    '{current_time.strftime('%Y-%m-%d %H:%M:%S')}',
                    0,
                    p.academic_year, p.semester, p.month,
                    p.student_id, p.class_id, p.class_teacher_id,
                    p.score, COALESCE(p.remark, '')
                FROM {self.temp_parsed_table} p
                INNER JOIN {temp_id_table} t ON p.id = t.record_id
                WHERE p.class_teacher_id IS NOT NULL
                """

                cursor.execute(insert_sql)
                imported_count = cursor.rowcount

                # 计算重复和错误数量
                total_processed = len(record_ids)
                duplicates = total_processed - imported_count
                errors = 0  # INSERT IGNORE不会产生错误，只有重复

                self.connection.commit()

                # 清理临时表
                cursor.execute(f"DROP TEMPORARY TABLE {temp_id_table}")

                return {
                    'imported': imported_count,
                    'duplicates': duplicates,
                    'errors': errors
                }

        except Exception as e:
            logger.error(f"直接导入数据失败: {e}")
            self.connection.rollback()
            return {'imported': 0, 'duplicates': 0, 'errors': len(record_ids)}

    def update_status_after_import(self, record_ids: List[int], imported_count: int, duplicate_count: int) -> None:
        """根据导入结果更新状态"""
        try:
            with self.connection.cursor() as cursor:
                # 更新已导入的记录状态
                if imported_count > 0:
                    imported_ids = record_ids[:imported_count]
                    placeholders = ','.join(['%s'] * len(imported_ids))
                    cursor.execute(f"""
                    UPDATE {self.temp_parsed_table}
                    SET processing_status = 'imported', update_time = NOW()
                    WHERE id IN ({placeholders})
                    """, imported_ids)

                # 更新重复的记录状态
                if duplicate_count > 0:
                    duplicate_ids = record_ids[imported_count:imported_count + duplicate_count]
                    placeholders = ','.join(['%s'] * len(duplicate_ids))
                    cursor.execute(f"""
                    UPDATE {self.temp_parsed_table}
                    SET processing_status = 'duplicate', error_msg = '数据已存在', update_time = NOW()
                    WHERE id IN ({placeholders})
                    """, duplicate_ids)

                self.connection.commit()

        except Exception as e:
            logger.error(f"更新状态失败: {e}")
            self.connection.rollback()

    def get_statistics(self) -> Dict[str, Any]:
        """获取处理统计"""
        try:
            with self.connection.cursor() as cursor:
                # 中间表统计
                sql_parsed = f"""
                SELECT
                    COUNT(*) as total,
                    SUM(CASE WHEN processing_status = 'linked' THEN 1 ELSE 0 END) as pending,
                    SUM(CASE WHEN processing_status = 'imported' THEN 1 ELSE 0 END) as imported,
                    SUM(CASE WHEN processing_status = 'duplicate' THEN 1 ELSE 0 END) as duplicate,
                    SUM(CASE WHEN processing_status = 'import_error' THEN 1 ELSE 0 END) as error
                FROM {self.temp_parsed_table}
                """

                cursor.execute(sql_parsed)
                parsed_stats = cursor.fetchone()

                # 正式表统计
                sql_production = f"""
                SELECT COUNT(*) as total_production
                FROM {self.production_table}
                WHERE deleted = 0
                """

                cursor.execute(sql_production)
                production_stats = cursor.fetchone()

                return {
                    'parsed_total': parsed_stats[0] or 0,
                    'parsed_pending': parsed_stats[1] or 0,
                    'parsed_imported': parsed_stats[2] or 0,
                    'parsed_duplicate': parsed_stats[3] or 0,
                    'parsed_error': parsed_stats[4] or 0,
                    'production_total': production_stats[0] or 0
                }

        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {}

    def run(self, reprocess: bool = False) -> bool:
        """执行导入流程（简化版本）"""
        start_time = time.time()

        try:
            self.connect_db()

            # 显示初始统计
            initial_stats = self.get_statistics()
            if reprocess:
                logger.info("重新处理模式：将重新处理所有已关联的数据")
            else:
                logger.info("正常处理模式：仅处理待导入的数据")
            logger.info("正式表导入处理开始")
            logger.info(f"当前状态: 中间表总计{initial_stats.get('parsed_total', 0)}条记录，")
            logger.info(f"  待导入: {initial_stats.get('parsed_pending', 0)}")
            logger.info(f"  已导入: {initial_stats.get('parsed_imported', 0)}")
            logger.info(f"  重复数据: {initial_stats.get('parsed_duplicate', 0)}")
            logger.info(f"  导入错误: {initial_stats.get('parsed_error', 0)}")
            logger.info(f"正式表当前总记录数: {initial_stats.get('production_total', 0)}")

            # 获取所有待处理的记录ID
            record_ids = self.get_pending_ids(reprocess)
            if not record_ids:
                logger.info("没有待处理的记录")
                return True

            logger.info(f"开始直接导入 {len(record_ids)} 条记录...")

            # 直接导入所有数据
            result = self.import_data_directly(record_ids)

            # 更新状态
            self.update_status_after_import(record_ids, result['imported'], result['duplicates'])

            # 显示最终统计
            final_stats = self.get_statistics()
            logger.info("="*50)
            logger.info("正式表导入统计结果:")
            logger.info(f"  处理总记录: {len(record_ids)}")
            logger.info(f"  成功导入: {result['imported']}")
            logger.info(f"  重复跳过: {result['duplicates']}")
            logger.info(f"  导入失败: {result['errors']}")
            logger.info(f"  最终状态:")
            logger.info(f"    - 中间表已导入: {final_stats.get('parsed_imported', 0)}")
            logger.info(f"    - 中间表重复: {final_stats.get('parsed_duplicate', 0)}")
            logger.info(f"    - 中间表错误: {final_stats.get('parsed_error', 0)}")
            logger.info(f"    - 正式表总记录: {final_stats.get('production_total', 0)}")
            logger.info(f"导入总耗时: {elapsed_time:.1f}秒")
            logger.info("="*50)

            success = result['errors'] == 0
            if not success:
                logger.warning(f"有{result['errors']}条记录导入失败，请检查日志")

            return success

        except Exception as e:
            logger.error(f"导入过程出现错误: {e}")
            return False

        finally:
            self.disconnect_db()
            elapsed_time = time.time() - start_time
            logger.info(f"导入总耗时: {elapsed_time:.1f}秒")
def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='学生操行分历史数据 - 正式导入处理')
    parser.add_argument('--reprocess', action='store_true',
                       help='重新处理已导入的数据（忽略processing_status状态）')

    args = parser.parse_args()

    logger.info("="*60)
    logger.info("学生操行分历史数据迁移 - 脚本4：正式导入处理")
    if args.reprocess:
        logger.info("模式: 重新处理所有已关联的数据")
    else:
        logger.info("模式: 仅处理待导入的数据")
    logger.info("="*60)

    # 创建logs目录
    import os
    os.makedirs("logs", exist_ok=True)

    # 执行导入
    importer = ProductionImporter()
    success = importer.run(reprocess=args.reprocess)

    if success:
        logger.info("✅ 正式表导入处理完成")
        sys.exit(0)
    else:
        logger.error("❌ 正式表导入处理失败")
        sys.exit(1)

if __name__ == "__main__":
    main()