TJWaterServerBinary/app/api/v1/endpoints/network/pumps.py

from fastapi import APIRouter, Request, Query, Path, Body
from typing import Any, List, Dict, Union
from app.services.tjnetwork import (
    Any,
    ChangeSet,
    add_pump,
    delete_pump,
    get_all_pumps,
    get_pump,
    get_pump_schema,
    set_pump,
)

router = APIRouter()

@router.get("/getpumpschema", summary="获取水泵模式", description="获取水泵对象的模式定义，包含所有可用字段及其类型")
async def fastapi_get_pump_schema(
    network: str = Query(..., description="管网名称（或数据库名称）")
) -> dict[str, dict[str, Any]]:
    """
    获取水泵数据模式定义。
    
    Args:
        network: 管网名称（或数据库名称）
        
    Returns:
        包含水泵模式信息的字典
    """
    return get_pump_schema(network)

@router.post("/addpump/", response_model=None, summary="添加水泵", description="向网络中添加新的水泵，需要提供水泵的基本参数如功率等")
async def fastapi_add_pump(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵标识符"),
    node1: str = Query(..., description="水泵起始节点ID"),
    node2: str = Query(..., description="水泵终止节点ID"),
    power: float = Query(0.0, description="水泵功率（单位：千瓦）")
) -> ChangeSet:
    """
    添加新水泵到网络。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        node1: 起始节点ID
        node2: 终止节点ID
        power: 水泵功率
        
    Returns:
        ChangeSet对象，包含本次操作的变更信息
    """
    ps = {"id": pump, "node1": node1, "node2": node2, "power": power}
    return add_pump(network, ChangeSet(ps))

@router.post("/deletepump/", response_model=None, summary="删除水泵", description="从网络中删除指定的水泵")
async def fastapi_delete_pump(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="要删除的水泵ID")
) -> ChangeSet:
    """
    删除水泵。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        
    Returns:
        ChangeSet对象，包含本次删除操作的变更信息
    """
    ps = {"id": pump}
    return delete_pump(network, ChangeSet(ps))

@router.get("/getpumpnode1/", summary="获取水泵起始节点", description="获取指定水泵的起始节点ID")
async def fastapi_get_pump_node1(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID")
) -> str | None:
    """
    获取水泵的起始节点。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        
    Returns:
        起始节点ID，如果不存在则返回None
    """
    ps = get_pump(network, pump)
    return ps["node1"]

@router.get("/getpumpnode2/", summary="获取水泵终止节点", description="获取指定水泵的终止节点ID")
async def fastapi_get_pump_node2(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID")
) -> str | None:
    """
    获取水泵的终止节点。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        
    Returns:
        终止节点ID，如果不存在则返回None
    """
    ps = get_pump(network, pump)
    return ps["node2"]

@router.post("/setpumpnode1/", response_model=None, summary="设置水泵起始节点", description="设置指定水泵的起始节点")
async def fastapi_set_pump_node1(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID"),
    node1: str = Query(..., description="新的起始节点ID")
) -> ChangeSet:
    """
    设置水泵起始节点。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        node1: 新的起始节点ID
        
    Returns:
        ChangeSet对象，包含本次修改的变更信息
    """
    ps = {"id": pump, "node1": node1}
    return set_pump(network, ChangeSet(ps))

@router.post("/setpumpnode2/", response_model=None, summary="设置水泵终止节点", description="设置指定水泵的终止节点")
async def fastapi_set_pump_node2(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID"),
    node2: str = Query(..., description="新的终止节点ID")
) -> ChangeSet:
    """
    设置水泵终止节点。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        node2: 新的终止节点ID
        
    Returns:
        ChangeSet对象，包含本次修改的变更信息
    """
    ps = {"id": pump, "node2": node2}
    return set_pump(network, ChangeSet(ps))

@router.get("/getpumpproperties/", summary="获取水泵属性", description="获取指定水泵的所有属性信息")
async def fastapi_get_pump_properties(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID")
) -> dict[str, Any]:
    """
    获取水泵的所有属性。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        
    Returns:
        包含水泵所有属性的字典
    """
    return get_pump(network, pump)

@router.get("/getallpumpproperties/", summary="获取所有水泵属性", description="获取网络中所有水泵的属性信息列表")
async def fastapi_get_all_pump_properties(
    network: str = Query(..., description="管网名称（或数据库名称）")
) -> list[dict[str, Any]]:
    """
    获取网络中所有水泵的属性。
    
    Args:
        network: 管网名称（或数据库名称）
        
    Returns:
        包含所有水泵属性的字典列表
    """
    # 缓存查询结果提高性能
    # global redis_client
    results = get_all_pumps(network)
    return results

@router.post("/setpumpproperties/", response_model=None, summary="设置水泵属性", description="批量设置指定水泵的多个属性")
async def fastapi_set_pump_properties(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    pump: str = Query(..., description="水泵ID"),
    req: Request = None
) -> ChangeSet:
    """
    批量设置水泵属性。
    
    Args:
        network: 管网名称（或数据库名称）
        pump: 水泵ID
        req: 请求体，包含要设置的属性及其值
        
    Returns:
        ChangeSet对象，包含本次修改的变更信息
    """
    props = await req.json()
    ps = {"id": pump} | props
    return set_pump(network, ChangeSet(ps))