原始数据建表优化

2025-11-18 09:54:15 +08:00
parent d8b6247094
commit 0e782baed9
1 changed files with 99 additions and 94 deletions
--- a/app/services/original_data.py
+++ b/app/services/original_data.py
@@ -337,101 +337,46 @@ class OriginalDataService(BaseService[OriginalData]):

        table_name = self._get_table_name(account_id)

-        for attempt in range(2):  # 最多重试1次
+        # 检查是否已在事务中（避免重复开始事务）
+        in_transaction = db.in_transaction()
+
+        # 如果不在事务中，才需要手动管理事务
+        if not in_transaction:
+            for attempt in range(2):  # 最多重试1次
+                try:
+                    db.begin()
+                    success_count = 0
+                    failed_count = 0
+                    failed_items = []
+
+                    # 执行数据导入操作
+                    success_count = self._execute_import(db, table_name, data, account_id)
+                    db.commit()
+                    logger.info(f"Batch import original data completed. Success: {success_count}, Failed: {failed_count}")
+                    break
+
+                except Exception as e:
+                    db.rollback()
+                    logger.warning(f"Batch import attempt {attempt + 1} failed: {str(e)}")
+                    if attempt == 1:  # 最后一次重试失败
+                        logger.error("Batch import original data failed after retries")
+                        return {
+                            'success': False,
+                            'message': f'批量导入失败: {str(e)}',
+                            'total_count': total_count,
+                            'success_count': 0,
+                            'failed_count': total_count,
+                            'failed_items': failed_items
+                        }
+        else:
+            # 如果已在事务中，直接执行操作（不管理事务）
            try:
-                db.begin()
-                success_count = 0
-                failed_count = 0
-                failed_items = []
-
-                nyid = str(data[0].get('NYID'))  # 统一转换为字符串
-                # 检查该期数数据是否已存在
-                check_query = text(f"SELECT COUNT(*) as cnt FROM `{table_name}` WHERE NYID = :nyid")
-                is_exists = db.execute(check_query, {"nyid": nyid}).fetchone()[0]
-
-                if is_exists > 0:
-                    db.rollback()
-                    return {
-                        'success': True,
-                        'message': '数据已存在',
-                        'total_count': 0,
-                        'success_count': success_count,
-                        'failed_count': failed_count,
-                        'failed_items': failed_items
-                    }
-
-                # ===== 性能优化：批量查询沉降数据 =====
-                # 统一转换为字符串处理（数据库NYID字段是VARCHAR类型）
-                nyid_list = list(set(str(item.get('NYID')) for item in data if item.get('NYID')))
-                settlements = db.query(SettlementData).filter(SettlementData.NYID.in_(nyid_list)).all()
-                settlement_map = {s.NYID: s for s in settlements}
-                missing_nyids = set(nyid_list) - set(settlement_map.keys())
-
-                if missing_nyids:
-                    logger.warning(f"[批量导入原始数据] 批量查询settlement数据失败， Nyid: {list(missing_nyids)}")
-                    db.rollback()
-                    return {
-                        'success': False,
-                        'message': f'以下期数在沉降表中不存在: {list(missing_nyids)}',
-                        'total_count': total_count,
-                        'success_count': 0,
-                        'failed_count': total_count,
-                        'failed_items': []
-                    }
-
-                # ===== 性能优化：使用批量插入 =====
-                # 将数据分组，每组1000条（MySQL默认支持）
-                batch_size = 1000
-                for i in range(0, len(data), batch_size):
-                    batch_data = data[i:i + batch_size]
-
-                    # 构建批量参数
-                    values_list = []
-                    params = {}
-                    for idx, item_data in enumerate(batch_data):
-                        values_list.append(
-                            f"(:account_id_{idx}, :bfpcode_{idx}, :mtime_{idx}, :bffb_{idx}, "
-                            f":bfpl_{idx}, :bfpvalue_{idx}, :NYID_{idx}, :sort_{idx})"
-                        )
-                        params.update({
-                            f"account_id_{idx}": account_id,
-                            f"bfpcode_{idx}": item_data.get('bfpcode'),
-                            f"mtime_{idx}": item_data.get('mtime'),
-                            f"bffb_{idx}": item_data.get('bffb'),
-                            f"bfpl_{idx}": item_data.get('bfpl'),
-                            f"bfpvalue_{idx}": item_data.get('bfpvalue'),
-                            f"NYID_{idx}": item_data.get('NYID'),
-                            f"sort_{idx}": item_data.get('sort')
-                        })
-
-                    # 批量插入SQL - 使用字符串拼接（修复TextClause拼接问题）
-                    insert_sql = f"""
-                        INSERT INTO `{table_name}`
-                        (account_id, bfpcode, mtime, bffb, bfpl, bfpvalue, NYID, sort)
-                        VALUES {", ".join(values_list)}
-                    """
-                    final_sql = text(insert_sql)
-                    db.execute(final_sql, params)
-                    success_count += len(batch_data)
-                    logger.info(f"Inserted batch {i//batch_size + 1}: {len(batch_data)} records")
-
-                db.commit()
-                logger.info(f"Batch import original data completed. Success: {success_count}, Failed: {failed_count}")
-                break
-
+                success_count = self._execute_import(db, table_name, data, account_id)
+                logger.info(f"Batch import original data completed in existing transaction. Success: {success_count}")
            except Exception as e:
