Merge branch 'main' of http://www.yuxindazhineng.com:3000/lhx/project

app/services/original_data.py

@@ -337,6 +337,11 @@ class OriginalDataService(BaseService[OriginalData]):
 
         table_name = self._get_table_name(account_id)
 
+        # 检查是否已在事务中（避免重复开始事务）
+        in_transaction = db.in_transaction()
+
+        # 如果不在事务中，才需要手动管理事务
+        if not in_transaction:
             for attempt in range(2):  # 最多重试1次
                 try:
                     db.begin()
@@ -344,46 +349,104 @@ class OriginalDataService(BaseService[OriginalData]):
                     failed_count = 0
                     failed_items = []
 
-                nyid = str(data[0].get('NYID'))  # 统一转换为字符串
+                    # 执行数据导入操作
-                # 检查该期数数据是否已存在
+                    success_count = self._execute_import(db, table_name, data, account_id)
-                check_query = text(f"SELECT COUNT(*) as cnt FROM `{table_name}` WHERE NYID = :nyid")
+                    db.commit()
-                is_exists = db.execute(check_query, {"nyid": nyid}).fetchone()[0]
+                    logger.info(f"Batch import original data completed. Success: {success_count}, Failed: {failed_count}")
+                    break
 
-                if is_exists > 0:
+                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:
+                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:
+                logger.error(f"Batch import failed in existing transaction: {str(e)}")
+                # 抛出异常，让外部处理事务回滚
+                raise
+
         return {
             'success': True,
-                        'message': '数据已存在',
+            'message': '批量导入完成' if failed_count == 0 else f'部分导入失败',
-                        'total_count': 0,
+            'total_count': total_count,
             '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:
+        """执行数据导入操作（抽取的公共逻辑）"""
+        logger.info(f"Starting batch import into {table_name}, total records: {len(data)}")
+
         # ===== 性能优化：批量查询沉降数据 =====
         # 统一转换为字符串处理（数据库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:
-                    logger.warning(f"[批量导入原始数据] 批量查询settlement数据失败， Nyid: {list(missing_nyids)}")
+            raise Exception(f'以下期数在沉降表中不存在: {list(missing_nyids)}')
-                    db.rollback()
+
-                    return {
+        # ===== 性能优化：批量查询现有原始数据（IN查询）=====
-                        'success': False,
+        # 使用组合键 (NYID, sort) 查询现有数据，过滤重复数据
-                        'message': f'以下期数在沉降表中不存在: {list(missing_nyids)}',
+        existing_data = db.query(text("*")).from_statement(
-                        'total_count': total_count,
+            text(f"SELECT * FROM `{table_name}` WHERE account_id = :account_id")
-                        'success_count': 0,
+        ).params(account_id=account_id).all()
-                        'failed_count': total_count,
+
-                        'failed_items': []
+        # 使用组合键创建查找表：key = f"{NYID}_{sort}"
-                    }
+        existing_map = {
+            f"{item[7]}_{item[8]}": item  # NYID是第8个字段（索引7），sort是第9个字段（索引8）
+            for item in existing_data
+        }
+        logger.info(f"Found {len(existing_data)} existing records in {table_name}")
+
+        # ===== 批量处理插入和跳过 =====
+        to_insert = []
+        skipped_count = 0
+
+        for item_data in data:
+            nyid = str(item_data.get('NYID'))  # 统一转换为字符串
+            sort = item_data.get('sort')
+
+            # 构建组合键
+            key = f"{nyid}_{sort}"
+
+            if key in existing_map:
+                # 数据已存在，跳过
+                # logger.info(f"Skip existing data: NYID {nyid}, sort {sort}")
+                skipped_count += 1
+            else:
+                # 记录需要插入的数据
+                to_insert.append(item_data)
+            # logger.info(f"Skip existing data: {skipped_count} duplicates found so far")
+
+        logger.info(f"Filtered {skipped_count} duplicate records, {len(to_insert)} new records to insert")
+
+        # ===== 执行批量插入 =====
+        if not to_insert:
+            logger.info("No new records to insert, all data already exists")
+            return 0
 
-                # ===== 性能优化：使用批量插入 =====
         # 将数据分组，每组1000条（MySQL默认支持）
         batch_size = 1000
-                for i in range(0, len(data), batch_size):
+        total_inserted = 0
-                    batch_data = data[i:i + batch_size]
+        for i in range(0, len(to_insert), batch_size):
+            batch_data = to_insert[i:i + batch_size]
 
             # 构建批量参数
             values_list = []
@@ -412,32 +475,8 @@ class OriginalDataService(BaseService[OriginalData]):
             """
             final_sql = text(insert_sql)
             db.execute(final_sql, params)
-                    success_count += len(batch_data)
+            total_inserted += len(batch_data)
             logger.info(f"Inserted batch {i//batch_size + 1}: {len(batch_data)} records")
 
-                db.commit()
+        logger.info(f"Batch import completed: {total_inserted} records inserted, {skipped_count} duplicates skipped")
-                logger.info(f"Batch import original data completed. Success: {success_count}, Failed: {failed_count}")
+        return total_inserted
-                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
-                    }
-
-        return {
-            'success': True,
-            'message': '批量导入完成' if failed_count == 0 else f'部分导入失败',
-            'total_count': total_count,
-            'success_count': success_count,
-            'failed_count': failed_count,
-            'failed_items': failed_items
-        }