from sqlalchemy.orm import Session from typing import List, Optional, Dict, Any from ..models.level_data import LevelData from .base import BaseService from ..models.settlement_data import SettlementData from ..models.checkpoint import Checkpoint from ..models.section_data import SectionData from ..models.account import Account import logging logger = logging.getLogger(__name__) class LevelDataService(BaseService[LevelData]): def __init__(self): super().__init__(LevelData) def get_by_nyid(self, db: Session, nyid: str) -> List[LevelData]: """根据期数ID获取水准数据""" return self.get_by_field(db, "NYID", nyid) def get_by_nyids(self, db: Session, nyids: List[str]) -> List[LevelData]: """根据多个期数ID获取水准数据""" return db.query(LevelData).filter(LevelData.NYID.in_(nyids)).all() def get_by_linecode(self, db: Session, linecode: str) -> List[LevelData]: """根据水准线路编码获取水准数据""" return self.get_by_field(db, "linecode", linecode) def search_level_data(self, db: Session, id: Optional[str] = None, linecode: Optional[str] = None, nyid: Optional[str] = None, benchmarkids: Optional[str] = None) -> List[LevelData]: """根据多个条件搜索水准数据""" conditions = {} if linecode is not None: conditions["linecode"] = linecode if nyid is not None: conditions["NYID"] = nyid if benchmarkids is not None: conditions["benchmarkids"] = benchmarkids if id is not None: conditions["id"] = id return self.search_by_conditions(db, conditions) def get_by_nyid_and_linecode(self, db: Session, nyid: str, linecode: str = None) -> Optional[LevelData]: """根据期数ID和线路编码获取水准数据""" return db.query(LevelData).filter( LevelData.NYID == nyid if nyid else True, LevelData.linecode == linecode if linecode else True ).first() def _check_settlement_exists(self, db: Session, nyid: str) -> bool: """检查期数id沉降数据是否存在""" settlement = db.query(SettlementData).filter(SettlementData.NYID == nyid).first() return settlement is not None def batch_import_level_data(self, db: Session, data: List) -> Dict[str, Any]: """ 批量导入水准数据 - 性能优化版 使用批量查询和批量操作，大幅提升导入速度 根据期数ID和线路编码判断是否重复，重复数据改为更新操作 支持事务回滚，失败时重试一次 """ import logging logger = logging.getLogger(__name__) total_count = len(data) success_count = 0 failed_count = 0 failed_items = [] if total_count == 0: return { 'success': False, 'message': '导入数据不能为空', 'total_count': 0, 'success_count': 0, 'failed_count': 0, 'failed_items': [] } for attempt in range(2): # 最多重试1次 try: db.begin() success_count = 0 failed_count = 0 failed_items = [] # ===== 性能优化1：批量查询沉降数据（IN查询） ===== nyid_list = list(set(str(item.get('NYID')) for item in data if item.get('NYID'))) logger.info(f"Checking {len(nyid_list)} unique NYIDs in settlement data") 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: db.rollback() return { 'success': False, 'message': f'以下期数在沉降表中不存在: {list(missing_nyids)}', 'total_count': total_count, 'success_count': 0, 'failed_count': total_count, 'failed_items': [] } # ===== 性能优化2：批量查询现有水准数据（IN查询） ===== # 构建 (NYID, linecode) 组合键来查找重复数据 existing_data = db.query(LevelData).filter( LevelData.NYID.in_(nyid_list) ).all() # 使用组合键创建查找表：key = f"{NYID}_{linecode}" existing_map = { f"{item.NYID}_{item.linecode}": item for item in existing_data } logger.info(f"Found {len(existing_data)} existing level records") # ===== 性能优化3：批量处理插入和更新 ===== to_update = [] to_insert = [] for item_data in data: nyid = str(item_data.get('NYID')) linecode = item_data.get('linecode') # 构建组合键 key = f"{nyid}_{linecode}" if key in existing_map: # 记录需要更新的数据 existing_item = existing_map[key] to_update.append((existing_item, item_data)) else: # 记录需要插入的数据 to_insert.append(item_data) # ===== 执行批量更新 ===== if to_update: logger.info(f"Updating {len(to_update)} existing records") for existing_item, item_data in to_update: try: existing_item.benchmarkids = item_data.get('benchmarkids') existing_item.wsphigh = item_data.get('wsphigh') existing_item.mtype = item_data.get('mtype') existing_item.createDate = item_data.get('createDate') success_count += 1 except Exception as e: failed_count += 1 failed_items.append({ 'data': item_data, 'error': f'更新失败: {str(e)}' }) logger.error(f"Failed to update level data: {str(e)}") raise e # ===== 执行批量插入 ===== if to_insert: logger.info(f"Inserting {len(to_insert)} new records") # 分批插入，每批500条（避免SQL过长） batch_size = 500 for i in range(0, len(to_insert), batch_size): batch = to_insert[i:i + batch_size] try: level_data_list = [ LevelData( linecode=item.get('linecode'), benchmarkids=item.get('benchmarkids'), wsphigh=item.get('wsphigh'), mtype=item.get('mtype'), NYID=str(item.get('NYID')), createDate=item.get('createDate') ) for item in batch ] db.add_all(level_data_list) success_count += len(batch) logger.info(f"Inserted batch {i//batch_size + 1}: {len(batch)} records") except Exception as e: failed_count += len(batch) failed_items.extend([ { 'data': item, 'error': f'插入失败: {str(e)}' } for item in batch ]) logger.error(f"Failed to insert batch: {str(e)}") raise e # 如果有失败记录，不提交事务 if failed_items: db.rollback() return { 'success': False, 'message': f'批量导入失败: {len(failed_items)}条记录处理失败', 'total_count': total_count, 'success_count': success_count, 'failed_count': failed_count, 'failed_items': failed_items } db.commit() logger.info(f"Batch import level 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 level 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 } def get_level_data_by_project_name(self, db: Session, project_name: str, nyid_max: bool = False) -> List[Dict[str, Any]]: """ 通过project_name获取全部水准线路 业务逻辑: 1. 查询账号表获取账号数据 (通过project_name) 2. 查询断面表获取断面数据 (通过account_id) 3. 查询观测点表获取观测点数据 (通过section_id) 4. 查询沉降数据表获取沉降数据 (通过point_id) 5. 查询水准数据表获取水准数据 (通过NYID) 6. 将水准数据依照linecode去重（同linecode只需保留一个） """ try: logger.info(f"开始查询project_name={project_name}对应的水准线路数据") # 1. 查询账号表获取账号数据 accounts = db.query(Account).filter(Account.project_name.like(f"%{project_name}%")).all() if not accounts: logger.warning(f"未查询到project_name={project_name}对应的账号") return [] account_ids = [str(account.id) for account in accounts] logger.info(f"查询到{len(account_ids)}个账号: {account_ids}") # 2. 查询断面表获取断面数据 (通过account_id) sections = db.query(SectionData).filter(SectionData.account_id.in_(account_ids)).all() if not sections: logger.warning(f"未查询到对应的断面数据") return [] section_ids = [section.section_id for section in sections] logger.info(f"查询到{len(section_ids)}个断面: {section_ids}") # 3. 查询观测点表获取观测点数据 (通过section_id) checkpoints = db.query(Checkpoint).filter(Checkpoint.section_id.in_(section_ids)).all() if not checkpoints: logger.warning(f"未查询到对应的观测点数据") return [] point_ids = [checkpoint.point_id for checkpoint in checkpoints] logger.info(f"查询到{len(point_ids)}个观测点") # 4. 查询沉降数据表获取沉降数据 (通过point_id) settlements = db.query(SettlementData).filter(SettlementData.point_id.in_(point_ids)).all() if not settlements: logger.warning(f"未查询到对应的沉降数据") return [] nyid_list = list(set([settlement.NYID for settlement in settlements if settlement.NYID])) logger.info(f"查询到{len(nyid_list)}个期数ID") if nyid_max: # 只获取最新期数的水准数据 nyid_list = [max(nyid_list, key=int)] logger.info(f"筛选后只获取最新期数ID: {nyid_list}") level_data_list = db.query(LevelData).filter(LevelData.NYID.in_(nyid_list)).all() else: # 5. 查询水准数据表获取水准数据 (通过NYID) level_data_list = db.query(LevelData).filter(LevelData.NYID.in_(nyid_list)).all() if not level_data_list: logger.warning(f"未查询到对应的水准数据") return [] # 6. 将水准数据依照linecode去重（同linecode只需保留一个） linecode_seen = set() unique_level_data = [] for level in level_data_list: if level.linecode not in linecode_seen: linecode_seen.add(level.linecode) level_dict = { "id": level.id, "linecode": level.linecode, "benchmarkids": level.benchmarkids, "wsphigh": level.wsphigh, "NYID": level.NYID, "mtype": level.mtype, "createDate": level.createDate.strftime("%Y-%m-%d %H:%M:%S") if level.createDate else None } unique_level_data.append(level_dict) logger.info(f"查询完成，共{len(unique_level_data)}条去重后的水准数据") return unique_level_data except Exception as e: logger.error(f"查询project_name={project_name}的水准数据失败: {str(e)}", exc_info=True) raise e