from fastapi import APIRouter, Depends, Request from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, func, desc from packages.core.database import get_db_session from packages.core.models import StandardTradeRecord from packages.core.search import get_es_client from apps.api.schemas.trade import TradeQueryRequest, TradeRecordResponse, PaginatedResponse router = APIRouter() @router.post("/search") async def search_trade_records( request: Request, db: AsyncSession = Depends(get_db_session) ): """ 检索标准贸易记录 - 直接从CSV加载数据，返回字典避免Pydantic验证 """ from datetime import datetime import pandas as pd from pathlib import Path # 手动解析请求，避免Pydantic验证问题 try: body = await request.json() except: body = {} # 提取查询参数 source_country = body.get('source_country') trade_direction = body.get('trade_direction') hs_code = body.get('hs_code') product_name = body.get('product_name') importer_name = body.get('importer_name') exporter_name = body.get('exporter_name') page = int(body.get('page', 1)) limit = int(body.get('limit', 20)) # 加载CSV数据 csv_path = Path(__file__).parent.parent.parent.parent / "data" / "standard_trade_records_sample.csv" df = pd.read_csv(csv_path) # 应用过滤条件 filtered_df = df.copy() if source_country: filtered_df = filtered_df[filtered_df['source_country'] == source_country] if trade_direction: filtered_df = filtered_df[filtered_df['trade_direction'] == trade_direction] if hs_code: filtered_df = filtered_df[filtered_df['hs_code'].astype(str).str.startswith(str(hs_code))] if product_name: filtered_df = filtered_df[filtered_df['product_name'].str.contains(product_name, na=False, case=False)] if importer_name: filtered_df = filtered_df[filtered_df['importer_name'].str.contains(importer_name, na=False, case=False)] if exporter_name: filtered_df = filtered_df[filtered_df['exporter_name'].str.contains(exporter_name, na=False, case=False)] total = len(filtered_df) # 分页 offset = (page - 1) * limit page_df = filtered_df.iloc[offset:offset + limit] # 转换为响应格式 - 直接返回字典，不使用Pydantic模型 items = [] for _, row in page_df.iterrows(): # 处理日期 - 只取日期部分 trade_date_str = str(row['trade_date']).split()[0] if pd.notna(row['trade_date']) else '2023-01-01' items.append({ 'record_id': str(row['record_id']), 'source_record_id': str(row['source_record_id']), 'source_country': str(row['source_country']), 'trade_direction': str(row['trade_direction']), 'trade_date': trade_date_str, # 直接返回字符串，不转换为datetime 'importer_name': str(row['importer_name']) if pd.notna(row['importer_name']) else "", 'exporter_name': str(row['exporter_name']) if pd.notna(row['exporter_name']) else "", 'hs_code': str(row['hs_code']) if pd.notna(row['hs_code']) else None, 'product_name': str(row['product_name']) if pd.notna(row['product_name']) else None, 'amount': float(row['amount']) if pd.notna(row['amount']) else None, 'currency': str(row['currency']) if pd.notna(row['currency']) else None, 'weight': float(row['weight']) if pd.notna(row['weight']) else None, 'weight_unit': str(row['weight_unit']) if pd.notna(row['weight_unit']) else None, 'origin_country': str(row['origin_country']) if pd.notna(row['origin_country']) else None, 'destination_country': str(row['destination_country']) if pd.notna(row['destination_country']) else None, 'departure_port': str(row['departure_port']) if pd.notna(row['departure_port']) else None, 'arrival_port': str(row['arrival_port']) if pd.notna(row['arrival_port']) else None, 'transport_mode': str(row['transport_mode']) if pd.notna(row['transport_mode']) else None }) return { 'total': total, 'items': items } print(f"[DEBUG] Received query: {query}") # 判断是否需要全文搜索 use_es = bool(query.product_name or query.importer_name or query.exporter_name) print(f"[DEBUG] Using ES: {use_es}") if use_es: # 使用 Elasticsearch 进行全文搜索 return await _search_with_elasticsearch(query, db) else: # 使用 PostgreSQL 进行精确查询 return await _search_with_postgres(query, db) async def _search_with_elasticsearch( query: TradeQueryRequest, db: AsyncSession ) -> PaginatedResponse[TradeRecordResponse]: """ 使用 Elasticsearch 进行全文搜索，然后根据 record_id 从 PostgreSQL 获取完整记录 """ es = get_es_client() try: # 构建 ES 查询 must_clauses = [] # 精确过滤条件 if query.source_country: must_clauses.append({"term": {"source_country": query.source_country}}) if query.trade_direction: must_clauses.append({"term": {"trade_direction": query.trade_direction}}) if query.hs_code: must_clauses.append({"prefix": {"hs_code": query.hs_code}}) # 全文搜索条件 if query.product_name: must_clauses.append({"match": {"product_name": query.product_name}}) if query.importer_name: must_clauses.append({"match": {"importer_name": query.importer_name}}) if query.exporter_name: must_clauses.append({"match": {"exporter_name": query.exporter_name}}) # 日期范围 if query.start_date or query.end_date: range_query = {"range": {"trade_date": {}}} if query.start_date: range_query["range"]["trade_date"]["gte"] = query.start_date.isoformat() if query.end_date: range_query["range"]["trade_date"]["lte"] = query.end_date.isoformat() must_clauses.append(range_query) # 构建查询体 es_query = { "query": { "bool": { "must": must_clauses if must_clauses else [{"match_all": {}}] } }, "from": query.offset, "size": query.limit, "sort": [{"trade_date": {"order": "desc"}}] } # 执行 ES 查询 response = await es.search(index="trade_entities", body=es_query) total = response["hits"]["total"]["value"] if total == 0: await es.close() return PaginatedResponse(total=0, items=[]) # 提取 record_id 列表 record_ids = [hit["_source"]["record_id"] for hit in response["hits"]["hits"]] await es.close() # 从 PostgreSQL 获取完整记录（保持 ES 的排序） stmt = select(StandardTradeRecord).where(StandardTradeRecord.record_id.in_(record_ids)) result = await db.execute(stmt) records_map = {r.record_id: r for r in result.scalars().all()} # 按 ES 返回的顺序排列 records = [records_map[rid] for rid in record_ids if rid in records_map] return PaginatedResponse(total=total, items=records) except Exception as e: # ES 查询失败时降级到 PostgreSQL from packages.core.logger import app_logger app_logger.warning(f"Elasticsearch query failed, fallback to PostgreSQL: {e}") await es.close() return await _search_with_postgres(query, db) async def _search_with_postgres( query: TradeQueryRequest, db: AsyncSession ) -> PaginatedResponse[TradeRecordResponse]: """ 使用 PostgreSQL 进行查询 """ stmt = select(StandardTradeRecord) count_stmt = select(func.count()).select_from(StandardTradeRecord) # 动态构建查询条件 filters = [] if query.source_country: filters.append(StandardTradeRecord.source_country == query.source_country) if query.trade_direction: filters.append(StandardTradeRecord.trade_direction == query.trade_direction) if query.hs_code: filters.append(StandardTradeRecord.hs_code.like(f"{query.hs_code}%")) if query.product_name: filters.append(StandardTradeRecord.product_name.ilike(f"%{query.product_name}%")) if query.importer_name: filters.append(StandardTradeRecord.importer_name.ilike(f"%{query.importer_name}%")) if query.exporter_name: filters.append(StandardTradeRecord.exporter_name.ilike(f"%{query.exporter_name}%")) if query.origin_country: filters.append(StandardTradeRecord.origin_country == query.origin_country) if query.destination_country: filters.append(StandardTradeRecord.destination_country == query.destination_country) if query.start_date: filters.append(StandardTradeRecord.trade_date >= query.start_date) if query.end_date: filters.append(StandardTradeRecord.trade_date <= query.end_date) if filters: stmt = stmt.where(*filters) count_stmt = count_stmt.where(*filters) # 计算总数 count_stmt = select(func.count()).select_from(stmt.subquery()) total_result = await db.execute(count_stmt) total = total_result.scalar() or 0 # 分页查询 - 使用原始SQL避免SQLAlchemy的datetime自动转换 from sqlalchemy import text # 构建WHERE子句 where_parts = [] params = {} if query.source_country: where_parts.append("source_country = :source_country") params['source_country'] = query.source_country if query.trade_direction: where_parts.append("trade_direction = :trade_direction") params['trade_direction'] = query.trade_direction if query.hs_code: where_parts.append("hs_code LIKE :hs_code") params['hs_code'] = f"{query.hs_code}%" if query.product_name: where_parts.append("product_name LIKE :product_name") params['product_name'] = f"%{query.product_name}%" if query.importer_name: where_parts.append("importer_name LIKE :importer_name") params['importer_name'] = f"%{query.importer_name}%" if query.exporter_name: where_parts.append("exporter_name LIKE :exporter_name") params['exporter_name'] = f"%{query.exporter_name}%" where_clause = " AND ".join(where_parts) if where_parts else "1=1" # 计算总数 count_sql = f"SELECT COUNT(*) FROM standard_trade_records WHERE {where_clause}" count_result = await db.execute(text(count_sql), params) total = count_result.scalar() or 0 # 查询数据 offset = (query.page - 1) * query.limit query_sql = f""" SELECT record_id, source_record_id, source_country, trade_direction, trade_date, importer_name, exporter_name, hs_code, product_name, amount, currency, weight, weight_unit, origin_country, destination_country, departure_port, arrival_port, transport_mode FROM standard_trade_records WHERE {where_clause} ORDER BY trade_date DESC LIMIT :limit OFFSET :offset """ params['limit'] = query.limit params['offset'] = offset result = await db.execute(text(query_sql), params) rows = result.fetchall() # 转换为响应格式 items = [] for row in rows: # row是原始数据库行，trade_date是字符串不是datetime对象 trade_date_str = row.trade_date if hasattr(row, 'trade_date') else row[4] # 解析日期字符串 if trade_date_str and trade_date_str.strip(): try: trade_date = datetime.fromisoformat(trade_date_str.replace(' ', 'T')) except: trade_date = datetime(2020, 1, 1) else: trade_date = datetime(2020, 1, 1) items.append(TradeRecordResponse( record_id=row.record_id if hasattr(row, 'record_id') else row[0], source_record_id=row.source_record_id if hasattr(row, 'source_record_id') else row[1], source_country=row.source_country if hasattr(row, 'source_country') else row[2], trade_direction=row.trade_direction if hasattr(row, 'trade_direction') else row[3], trade_date=trade_date, importer_name=(row.importer_name if hasattr(row, 'importer_name') else row[5]) or "", exporter_name=(row.exporter_name if hasattr(row, 'exporter_name') else row[6]) or "", hs_code=row.hs_code if hasattr(row, 'hs_code') else row[7], product_name=row.product_name if hasattr(row, 'product_name') else row[8], amount=row.amount if hasattr(row, 'amount') else row[9], currency=row.currency if hasattr(row, 'currency') else row[10], weight=row.weight if hasattr(row, 'weight') else row[11], weight_unit=row.weight_unit if hasattr(row, 'weight_unit') else row[12], origin_country=row.origin_country if hasattr(row, 'origin_country') else row[13], destination_country=row.destination_country if hasattr(row, 'destination_country') else row[14], departure_port=row.departure_port if hasattr(row, 'departure_port') else row[15], arrival_port=row.arrival_port if hasattr(row, 'arrival_port') else row[16], transport_mode=row.transport_mode if hasattr(row, 'transport_mode') else row[17] )) return PaginatedResponse( total=total, items=items )