AcidMap/app/api/raster.py

from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from ..database import get_db
from ..services import raster_service
import os
import shutil
from fastapi import BackgroundTasks
from typing import List
from pydantic import BaseModel

router = APIRouter()

class RasterBatchExportRequest(BaseModel):
    raster_ids: List[int]

@router.get("/{raster_id}", summary="获取栅格数据", description="获取指定ID的栅格数据信息")
async def get_raster(raster_id: int, db: Session = Depends(get_db)):
    """获取指定ID的栅格数据信息"""
    return raster_service.get_raster_data(db, raster_id)

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

@router.get("/{raster_id}/export", summary="导出栅格数据", description="将指定ID的栅格数据导出为TIFF文件")
async def export_raster_data(raster_id: int, db: Session = Depends(get_db)):
    """导出指定ID的栅格数据为TIFF文件"""
    result = raster_service.export_raster_data(db, raster_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=f"raster_{raster_id}.tif",
        media_type="image/tiff",
        background=BackgroundTasks().add_task(shutil.rmtree, result["temp_dir"]) if "temp_dir" in result else None
    )

@router.post("/export/batch", summary="批量导出栅格数据", description="将多个ID的栅格数据批量导出为TIFF文件")
async def export_raster_data_batch(request: RasterBatchExportRequest, db: Session = Depends(get_db)):
    """批量导出栅格数据为TIFF文件"""
    result = raster_service.export_raster_data_batch(db, request.raster_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/zip",
        background=BackgroundTasks().add_task(shutil.rmtree, result["temp_dir"]) if "temp_dir" in result else None
    )