-                db.rollback()
-                logger.warning(f"Batch import attempt {attempt + 1} failed: {str(e)}")
-                if attempt == 1:  # 最后一次重试失败
-                    logger.error("Batch import original data failed after retries")
-                    return {
-                        'success': False,
-                        'message': f'批量导入失败: {str(e)}',
-                        'total_count': total_count,
-                        'success_count': 0,
-                        'failed_count': total_count,
-                        'failed_items': failed_items
-                    }
+                logger.error(f"Batch import failed in existing transaction: {str(e)}")
+                # 抛出异常，让外部处理事务回滚
+                raise

        return {
            'success': True,
@@ -440,4 +385,64 @@ class OriginalDataService(BaseService[OriginalData]):
            'success_count': success_count,
            'failed_count': failed_count,
            'failed_items': failed_items
-        }
+        }
+
+    def _execute_import(self, db: Session, table_name: str, data: List, account_id: int) -> int:
+        """执行数据导入操作（抽取的公共逻辑）"""
+        nyid = str(data[0].get('NYID'))  # 统一转换为字符串
+        # 检查该期数数据是否已存在
+        check_query = text(f"SELECT COUNT(*) as cnt FROM `{table_name}` WHERE NYID = :nyid")
+        is_exists = db.execute(check_query, {"nyid": nyid}).fetchone()[0]
+
+        if is_exists > 0:
+            return 0
+
+        # ===== 性能优化：批量查询沉降数据 =====
+        # 统一转换为字符串处理（数据库NYID字段是VARCHAR类型）
+        nyid_list = list(set(str(item.get('NYID')) for item in data if item.get('NYID')))
+        from ..models.settlement_data import SettlementData
+        settlements = db.query(SettlementData).filter(SettlementData.NYID.in_(nyid_list)).all()
+        settlement_map = {s.NYID: s for s in settlements}
+        missing_nyids = set(nyid_list) - set(settlement_map.keys())
+
+        if missing_nyids:
+            raise Exception(f'以下期数在沉降表中不存在: {list(missing_nyids)}')
+
+        # ===== 性能优化：使用批量插入 =====
+        # 将数据分组，每组1000条（MySQL默认支持）
+        batch_size = 1000
+        total_inserted = 0
+        for i in range(0, len(data), batch_size):
+            batch_data = data[i:i + batch_size]
+
+            # 构建批量参数
+            values_list = []
+            params = {}
+            for idx, item_data in enumerate(batch_data):
+                values_list.append(
+                    f"(:account_id_{idx}, :bfpcode_{idx}, :mtime_{idx}, :bffb_{idx}, "
+                    f":bfpl_{idx}, :bfpvalue_{idx}, :NYID_{idx}, :sort_{idx})"
+                )
+                params.update({
+                    f"account_id_{idx}": account_id,
+                    f"bfpcode_{idx}": item_data.get('bfpcode'),
+                    f"mtime_{idx}": item_data.get('mtime'),
+                    f"bffb_{idx}": item_data.get('bffb'),
+                    f"bfpl_{idx}": item_data.get('bfpl'),
+                    f"bfpvalue_{idx}": item_data.get('bfpvalue'),
+                    f"NYID_{idx}": item_data.get('NYID'),
+                    f"sort_{idx}": item_data.get('sort')
+                })
+
+            # 批量插入SQL - 使用字符串拼接（修复TextClause拼接问题）
+            insert_sql = f"""
+                INSERT INTO `{table_name}`
+                (account_id, bfpcode, mtime, bffb, bfpl, bfpvalue, NYID, sort)
+                VALUES {", ".join(values_list)}
+            """
+            final_sql = text(insert_sql)
+            db.execute(final_sql, params)
+            total_inserted += len(batch_data)
+            logger.info(f"Inserted batch {i//batch_size + 1}: {len(batch_data)} records")
+
+        return total_inserted