from sqlalchemy.orm import Session from typing import List, Dict, Any, Optional from ..models.account import Account from ..models.section_data import SectionData from ..models.checkpoint import Checkpoint from ..models.settlement_data import SettlementData from ..models.level_data import LevelData from ..services.section_data import SectionDataService from ..services.checkpoint import CheckpointService from ..services.settlement_data import SettlementDataService from ..services.level_data import LevelDataService from ..services.account import AccountService from ..core.exceptions import DataNotFoundException, AccountNotFoundException import pandas as pd import logging from datetime import datetime logger = logging.getLogger(__name__) class ExportExcelService: def __init__(self): self.account_service = AccountService() self.section_service = SectionDataService() self.checkpoint_service = CheckpointService() self.settlement_service = SettlementDataService() self.level_service = LevelDataService() def get_field_comments(self, model_class) -> Dict[str, str]: """获取模型字段的注释信息""" comments = {} for column in model_class.__table__.columns: if column.comment: comments[column.name] = column.comment return comments def merge_settlement_with_related_data(self, db: Session, settlement_data: SettlementData, section_data: SectionData, checkpoint_data: Checkpoint, level_data: Optional[LevelData]) -> Dict[str, Any]: """ 合并沉降数据与关联数据，去除重复和id字段 """ result = {} # 沉降数据字段映射（用注释名作为键） settlement_comments = self.get_field_comments(SettlementData) settlement_dict = settlement_data.to_dict() for field_name, value in settlement_dict.items(): # 跳过id字段 if field_name == 'id': continue # 使用注释名作为键，如果没有注释则使用字段名 key = settlement_comments.get(field_name, field_name) result[key] = value # 断面数据字段映射（添加前缀） section_comments = self.get_field_comments(SectionData) section_dict = section_data.to_dict() for field_name, value in section_dict.items(): # 跳过id和account_id字段 if field_name in ['id', 'account_id']: continue key = section_comments.get(field_name, field_name) result[f"断面_{key}"] = value # 观测点数据字段映射（添加前缀） checkpoint_comments = self.get_field_comments(Checkpoint) checkpoint_dict = checkpoint_data.to_dict() for field_name, value in checkpoint_dict.items(): # 跳过id和section_id字段（section_id可能重复） if field_name in ['id', 'section_id']: continue key = checkpoint_comments.get(field_name, field_name) result[f"观测点_{key}"] = value # 水准数据字段映射（添加前缀） if level_data is not None: level_comments = self.get_field_comments(LevelData) level_dict = level_data.to_dict() for field_name, value in level_dict.items(): # 跳过id和NYID字段（NYID可能重复） if field_name in ['id', 'NYID']: continue key = level_comments.get(field_name, field_name) result[f"水准_{key}"] = value return result def export_settlement_data_to_file(self, db: Session, project_name: str, file_path: str): """ 根据项目名称导出沉降数据Excel文件到指定路径（批量查询优化版本） """ logger.info(f"开始导出项目 '{project_name}' 的沉降数据到文件: {file_path}") # 1. 在账号表查询到账号id作为account_id account_responses = self.account_service.search_accounts(db, project_name=project_name) if not account_responses: logger.warning(f"未找到项目名称为 '{project_name}' 的账号") raise AccountNotFoundException(f"未找到项目名称为 '{project_name}' 的账号") account_response = account_responses[0] account_id = str(account_response.account_id) logger.info(f"找到账号 ID: {account_id}") # 2. 通过 account_id 查询断面数据 sections = self.section_service.search_section_data(db, account_id=account_id, limit=10000) if not sections: logger.warning(f"账号 {account_id} 下未找到断面数据") raise DataNotFoundException(f"账号 {account_id} 下未找到断面数据") logger.info(f"找到 {len(sections)} 个断面") # 3. 收集所有观测点数据，建立断面->观测点映射 section_dict = {section.section_id: section for section in sections} section_checkpoint_map = {} # section_id -> [checkpoints] all_checkpoints = [] for section in sections: checkpoints = self.checkpoint_service.get_by_section_id(db, section.section_id) if checkpoints: section_checkpoint_map[section.section_id] = checkpoints all_checkpoints.extend(checkpoints) if not all_checkpoints: logger.warning("未找到任何观测点数据") raise DataNotFoundException("未找到任何观测点数据") logger.info(f"找到 {len(all_checkpoints)} 个观测点") # 4. 批量查询沉降数据（关键优化点） point_ids = [cp.point_id for cp in all_checkpoints] logger.info(f"开始批量查询 {len(point_ids)} 个观测点的沉降数据") all_settlements = self.settlement_service.get_by_point_ids(db, point_ids) if not all_settlements: logger.warning("未找到任何沉降数据") logger.info(f"观测点id集合{point_ids}") raise DataNotFoundException("未找到任何沉降数据") logger.info(f"批量查询到 {len(all_settlements)} 条沉降数据") # 5. 建立观测点->沉降数据映射 checkpoint_dict = {cp.point_id: cp for cp in all_checkpoints} point_settlement_map = {} # point_id -> [settlements] nyid_set = set() for settlement in all_settlements: if settlement.point_id not in point_settlement_map: point_settlement_map[settlement.point_id] = [] point_settlement_map[settlement.point_id].append(settlement) if settlement.NYID: nyid_set.add(settlement.NYID) # 6. 批量查询水准数据（关键优化点） nyid_list = list(nyid_set) logger.info(f"开始批量查询 {len(nyid_list)} 个期数的水准数据") all_level_data = self.level_service.get_by_nyids(db, nyid_list) logger.info(f"批量查询到 {len(all_level_data)} 条水准数据") # 建立NYID->水准数据映射 nyid_level_map = {} for level_data in all_level_data: if level_data.NYID not in nyid_level_map: nyid_level_map[level_data.NYID] = level_data # 7. 合并数据 all_settlement_records = [] for section in sections: checkpoints = section_checkpoint_map.get(section.section_id, []) for checkpoint in checkpoints: settlements = point_settlement_map.get(checkpoint.point_id, []) for settlement in settlements: # 从映射中获取水准数据 level_data = nyid_level_map.get(settlement.NYID) # 合并数据 merged_record = self.merge_settlement_with_related_data( db, settlement, section, checkpoint, level_data ) all_settlement_records.append(merged_record) if not all_settlement_records: logger.warning("未能合并任何数据记录") raise DataNotFoundException("未能合并任何数据记录") logger.info(f"共找到 {len(all_settlement_records)} 条沉降数据记录") # 转换为DataFrame df = pd.DataFrame(all_settlement_records) # 导出到Excel文件 with pd.ExcelWriter(file_path, engine='openpyxl') as writer: df.to_excel(writer, index=False, sheet_name='沉降数据') # 自动调整列宽 worksheet = writer.sheets['沉降数据'] for column in worksheet.columns: max_length = 0 column_letter = column[0].column_letter for cell in column: try: if len(str(cell.value)) > max_length: max_length = len(str(cell.value)) except: pass adjusted_width = min(max_length + 2, 50) worksheet.column_dimensions[column_letter].width = adjusted_width logger.info("Excel文件生成完成")