AcidMap/app/api/vector.py

# 矢量数据API
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, BackgroundTasks, Query
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from ..database import get_db
from ..services import vector_service
from ..services.admin_boundary_service import get_boundary_geojson_by_name
import os
import shutil
from typing import List, Optional

router = APIRouter()

@router.post("/import", summary="导入GeoJSON文件", description="将GeoJSON文件导入到数据库中")
async def import_vector_data(file: UploadFile = File(...), db: Session = Depends(get_db)):
    """导入GeoJSON文件到数据库"""
    # 检查文件类型
    if not file.filename.endswith('.geojson'):
        raise HTTPException(status_code=400, detail="只支持GeoJSON文件")
    
    return await vector_service.import_vector_data(file, db)

@router.get("/boundary", summary="按名称获取边界", description="按县/市/省名称获取边界GeoJSON")
def get_boundary(name: str = Query(...), level: str = Query("auto"), db: Session = Depends(get_db)):
    """按名称和层级获取边界。level: county/city/province/auto"""
    try:
        return get_boundary_geojson_by_name(db, name, level)
    except Exception as e:
        raise HTTPException(status_code=404, detail=str(e))

@router.get("/{vector_id}", summary="获取矢量数据", description="根据ID获取一条矢量数据记录")
def get_vector_data(vector_id: int, db: Session = Depends(get_db)):
    """获取指定ID的矢量数据"""
    return vector_service.get_vector_data(db, vector_id)

@router.get("/{vector_id}/export", summary="导出矢量数据", description="将指定ID的矢量数据导出为GeoJSON文件")
async def export_vector_data(vector_id: int, db: Session = Depends(get_db)):
    """导出指定ID的矢量数据为GeoJSON文件"""
    result = vector_service.export_vector_data(db, vector_id)
    
    # 检查文件是否存在
    if not os.path.exists(result["file_path"]):
        if "temp_dir" in result and os.path.exists(result["temp_dir"]):
            shutil.rmtree(result["temp_dir"])
        raise HTTPException(status_code=404, detail="导出文件不存在")
    
    # 返回文件下载
    return FileResponse(
        path=result["file_path"],
        filename=os.path.basename(result["file_path"]),
        media_type="application/json",
        background=BackgroundTasks().add_task(shutil.rmtree, result["temp_dir"]) if "temp_dir" in result else None
    )

@router.get("/export/all", summary="导出矢量数据", description="将指定的矢量数据表导出为GeoJSON或JSON文件")
async def export_all_vector_data_api(table_name: str = "vector_data", db: Session = Depends(get_db),format:Optional[str] = "geojson"):
    """导出指定矢量数据表为GeoJSON文件或普通JSON文件
    
    Args:
        table_name (str): 要导出的矢量数据表名，默认为'vector_data'
        db (Session): 数据库会话
        format (str,optional):导出格式支持'geojson'和'json'，默认为'geojson'
        
    Returns:
        FileResponse: GeoJSON文件下载响应
    """
    #验证格式参数
    if format not in ["geojson" , "json"]:
        raise HTTPException(
            status_code=400,
            detail="只支持'geojson'和'json'格式"
        )
    result = vector_service.export_all_vector_data(db, table_name,export_format=format)#传递导出格式
    
    # 检查文件是否存在
    if not os.path.exists(result["file_path"]):
        if "temp_dir" in result and os.path.exists(result["temp_dir"]):
            shutil.rmtree(result["temp_dir"])
        raise HTTPException(status_code=404, detail="导出文件不存在")
    
    #根据格式设置正确的媒体类型
    media_type = "application/geo+json" if format == "geojson" else "application/json"

    # 返回文件下载
    return FileResponse(
        path=result["file_path"],
        filename=os.path.basename(result["file_path"]),
        media_type=media_type,
        background=BackgroundTasks().add_task(shutil.rmtree, result["temp_dir"]) if "temp_dir" in result else None
    )

@router.post("/export/batch", summary="批量导出矢量数据", description="将多个ID的矢量数据批量导出为GeoJSON文件")
async def export_vector_data_batch(vector_ids: List[int], db: Session = Depends(get_db)):
    """批量导出矢量数据为GeoJSON文件"""
    result = vector_service.export_vector_data_batch(db, vector_ids)
    
    # 检查文件是否存在
    if not os.path.exists(result["file_path"]):
        if "temp_dir" in result and os.path.exists(result["temp_dir"]):
            shutil.rmtree(result["temp_dir"])
        raise HTTPException(status_code=404, detail="导出文件不存在")
    
    # 返回文件下载
    return FileResponse(
        path=result["file_path"],
        filename=os.path.basename(result["file_path"]),
        media_type="application/json",
        background=BackgroundTasks().add_task(shutil.rmtree, result["temp_dir"]) if "temp_dir" in result else None
    )