Compare commits
10
Commits
9f7f5536d7
...
18253f2fe0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
18253f2fe0 | ||
|
|
02a8686222 | ||
|
|
50d823ca58 | ||
|
|
9f225374de | ||
|
|
9b16e4e0a5 | ||
|
|
810a39a1dc | ||
|
|
e49d21cd1b | ||
|
|
0a2ce81753 | ||
|
|
b8f2f70013 | ||
|
|
6f9c94a4dd |
+1
-1
@@ -13,7 +13,7 @@ temp/
|
||||
data/
|
||||
# db_inp/
|
||||
inp/
|
||||
# .env
|
||||
.env
|
||||
*.pyc
|
||||
*.dump
|
||||
app/algorithms/health/model/my_survival_forest_model_quxi.joblib
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
# tjwater-server Customer Image Packaging Notes
|
||||
|
||||
This repository is the customer-delivery backend package. Build delivery images from this repository root.
|
||||
|
||||
## Image Build
|
||||
|
||||
Build the customer backend image:
|
||||
|
||||
```bash
|
||||
docker build -t tjwater-server:latest .
|
||||
```
|
||||
|
||||
The `Dockerfile` already performs source encapsulation in the builder stage:
|
||||
|
||||
```bash
|
||||
python scripts/compile.py
|
||||
python scripts/compile.py --delete-source
|
||||
```
|
||||
|
||||
The compiled and source-deleted `app/` directory is then copied into the final runner stage. The source deletion happens inside the Docker builder layer only; it does not delete local workspace files.
|
||||
|
||||
## Encapsulation Scope
|
||||
|
||||
`scripts/compile.py` defaults to compiling these sensitive areas:
|
||||
|
||||
- `app/services`
|
||||
- `app/native/wndb`
|
||||
- `app/algorithms`
|
||||
- `app/infra/epanet/epanet.py`
|
||||
|
||||
These areas should not contain uncompiled `.py` source files in the delivered image. Public entry points, API route wiring, schemas, configuration, and operational files may remain readable when required for runtime.
|
||||
|
||||
## Verification
|
||||
|
||||
After building, verify the image tag:
|
||||
|
||||
```bash
|
||||
docker image ls tjwater-server
|
||||
```
|
||||
|
||||
Verify that core source files were removed and compiled extensions exist:
|
||||
|
||||
```bash
|
||||
docker run --rm --entrypoint sh tjwater-server:latest -c "\
|
||||
printf 'core_py_count='; \
|
||||
find /app/app/services /app/app/native/wndb /app/app/algorithms /app/app/infra/epanet/epanet.py -name '*.py' 2>/dev/null | wc -l; \
|
||||
printf 'core_so_count='; \
|
||||
find /app/app/services /app/app/native/wndb /app/app/algorithms /app/app/infra/epanet -name '*.so' 2>/dev/null | wc -l; \
|
||||
python -c 'import app.main; print(\"import_app_main=ok\")'"
|
||||
```
|
||||
|
||||
Expected delivery result:
|
||||
|
||||
```text
|
||||
core_py_count=0
|
||||
core_so_count=<non-zero>
|
||||
import_app_main=ok
|
||||
```
|
||||
|
||||
Warnings from third-party packages during import are not necessarily build failures. Treat non-zero exit codes, failed imports, or leftover core `.py` files as blockers.
|
||||
|
||||
## Export For Windows Delivery
|
||||
|
||||
Export the image tarball to the Windows desktop from WSL:
|
||||
|
||||
```bash
|
||||
docker save -o /mnt/c/Users/admin/Desktop/tjwater-server-latest.tar tjwater-server:latest
|
||||
```
|
||||
|
||||
Adjust the Windows username if needed. To find available desktop paths:
|
||||
|
||||
```bash
|
||||
find /mnt/c/Users -maxdepth 2 -type d \( -name Desktop -o -name 桌面 \) 2>/dev/null
|
||||
```
|
||||
|
||||
On the target machine, import the image with:
|
||||
|
||||
```bash
|
||||
docker load -i tjwater-server-latest.tar
|
||||
```
|
||||
|
||||
## Deployment Caution
|
||||
|
||||
The development `infra/docker/docker-compose.yml` bind-mounts local source:
|
||||
|
||||
```yaml
|
||||
volumes:
|
||||
- ../../app:/app/app
|
||||
- ../../resources:/app/resources
|
||||
```
|
||||
|
||||
Do not use that source mount for customer delivery of the encapsulated image. Mounting `../../app` over `/app/app` replaces the compiled code inside the image with local source files and defeats the encapsulation. For delivery compose files, use the built image directly and mount only required runtime data/config paths.
|
||||
|
||||
## Local Safety
|
||||
|
||||
Do not run this destructive command in the normal working tree:
|
||||
|
||||
```bash
|
||||
python scripts/compile.py --delete-source
|
||||
```
|
||||
|
||||
Only use `--delete-source` in Docker builder stages or in a disposable delivery copy. The normal delivery image build already handles this safely.
|
||||
+15
-3
@@ -1,4 +1,4 @@
|
||||
FROM condaforge/miniforge3:latest
|
||||
FROM condaforge/miniforge3:latest AS runtime-base
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
@@ -14,13 +14,25 @@ COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir uv
|
||||
RUN uv pip install --system --no-cache-dir -r requirements.txt
|
||||
|
||||
FROM runtime-base AS builder
|
||||
|
||||
RUN mamba install -y c-compiler cxx-compiler && \
|
||||
mamba clean -afy
|
||||
|
||||
COPY app ./app
|
||||
COPY scripts ./scripts
|
||||
|
||||
RUN python scripts/compile.py && \
|
||||
python scripts/compile.py --delete-source
|
||||
|
||||
FROM runtime-base AS runner
|
||||
|
||||
# 将代码放入子目录 'app',将数据放入子目录 'db_inp'
|
||||
# 这样临时文件默认会生成在 /app 下,而代码在 /app/app 下,实现了分离
|
||||
COPY app ./app
|
||||
COPY --from=builder /app/app ./app
|
||||
RUN python -c "from pathlib import Path; from zipfile import ZipFile; model_dir = Path('app/algorithms/health/model'); zip_path = model_dir / 'my_survival_forest_model_quxi.zip'; joblib_name = 'my_survival_forest_model_quxi.joblib'; joblib_path = model_dir / joblib_name; assert zip_path.exists(), f'Model archive not found: {zip_path}'; archive = ZipFile(zip_path); archive.extract(joblib_name, model_dir); archive.close(); assert joblib_path.exists(), f'Model file not extracted: {joblib_path}'" && \
|
||||
rm -f app/algorithms/health/model/my_survival_forest_model_quxi.zip
|
||||
# COPY db_inp ./db_inp
|
||||
COPY .env .
|
||||
RUN mkdir -p db_inp temp data inp
|
||||
|
||||
# 设置 PYTHONPATH 以便 uvicorn 找到 app 模块
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
# TJWaterServerCustomer 客户版后端
|
||||
|
||||
`TJWaterServerCustomer` 是 TJWater 客户交付版 Python 后端。该仓库应被视为可部署交付包,只保留客户运行、配置、部署、诊断和交付说明所需内容。
|
||||
|
||||
## 技术栈
|
||||
|
||||
- Python 3.12
|
||||
- FastAPI / Uvicorn
|
||||
- Pydantic / SQLAlchemy / psycopg
|
||||
- Redis、PostgreSQL、PostGIS、TimescaleDB
|
||||
- WNTR、EPANET、Cython、科学计算与空间分析依赖
|
||||
- pytest
|
||||
|
||||
## 目录结构
|
||||
|
||||
```text
|
||||
app/main.py FastAPI 入口
|
||||
app/api/ HTTP API 路由
|
||||
app/auth/ 认证和权限上下文
|
||||
app/core/ 配置、日志和基础设施初始化
|
||||
app/domain/ 领域模型和 Pydantic schema
|
||||
app/infra/ 数据库、缓存、EPANET 和外部集成
|
||||
app/services/ 核心业务服务,交付镜像中必须封装
|
||||
app/algorithms/ 核心算法,交付镜像中必须封装
|
||||
app/native/ 本地管网数据读写与转换,交付镜像中必须封装
|
||||
scripts/compile.py Cython 封装脚本
|
||||
infra/docker/ Docker Compose 编排
|
||||
tests/ 后端测试
|
||||
```
|
||||
|
||||
## 本地开发
|
||||
|
||||
推荐使用已有 conda 环境:
|
||||
|
||||
```bash
|
||||
conda run -n server python -m pytest tests -q
|
||||
conda run -n server python scripts/run_server.py
|
||||
```
|
||||
|
||||
本地调试不要在正常工作树执行源码删除命令。
|
||||
|
||||
## 客户版镜像打包
|
||||
|
||||
交付镜像标签通常为:
|
||||
|
||||
```bash
|
||||
docker build -t tjwater-server:latest .
|
||||
```
|
||||
|
||||
`Dockerfile` 会在 builder 阶段执行:
|
||||
|
||||
```bash
|
||||
python scripts/compile.py
|
||||
python scripts/compile.py --delete-source
|
||||
```
|
||||
|
||||
源码删除只发生在 Docker 构建层内,不会删除本地工作树源码。详细封装、验证和 Windows 导出说明见:
|
||||
|
||||
```text
|
||||
DELIVERY_PACKAGING_NOTES.md
|
||||
```
|
||||
|
||||
## 封装范围
|
||||
|
||||
`scripts/compile.py` 默认封装:
|
||||
|
||||
- `app/services`
|
||||
- `app/native/wndb`
|
||||
- `app/algorithms`
|
||||
- `app/infra/epanet/epanet.py`
|
||||
|
||||
交付镜像中这些核心目录不应残留未编译的 `.py` 源码,应以 `.so` 扩展模块运行。
|
||||
|
||||
## 验证命令
|
||||
|
||||
构建后检查镜像:
|
||||
|
||||
```bash
|
||||
docker image ls tjwater-server
|
||||
```
|
||||
|
||||
检查核心源码是否已删除、FastAPI 入口是否可导入:
|
||||
|
||||
```bash
|
||||
docker run --rm --entrypoint sh tjwater-server:latest -c "\
|
||||
printf 'core_py_count='; \
|
||||
find /app/app/services /app/app/native/wndb /app/app/algorithms /app/app/infra/epanet/epanet.py -name '*.py' 2>/dev/null | wc -l; \
|
||||
printf 'core_so_count='; \
|
||||
find /app/app/services /app/app/native/wndb /app/app/algorithms /app/app/infra/epanet -name '*.so' 2>/dev/null | wc -l; \
|
||||
python -c 'import app.main; print(\"import_app_main=ok\")'"
|
||||
```
|
||||
|
||||
期望结果:
|
||||
|
||||
```text
|
||||
core_py_count=0
|
||||
core_so_count=<非零>
|
||||
import_app_main=ok
|
||||
```
|
||||
|
||||
## 部署注意
|
||||
|
||||
客户交付时不要把本地 `../../app` 挂载到容器 `/app/app`,否则会覆盖镜像内已封装代码。交付 compose 文件应使用构建好的镜像,只挂载必要的运行数据和配置。
|
||||
|
||||
## 安全规则
|
||||
|
||||
不要提交 `.env`、生产凭据、客户数据、数据库 dump、日志、生成缓存、临时交付压缩包或本地运行目录。客户版仓库不应加入内部实验、调试工具或非交付源码材料。
|
||||
+37
-28
@@ -1,27 +1,21 @@
|
||||
from app.algorithms.cleaning import flow_data_clean, pressure_data_clean
|
||||
from app.algorithms.sensor import (
|
||||
pressure_sensor_placement_sensitivity,
|
||||
pressure_sensor_placement_kmeans,
|
||||
)
|
||||
from app.algorithms.isolation.valve import valve_isolation_analysis
|
||||
from app.algorithms.leakage import LeakageIdentifier
|
||||
from app.algorithms.health import PipelineHealthAnalyzer
|
||||
from app.algorithms.burst_location import run_burst_location
|
||||
from app.algorithms.simulation.scenarios import (
|
||||
convert_to_local_unit,
|
||||
burst_analysis,
|
||||
valve_close_analysis,
|
||||
flushing_analysis,
|
||||
contaminant_simulation,
|
||||
age_analysis,
|
||||
pressure_regulation,
|
||||
)
|
||||
"""Algorithm package with side-effect-free, lazy compatibility exports."""
|
||||
|
||||
__all__ = [
|
||||
"flow_data_clean",
|
||||
"pressure_data_clean",
|
||||
"pressure_sensor_placement_sensitivity",
|
||||
"pressure_sensor_placement_kmeans",
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
|
||||
_EXPORT_MODULES = {
|
||||
"flow_data_clean": "app.algorithms.cleaning",
|
||||
"pressure_data_clean": "app.algorithms.cleaning",
|
||||
"pressure_sensor_placement_sensitivity": "app.algorithms.sensor",
|
||||
"pressure_sensor_placement_kmeans": "app.algorithms.sensor",
|
||||
"valve_isolation_analysis": "app.algorithms.isolation.valve",
|
||||
"LeakageIdentifier": "app.algorithms.leakage",
|
||||
"PipelineHealthAnalyzer": "app.algorithms.health",
|
||||
"run_burst_location": "app.algorithms.burst_location",
|
||||
**{
|
||||
name: "app.algorithms.simulation.scenarios"
|
||||
for name in (
|
||||
"convert_to_local_unit",
|
||||
"burst_analysis",
|
||||
"valve_close_analysis",
|
||||
@@ -29,8 +23,23 @@ __all__ = [
|
||||
"contaminant_simulation",
|
||||
"age_analysis",
|
||||
"pressure_regulation",
|
||||
"valve_isolation_analysis",
|
||||
"LeakageIdentifier",
|
||||
"PipelineHealthAnalyzer",
|
||||
"run_burst_location",
|
||||
]
|
||||
)
|
||||
},
|
||||
}
|
||||
|
||||
__all__ = list(_EXPORT_MODULES)
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
try:
|
||||
module_name = _EXPORT_MODULES[name]
|
||||
except KeyError as exc:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
|
||||
|
||||
value = getattr(import_module(module_name), name)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted({*globals(), *__all__})
|
||||
|
||||
@@ -124,7 +124,9 @@ def _worker_evaluate(raw_ratios: np.ndarray) -> float:
|
||||
class LeakageIdentifier:
|
||||
FLOW_UNIT_TO_M3S = {
|
||||
"m3/s": 1.0,
|
||||
"m³/s": 1.0,
|
||||
"m3/h": 1.0 / 3600.0,
|
||||
"m³/h": 1.0 / 3600.0,
|
||||
"L/s": 1.0 / 1000.0,
|
||||
"L/min": 1.0 / 60000.0,
|
||||
}
|
||||
|
||||
@@ -72,6 +72,7 @@ def burst_analysis(
|
||||
modify_variable_pump_pattern: dict[str, list] = None,
|
||||
modify_valve_opening: dict[str, float] = None,
|
||||
scheme_name: str = None,
|
||||
username: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
爆管模拟
|
||||
@@ -86,6 +87,9 @@ def burst_analysis(
|
||||
:param scheme_name: 方案名称
|
||||
:return:
|
||||
"""
|
||||
if not username:
|
||||
raise ValueError("username is required when storing burst analysis scheme")
|
||||
|
||||
scheme_detail: dict = {
|
||||
"burst_ID": burst_ID,
|
||||
"burst_size": burst_size,
|
||||
@@ -211,7 +215,7 @@ def burst_analysis(
|
||||
name=name,
|
||||
scheme_name=scheme_name,
|
||||
scheme_type="burst_analysis",
|
||||
username="admin",
|
||||
username=username,
|
||||
scheme_start_time=modify_pattern_start_time,
|
||||
scheme_detail=scheme_detail,
|
||||
)
|
||||
@@ -311,6 +315,7 @@ def flushing_analysis(
|
||||
drainage_node_ID: str = None,
|
||||
flushing_flow: float = 0,
|
||||
scheme_name: str = None,
|
||||
username: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
管道冲洗模拟
|
||||
@@ -323,6 +328,9 @@ def flushing_analysis(
|
||||
:param scheme_name: 方案名称
|
||||
:return:
|
||||
"""
|
||||
if not username:
|
||||
raise ValueError("username is required when storing flushing analysis scheme")
|
||||
|
||||
scheme_detail: dict = {
|
||||
"duration": modify_total_duration,
|
||||
"valve_opening": modify_valve_opening,
|
||||
@@ -455,7 +463,7 @@ def flushing_analysis(
|
||||
name=name,
|
||||
scheme_name=scheme_name,
|
||||
scheme_type="flushing_analysis",
|
||||
username="admin",
|
||||
username=username,
|
||||
scheme_start_time=modify_pattern_start_time,
|
||||
scheme_detail=scheme_detail,
|
||||
)
|
||||
@@ -473,6 +481,7 @@ def contaminant_simulation(
|
||||
concentration: float, # 污染源浓度,单位mg/L
|
||||
scheme_name: str = None,
|
||||
source_pattern: str = None, # 污染源时间变化模式名称
|
||||
username: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
污染模拟
|
||||
@@ -486,6 +495,9 @@ def contaminant_simulation(
|
||||
:param scheme_name: 方案名称
|
||||
:return:
|
||||
"""
|
||||
if not username:
|
||||
raise ValueError("username is required when storing contaminant analysis scheme")
|
||||
|
||||
scheme_detail: dict = {
|
||||
"source": source,
|
||||
"concentration": concentration,
|
||||
@@ -608,7 +620,7 @@ def contaminant_simulation(
|
||||
name=name,
|
||||
scheme_name=scheme_name,
|
||||
scheme_type="contaminant_analysis",
|
||||
username="admin",
|
||||
username=username,
|
||||
scheme_start_time=modify_pattern_start_time,
|
||||
scheme_detail=scheme_detail,
|
||||
)
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Query
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
import app.native.wndb as wndb
|
||||
from app.infra.db.postgresql.scada import ScadaInfoRepository
|
||||
from app.infra.db.postgresql.scheme import SchemeRepository
|
||||
from app.auth.project_dependencies import get_project_pg_connection
|
||||
from app.services import project_info
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@@ -26,9 +25,7 @@ async def get_scada_info_with_connection(
|
||||
返回项目中所有的SCADA设备信息
|
||||
"""
|
||||
try:
|
||||
_ = conn
|
||||
network_name = project_info.name
|
||||
scada_data = wndb.get_all_scada_info(network_name) if network_name else []
|
||||
scada_data = await ScadaInfoRepository.get_scadas(conn)
|
||||
return {"success": True, "data": scada_data, "count": len(scada_data)}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
|
||||
@@ -4,8 +4,9 @@ import json
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
from fastapi import APIRouter, HTTPException, File, UploadFile, Query, Path, Body
|
||||
from fastapi import APIRouter, Depends, HTTPException, File, UploadFile, Query, Path, Body
|
||||
from fastapi.responses import PlainTextResponse
|
||||
from app.auth.keycloak_dependencies import get_current_keycloak_username
|
||||
import app.services.simulation as simulation
|
||||
import app.services.globals as globals
|
||||
from app.services.tjnetwork import (
|
||||
@@ -209,6 +210,7 @@ async def fastapi_burst_analysis(
|
||||
burst_size: list[float] = Query(..., description="对应各爆管点的爆管流量大小列表(L/s)"),
|
||||
modify_total_duration: int = Query(..., description="模拟总时长(秒)"),
|
||||
scheme_name: str = Query(..., description="分析方案名称"),
|
||||
username: str = Depends(get_current_keycloak_username),
|
||||
) -> str:
|
||||
"""
|
||||
爆管分析(高级版本)
|
||||
@@ -229,6 +231,7 @@ async def fastapi_burst_analysis(
|
||||
burst_size=burst_size,
|
||||
modify_total_duration=modify_total_duration,
|
||||
scheme_name=scheme_name,
|
||||
username=username,
|
||||
)
|
||||
return "success"
|
||||
|
||||
@@ -314,6 +317,7 @@ async def fastapi_flushing_analysis(
|
||||
flush_flow: float = Query(0, description="冲洗流量(L/s),0表示自动计算"),
|
||||
duration: int | None = Query(None, description="模拟持续时间(秒),默认900秒"),
|
||||
scheme_name: str = Query(..., description="冲洗方案名称"),
|
||||
username: str = Depends(get_current_keycloak_username),
|
||||
) -> str:
|
||||
"""
|
||||
冲洗分析(高级版本)
|
||||
@@ -340,6 +344,7 @@ async def fastapi_flushing_analysis(
|
||||
drainage_node_ID=drainage_node_ID,
|
||||
flushing_flow=flush_flow,
|
||||
scheme_name=scheme_name,
|
||||
username=username,
|
||||
)
|
||||
return result or "success"
|
||||
|
||||
@@ -354,6 +359,7 @@ async def fastapi_contaminant_simulation(
|
||||
duration: int = Query(..., description="模拟持续时间(秒)"),
|
||||
scheme_name: str = Query(..., description="模拟方案名称"),
|
||||
pattern: str | None = Query(None, description="污染源模式ID(可选)"),
|
||||
username: str = Depends(get_current_keycloak_username),
|
||||
) -> str:
|
||||
"""
|
||||
污染物模拟
|
||||
@@ -376,6 +382,7 @@ async def fastapi_contaminant_simulation(
|
||||
source=source,
|
||||
concentration=concentration,
|
||||
source_pattern=pattern,
|
||||
username=username,
|
||||
)
|
||||
return result or "success"
|
||||
|
||||
|
||||
@@ -7,7 +7,6 @@ from app.api.v1.endpoints import (
|
||||
scada,
|
||||
extension,
|
||||
snapshots,
|
||||
# data_query,
|
||||
users,
|
||||
schemes,
|
||||
misc,
|
||||
@@ -87,7 +86,6 @@ api_router.include_router(visuals.router, tags=["Visuals"])
|
||||
|
||||
# Simulation & Data
|
||||
api_router.include_router(simulation.router, tags=["Simulation Control"])
|
||||
# api_router.include_router(data_query.router, tags=["Data Query & InfluxDB"])
|
||||
api_router.include_router(scada.router)
|
||||
api_router.include_router(snapshots.router, tags=["Snapshots"])
|
||||
api_router.include_router(users.router, tags=["Users"])
|
||||
|
||||
@@ -27,11 +27,6 @@ class Settings(BaseSettings):
|
||||
TIMESCALEDB_DB_PORT: str = "5433"
|
||||
TIMESCALEDB_DB_USER: str = "postgres"
|
||||
TIMESCALEDB_DB_PASSWORD: str = "password"
|
||||
# InfluxDB
|
||||
INFLUXDB_URL: str = "http://localhost:8086"
|
||||
INFLUXDB_TOKEN: str = "token"
|
||||
INFLUXDB_ORG: str = "org"
|
||||
INFLUXDB_BUCKET: str = "bucket"
|
||||
|
||||
# Metadata Database Config (PostgreSQL)
|
||||
METADATA_DB_NAME: str = "system_hub"
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
from typing import Any
|
||||
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
|
||||
def _optional_text(value: Any) -> str | None:
|
||||
return str(value).strip() if value is not None else None
|
||||
|
||||
|
||||
def _optional_float(value: Any) -> float | None:
|
||||
return float(value) if value is not None else None
|
||||
|
||||
|
||||
class ScadaInfoRepository:
|
||||
"""Read SCADA metadata from the current project's business database."""
|
||||
|
||||
@staticmethod
|
||||
async def get_scadas(conn: AsyncConnection) -> list[dict[str, Any]]:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
SELECT id,
|
||||
type,
|
||||
associated_element_id,
|
||||
api_query_id,
|
||||
transmission_mode,
|
||||
transmission_frequency,
|
||||
reliability,
|
||||
x_coor,
|
||||
y_coor
|
||||
FROM public.scada_info
|
||||
"""
|
||||
)
|
||||
records = await cur.fetchall()
|
||||
|
||||
return [
|
||||
{
|
||||
"id": str(record["id"]).strip(),
|
||||
"type": str(record["type"]).strip().lower(),
|
||||
"associated_element_id": _optional_text(
|
||||
record["associated_element_id"]
|
||||
),
|
||||
"api_query_id": record["api_query_id"],
|
||||
"transmission_mode": record["transmission_mode"],
|
||||
"transmission_frequency": record["transmission_frequency"],
|
||||
"reliability": _optional_float(record["reliability"]),
|
||||
"x": _optional_float(record["x_coor"]),
|
||||
"y": _optional_float(record["y_coor"]),
|
||||
}
|
||||
for record in records
|
||||
]
|
||||
@@ -1,18 +1,19 @@
|
||||
import time
|
||||
from typing import List, Optional, Any, Dict, Tuple
|
||||
from datetime import datetime, timedelta
|
||||
from psycopg import AsyncConnection
|
||||
import pandas as pd
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
import app.native.wndb as wndb
|
||||
from app.algorithms.cleaning.flow import clean_flow_data_df_kf
|
||||
from app.algorithms.cleaning.pressure import clean_pressure_data_df_km
|
||||
from app.algorithms.health.analyzer import PipelineHealthAnalyzer
|
||||
import app.native.wndb as wndb
|
||||
|
||||
from app.infra.db.postgresql.scada import ScadaInfoRepository
|
||||
from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository
|
||||
from app.infra.db.timescaledb.repositories.scheme import SchemeRepository
|
||||
from app.infra.db.timescaledb.repositories.scada import ScadaRepository
|
||||
from app.services import project_info
|
||||
|
||||
|
||||
class CompositeQueries:
|
||||
@@ -20,6 +21,13 @@ class CompositeQueries:
|
||||
复合查询类,提供跨表查询功能
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
async def _get_project_scada_index(
|
||||
postgres_conn: AsyncConnection,
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
scadas = await ScadaInfoRepository.get_scadas(postgres_conn)
|
||||
return {scada["id"]: scada for scada in scadas}
|
||||
|
||||
@staticmethod
|
||||
async def get_scada_associated_realtime_simulation_data(
|
||||
timescale_conn: AsyncConnection,
|
||||
@@ -48,31 +56,22 @@ class CompositeQueries:
|
||||
ValueError: 当 SCADA 设备未找到或字段无效时
|
||||
"""
|
||||
result = {}
|
||||
# 1. 查询所有 SCADA 信息
|
||||
network_name = project_info.name
|
||||
scada_infos = wndb.get_all_scada_info(network_name) if network_name else []
|
||||
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
|
||||
|
||||
for device_id in device_ids:
|
||||
# 2. 根据 device_id 找到对应的 SCADA 信息
|
||||
target_scada = None
|
||||
for scada in scada_infos:
|
||||
if scada["id"] == device_id:
|
||||
target_scada = scada
|
||||
break
|
||||
|
||||
target_scada = scada_by_id.get(device_id)
|
||||
if not target_scada:
|
||||
raise ValueError(f"SCADA device {device_id} not found")
|
||||
|
||||
# 3. 根据 type 和 associated_element_id 查询对应的模拟数据
|
||||
element_id = target_scada["associated_element_id"]
|
||||
scada_type = target_scada["type"]
|
||||
|
||||
if scada_type.lower() == "pipe_flow":
|
||||
if scada_type == "pipe_flow":
|
||||
# 查询 link 模拟数据
|
||||
res = await RealtimeRepository.get_link_field_by_time_range(
|
||||
timescale_conn, start_time, end_time, element_id, "flow"
|
||||
)
|
||||
elif scada_type.lower() == "pressure":
|
||||
elif scada_type == "pressure":
|
||||
# 查询 node 模拟数据
|
||||
res = await RealtimeRepository.get_node_field_by_time_range(
|
||||
timescale_conn, start_time, end_time, element_id, "pressure"
|
||||
@@ -115,26 +114,17 @@ class CompositeQueries:
|
||||
ValueError: 当 SCADA 设备未找到或字段无效时
|
||||
"""
|
||||
result = {}
|
||||
# 1. 查询所有 SCADA 信息
|
||||
network_name = project_info.name
|
||||
scada_infos = wndb.get_all_scada_info(network_name) if network_name else []
|
||||
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
|
||||
|
||||
for device_id in device_ids:
|
||||
# 2. 根据 device_id 找到对应的 SCADA 信息
|
||||
target_scada = None
|
||||
for scada in scada_infos:
|
||||
if scada["id"] == device_id:
|
||||
target_scada = scada
|
||||
break
|
||||
|
||||
target_scada = scada_by_id.get(device_id)
|
||||
if not target_scada:
|
||||
raise ValueError(f"SCADA device {device_id} not found")
|
||||
|
||||
# 3. 根据 type 和 associated_element_id 查询对应的模拟数据
|
||||
element_id = target_scada["associated_element_id"]
|
||||
scada_type = target_scada["type"]
|
||||
|
||||
if scada_type.lower() == "pipe_flow":
|
||||
if scada_type == "pipe_flow":
|
||||
# 查询 link 模拟数据
|
||||
res = await SchemeRepository.get_link_field_by_scheme_and_time_range(
|
||||
timescale_conn,
|
||||
@@ -145,7 +135,7 @@ class CompositeQueries:
|
||||
element_id,
|
||||
"flow",
|
||||
)
|
||||
elif scada_type.lower() == "pressure":
|
||||
elif scada_type == "pressure":
|
||||
# 查询 node 模拟数据
|
||||
res = await SchemeRepository.get_node_field_by_scheme_and_time_range(
|
||||
timescale_conn,
|
||||
@@ -167,19 +157,19 @@ class CompositeQueries:
|
||||
@staticmethod
|
||||
async def get_realtime_simulation_data(
|
||||
timescale_conn: AsyncConnection,
|
||||
featureInfos: List[Tuple[str, str]],
|
||||
feature_infos: List[Tuple[str, str]],
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> Dict[str, List[Dict[str, Any]]]:
|
||||
"""
|
||||
获取 link/node 模拟值
|
||||
|
||||
根据传入的 featureInfos,找到关联的 link/node,
|
||||
根据传入的 feature_infos,找到关联的 link/node,
|
||||
并根据对应的 type,查询对应的模拟数据
|
||||
|
||||
Args:
|
||||
timescale_conn: TimescaleDB 异步连接
|
||||
featureInfos: 传入的 feature 信息列表,包含 (element_id, type)
|
||||
feature_infos: 传入的 feature 信息列表,包含 (element_id, type)
|
||||
start_time: 开始时间
|
||||
end_time: 结束时间
|
||||
|
||||
@@ -190,20 +180,20 @@ class CompositeQueries:
|
||||
ValueError: 当 SCADA 设备未找到或字段无效时
|
||||
"""
|
||||
result = {}
|
||||
for feature_id, type in featureInfos:
|
||||
for feature_id, feature_type in feature_infos:
|
||||
|
||||
if type.lower() == "pipe":
|
||||
if feature_type.lower() == "pipe":
|
||||
# 查询 link 模拟数据
|
||||
res = await RealtimeRepository.get_link_field_by_time_range(
|
||||
timescale_conn, start_time, end_time, feature_id, "flow"
|
||||
)
|
||||
elif type.lower() == "junction":
|
||||
elif feature_type.lower() == "junction":
|
||||
# 查询 node 模拟数据
|
||||
res = await RealtimeRepository.get_node_field_by_time_range(
|
||||
timescale_conn, start_time, end_time, feature_id, "pressure"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown type: {type}")
|
||||
raise ValueError(f"Unknown type: {feature_type}")
|
||||
# 添加 scada_id 到每个数据项
|
||||
for item in res:
|
||||
item["feature_id"] = feature_id
|
||||
@@ -213,7 +203,7 @@ class CompositeQueries:
|
||||
@staticmethod
|
||||
async def get_scheme_simulation_data(
|
||||
timescale_conn: AsyncConnection,
|
||||
featureInfos: List[Tuple[str, str]],
|
||||
feature_infos: List[Tuple[str, str]],
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
scheme_type: str,
|
||||
@@ -222,12 +212,12 @@ class CompositeQueries:
|
||||
"""
|
||||
获取 link/node scheme 模拟值
|
||||
|
||||
根据传入的 featureInfos,找到关联的 link/node,
|
||||
根据传入的 feature_infos,找到关联的 link/node,
|
||||
并根据对应的 type,查询对应的模拟数据
|
||||
|
||||
Args:
|
||||
timescale_conn: TimescaleDB 异步连接
|
||||
featureInfos: 传入的 feature 信息列表,包含 (element_id, type)
|
||||
feature_infos: 传入的 feature 信息列表,包含 (element_id, type)
|
||||
start_time: 开始时间
|
||||
end_time: 结束时间
|
||||
scheme_type: 工况类型
|
||||
@@ -240,8 +230,8 @@ class CompositeQueries:
|
||||
ValueError: 当类型无效时
|
||||
"""
|
||||
result = {}
|
||||
for feature_id, type in featureInfos:
|
||||
if type.lower() == "pipe":
|
||||
for feature_id, feature_type in feature_infos:
|
||||
if feature_type.lower() == "pipe":
|
||||
# 查询 link 模拟数据
|
||||
res = await SchemeRepository.get_link_field_by_scheme_and_time_range(
|
||||
timescale_conn,
|
||||
@@ -252,7 +242,7 @@ class CompositeQueries:
|
||||
feature_id,
|
||||
"flow",
|
||||
)
|
||||
elif type.lower() == "junction":
|
||||
elif feature_type.lower() == "junction":
|
||||
# 查询 node 模拟数据
|
||||
res = await SchemeRepository.get_node_field_by_scheme_and_time_range(
|
||||
timescale_conn,
|
||||
@@ -264,7 +254,7 @@ class CompositeQueries:
|
||||
"pressure",
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown type: {type}")
|
||||
raise ValueError(f"Unknown type: {feature_type}")
|
||||
# 添加 feature_id 到每个数据项
|
||||
for item in res:
|
||||
item["feature_id"] = feature_id
|
||||
@@ -301,33 +291,27 @@ class CompositeQueries:
|
||||
ValueError: 当元素类型无效时
|
||||
"""
|
||||
|
||||
# 1. 查询所有 SCADA 信息
|
||||
network_name = project_info.name
|
||||
scada_infos = wndb.get_all_scada_info(network_name) if network_name else []
|
||||
|
||||
# 2. 根据 element_type 和 element_id 找到关联的 SCADA 设备
|
||||
associated_scada = None
|
||||
for scada in scada_infos:
|
||||
if scada["associated_element_id"] == element_id:
|
||||
associated_scada = scada
|
||||
break
|
||||
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
|
||||
associated_scada = next(
|
||||
(
|
||||
scada
|
||||
for scada in scada_by_id.values()
|
||||
if scada["associated_element_id"] == element_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
if not associated_scada:
|
||||
# 没有找到关联的 SCADA 设备
|
||||
return None
|
||||
|
||||
# 3. 通过 SCADA device_id 获取监测数据
|
||||
device_id = associated_scada["id"]
|
||||
|
||||
# 根据 use_cleaned 参数选择字段
|
||||
data_field = "cleaned_value" if use_cleaned else "monitored_value"
|
||||
|
||||
# 保证 device_id 以列表形式传递
|
||||
res = await ScadaRepository.get_scada_field_by_id_time_range(
|
||||
timescale_conn, [device_id], start_time, end_time, data_field
|
||||
)
|
||||
|
||||
# 将 device_id 替换为 element_id 返回
|
||||
return {element_id: res.get(device_id, [])}
|
||||
|
||||
@staticmethod
|
||||
@@ -351,97 +335,111 @@ class CompositeQueries:
|
||||
end_time: 结束时间
|
||||
|
||||
Returns:
|
||||
"success" 或错误信息
|
||||
"success"
|
||||
|
||||
Raises:
|
||||
ValueError: 当前项目没有可清洗设备或指定时间范围内没有监测数据
|
||||
"""
|
||||
try:
|
||||
# 获取所有 SCADA 信息
|
||||
network_name = project_info.name
|
||||
scada_infos = wndb.get_all_scada_info(network_name) if network_name else []
|
||||
# 将列表转换为字典,以 device_id 为键
|
||||
scada_device_info_dict = {info["id"]: info for info in scada_infos}
|
||||
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
|
||||
supported_types = {"pressure", "pipe_flow", "flow"}
|
||||
|
||||
if device_ids:
|
||||
device_ids = [str(device_id).strip() for device_id in device_ids]
|
||||
missing_metadata_ids = [
|
||||
device_id
|
||||
for device_id in device_ids
|
||||
if device_id not in scada_by_id
|
||||
]
|
||||
if missing_metadata_ids:
|
||||
raise ValueError(
|
||||
f"当前项目中有 {len(missing_metadata_ids)} 个 SCADA 设备缺少元数据"
|
||||
)
|
||||
|
||||
unsupported_ids = [
|
||||
device_id
|
||||
for device_id in device_ids
|
||||
if scada_by_id[device_id]["type"] not in supported_types
|
||||
]
|
||||
if unsupported_ids:
|
||||
raise ValueError(
|
||||
f"当前项目中有 {len(unsupported_ids)} 个 SCADA 设备类型不支持清洗"
|
||||
)
|
||||
else:
|
||||
device_ids = [
|
||||
device_id
|
||||
for device_id, info in scada_by_id.items()
|
||||
if info["type"] in supported_types
|
||||
]
|
||||
|
||||
# 如果 device_ids 为空,则处理所有 SCADA 设备
|
||||
if not device_ids:
|
||||
device_ids = list(scada_device_info_dict.keys())
|
||||
raise ValueError("当前项目没有可清洗的 SCADA 设备")
|
||||
|
||||
# 批量查询所有设备的数据
|
||||
data = await ScadaRepository.get_scada_field_by_id_time_range(
|
||||
timescale_conn, device_ids, start_time, end_time, "monitored_value"
|
||||
)
|
||||
|
||||
if not data:
|
||||
return "error: fetch none scada data" # 没有数据,直接返回
|
||||
raise ValueError("指定时间范围内没有 SCADA 监测数据")
|
||||
|
||||
# 将嵌套字典转换为 DataFrame,使用 time 作为索引
|
||||
# data 格式: {device_id: [{"time": "...", "value": ...}, ...]}
|
||||
all_records = []
|
||||
for device_id, records in data.items():
|
||||
for record in records:
|
||||
all_records.append(
|
||||
normalized_data = {
|
||||
str(device_id): records for device_id, records in data.items()
|
||||
}
|
||||
missing_data_ids = [
|
||||
device_id for device_id in device_ids if not normalized_data.get(device_id)
|
||||
]
|
||||
if missing_data_ids:
|
||||
raise ValueError(
|
||||
f"指定时间范围内有 {len(missing_data_ids)} 个 SCADA 设备没有监测数据"
|
||||
)
|
||||
|
||||
all_records = [
|
||||
{
|
||||
"time": record["time"],
|
||||
"device_id": device_id,
|
||||
"value": record["value"],
|
||||
}
|
||||
)
|
||||
|
||||
for device_id, records in normalized_data.items()
|
||||
for record in records
|
||||
]
|
||||
if not all_records:
|
||||
return "error: fetch none scada data" # 没有数据,直接返回
|
||||
raise ValueError("指定时间范围内没有 SCADA 监测数据")
|
||||
|
||||
# 创建 DataFrame 并透视,使 device_id 成为列
|
||||
df_long = pd.DataFrame(all_records)
|
||||
df = df_long.pivot(index="time", columns="device_id", values="value")
|
||||
|
||||
# 根据type分类设备
|
||||
pressure_ids = [
|
||||
id
|
||||
for id in df.columns
|
||||
if scada_device_info_dict.get(id, {}).get("type") == "pressure"
|
||||
device_id
|
||||
for device_id in df.columns
|
||||
if scada_by_id[device_id]["type"] == "pressure"
|
||||
]
|
||||
flow_ids = [
|
||||
id
|
||||
for id in df.columns
|
||||
if scada_device_info_dict.get(id, {}).get("type") == "pipe_flow"
|
||||
device_id
|
||||
for device_id in df.columns
|
||||
if scada_by_id[device_id]["type"] in {"pipe_flow", "flow"}
|
||||
]
|
||||
|
||||
# 处理pressure数据
|
||||
if pressure_ids:
|
||||
pressure_df = df[pressure_ids]
|
||||
# 重置索引,将 time 变为普通列
|
||||
pressure_df = pressure_df.reset_index()
|
||||
# 调用清洗方法
|
||||
cleaned_df = clean_pressure_data_df_km(pressure_df)
|
||||
# 将清洗后的数据写回数据库
|
||||
for device_id in pressure_ids:
|
||||
if device_id in cleaned_df.columns:
|
||||
cleaned_values = cleaned_df[device_id].tolist()
|
||||
time_values = cleaned_df["time"].tolist()
|
||||
for i, time_str in enumerate(time_values):
|
||||
time_dt = datetime.fromisoformat(time_str)
|
||||
value = cleaned_values[i]
|
||||
await ScadaRepository.update_scada_field(
|
||||
timescale_conn,
|
||||
time_dt,
|
||||
device_id,
|
||||
"cleaned_value",
|
||||
value,
|
||||
)
|
||||
updated_rows = 0
|
||||
for grouped_ids, cleaning_function in (
|
||||
(pressure_ids, clean_pressure_data_df_km),
|
||||
(flow_ids, clean_flow_data_df_kf),
|
||||
):
|
||||
if not grouped_ids:
|
||||
continue
|
||||
|
||||
# 处理flow数据
|
||||
if flow_ids:
|
||||
flow_df = df[flow_ids]
|
||||
# 重置索引,将 time 变为普通列
|
||||
flow_df = flow_df.reset_index()
|
||||
# 调用清洗方法
|
||||
cleaned_df = clean_flow_data_df_kf(flow_df)
|
||||
# 将清洗后的数据写回数据库
|
||||
for device_id in flow_ids:
|
||||
if device_id in cleaned_df.columns:
|
||||
cleaned_values = cleaned_df[device_id].tolist()
|
||||
source_df = df[grouped_ids].reset_index()
|
||||
cleaned_df = cleaning_function(source_df)
|
||||
time_values = cleaned_df["time"].tolist()
|
||||
for i, time_str in enumerate(time_values):
|
||||
time_dt = datetime.fromisoformat(time_str)
|
||||
value = cleaned_values[i]
|
||||
|
||||
for device_id in grouped_ids:
|
||||
if device_id not in cleaned_df.columns:
|
||||
raise ValueError(f"设备 {device_id} 的清洗结果缺少数据列")
|
||||
|
||||
cleaned_values = cleaned_df[device_id].tolist()
|
||||
for time_value, value in zip(time_values, cleaned_values):
|
||||
time_dt = (
|
||||
time_value
|
||||
if isinstance(time_value, datetime)
|
||||
else datetime.fromisoformat(str(time_value))
|
||||
)
|
||||
await ScadaRepository.update_scada_field(
|
||||
timescale_conn,
|
||||
time_dt,
|
||||
@@ -449,10 +447,12 @@ class CompositeQueries:
|
||||
"cleaned_value",
|
||||
value,
|
||||
)
|
||||
updated_rows += 1
|
||||
|
||||
if updated_rows == 0:
|
||||
raise ValueError("SCADA 数据清洗未产生任何数据库更新")
|
||||
|
||||
return "success"
|
||||
except Exception as e:
|
||||
return f"error: {str(e)}"
|
||||
|
||||
@staticmethod
|
||||
async def predict_pipeline_health(
|
||||
|
||||
@@ -50,6 +50,7 @@ class InternalStorage:
|
||||
link_result_list: List[dict],
|
||||
result_start_time: str,
|
||||
num_periods: int = 1,
|
||||
result_timestep_seconds: int | None = None,
|
||||
db_name: str = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
@@ -70,6 +71,7 @@ class InternalStorage:
|
||||
link_result_list,
|
||||
result_start_time,
|
||||
num_periods,
|
||||
result_timestep_seconds,
|
||||
)
|
||||
break # 成功
|
||||
except Exception as e:
|
||||
|
||||
@@ -3,10 +3,24 @@ from datetime import datetime, timedelta
|
||||
from collections import defaultdict
|
||||
from psycopg import AsyncConnection, Connection, sql
|
||||
import app.services.globals as globals
|
||||
from app.services.time_api import parse_utc_time
|
||||
from app.services.time_api import parse_clock_duration_seconds, parse_utc_time
|
||||
|
||||
|
||||
class SchemeRepository:
|
||||
@staticmethod
|
||||
def _get_result_timestep(result_timestep_seconds: int | None) -> timedelta:
|
||||
if result_timestep_seconds is not None:
|
||||
if result_timestep_seconds <= 0:
|
||||
raise ValueError("result_timestep_seconds must be greater than 0.")
|
||||
return timedelta(seconds=result_timestep_seconds)
|
||||
|
||||
timestep_seconds = parse_clock_duration_seconds(
|
||||
globals.hydraulic_timestep,
|
||||
field_name="HYDRAULIC TIMESTEP",
|
||||
)
|
||||
if timestep_seconds <= 0:
|
||||
raise ValueError("HYDRAULIC TIMESTEP must be greater than 0.")
|
||||
return timedelta(seconds=timestep_seconds)
|
||||
|
||||
# --- Link Simulation ---
|
||||
|
||||
@@ -452,6 +466,7 @@ class SchemeRepository:
|
||||
link_result_list: List[Dict[str, any]],
|
||||
result_start_time: str,
|
||||
num_periods: int = 1,
|
||||
result_timestep_seconds: int | None = None,
|
||||
):
|
||||
"""
|
||||
Store scheme simulation results to TimescaleDB.
|
||||
@@ -468,20 +483,16 @@ class SchemeRepository:
|
||||
result_start_time, field_name="result_start_time"
|
||||
)
|
||||
|
||||
timestep_parts = globals.hydraulic_timestep.split(":")
|
||||
timestep = timedelta(
|
||||
hours=int(timestep_parts[0]),
|
||||
minutes=int(timestep_parts[1]),
|
||||
seconds=int(timestep_parts[2]),
|
||||
)
|
||||
timestep = SchemeRepository._get_result_timestep(result_timestep_seconds)
|
||||
|
||||
# Prepare node data for batch insert
|
||||
node_data = []
|
||||
for node_result in node_result_list:
|
||||
node_id = node_result.get("node")
|
||||
for period_index in range(num_periods):
|
||||
result_rows = node_result.get("result", [])
|
||||
for period_index in range(min(num_periods, len(result_rows))):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = node_result.get("result", [])[period_index]
|
||||
data = result_rows[period_index]
|
||||
node_data.append(
|
||||
{
|
||||
"time": current_time,
|
||||
@@ -499,9 +510,10 @@ class SchemeRepository:
|
||||
link_data = []
|
||||
for link_result in link_result_list:
|
||||
link_id = link_result.get("link")
|
||||
for period_index in range(num_periods):
|
||||
result_rows = link_result.get("result", [])
|
||||
for period_index in range(min(num_periods, len(result_rows))):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = link_result.get("result", [])[period_index]
|
||||
data = result_rows[period_index]
|
||||
link_data.append(
|
||||
{
|
||||
"time": current_time,
|
||||
@@ -535,6 +547,7 @@ class SchemeRepository:
|
||||
link_result_list: List[Dict[str, any]],
|
||||
result_start_time: str,
|
||||
num_periods: int = 1,
|
||||
result_timestep_seconds: int | None = None,
|
||||
):
|
||||
"""
|
||||
Store scheme simulation results to TimescaleDB (sync version).
|
||||
@@ -551,20 +564,16 @@ class SchemeRepository:
|
||||
result_start_time, field_name="result_start_time"
|
||||
)
|
||||
|
||||
timestep_parts = globals.hydraulic_timestep.split(":")
|
||||
timestep = timedelta(
|
||||
hours=int(timestep_parts[0]),
|
||||
minutes=int(timestep_parts[1]),
|
||||
seconds=int(timestep_parts[2]),
|
||||
)
|
||||
timestep = SchemeRepository._get_result_timestep(result_timestep_seconds)
|
||||
|
||||
# Prepare node data for batch insert
|
||||
node_data = []
|
||||
for node_result in node_result_list:
|
||||
node_id = node_result.get("node")
|
||||
for period_index in range(num_periods):
|
||||
result_rows = node_result.get("result", [])
|
||||
for period_index in range(min(num_periods, len(result_rows))):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = node_result.get("result", [])[period_index]
|
||||
data = result_rows[period_index]
|
||||
node_data.append(
|
||||
{
|
||||
"time": current_time,
|
||||
@@ -582,9 +591,10 @@ class SchemeRepository:
|
||||
link_data = []
|
||||
for link_result in link_result_list:
|
||||
link_id = link_result.get("link")
|
||||
for period_index in range(num_periods):
|
||||
result_rows = link_result.get("result", [])
|
||||
for period_index in range(min(num_periods, len(result_rows))):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = link_result.get("result", [])[period_index]
|
||||
data = result_rows[period_index]
|
||||
link_data.append(
|
||||
{
|
||||
"time": current_time,
|
||||
|
||||
@@ -1,3 +1,78 @@
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from threading import RLock
|
||||
|
||||
import psycopg as pg
|
||||
|
||||
from app.core.config import get_pgconn_string
|
||||
|
||||
g_conn_dict: dict[str, pg.Connection] = {}
|
||||
_registry_lock = RLock()
|
||||
_project_locks: dict[str, RLock] = {}
|
||||
|
||||
|
||||
def _is_closed(connection: pg.Connection) -> bool:
|
||||
return bool(getattr(connection, "closed", False))
|
||||
|
||||
|
||||
def _close_connection(connection: pg.Connection) -> None:
|
||||
if not _is_closed(connection):
|
||||
connection.close()
|
||||
|
||||
|
||||
def _is_healthy(connection: pg.Connection) -> bool:
|
||||
if _is_closed(connection):
|
||||
return False
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute("SELECT 1")
|
||||
except pg.Error:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _get_project_lock(name: str) -> RLock:
|
||||
with _registry_lock:
|
||||
lock = _project_locks.get(name)
|
||||
if lock is None:
|
||||
lock = RLock()
|
||||
_project_locks[name] = lock
|
||||
return lock
|
||||
|
||||
|
||||
def open_connection(name: str) -> pg.Connection:
|
||||
with _get_project_lock(name):
|
||||
connection = g_conn_dict.get(name)
|
||||
if connection is None or not _is_healthy(connection):
|
||||
if connection is not None:
|
||||
_close_connection(connection)
|
||||
connection = pg.connect(
|
||||
conninfo=get_pgconn_string(db_name=name), autocommit=True
|
||||
)
|
||||
g_conn_dict[name] = connection
|
||||
return connection
|
||||
|
||||
|
||||
def is_connection_open(name: str) -> bool:
|
||||
with _get_project_lock(name):
|
||||
connection = g_conn_dict.get(name)
|
||||
if connection is None:
|
||||
return False
|
||||
if not _is_healthy(connection):
|
||||
del g_conn_dict[name]
|
||||
_close_connection(connection)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def close_connection(name: str) -> None:
|
||||
with _get_project_lock(name):
|
||||
connection = g_conn_dict.pop(name, None)
|
||||
if connection is not None:
|
||||
_close_connection(connection)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def project_connection(name: str) -> Iterator[pg.Connection]:
|
||||
with _get_project_lock(name):
|
||||
yield open_connection(name)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from typing import Any
|
||||
from psycopg.rows import dict_row, Row
|
||||
from .connection import g_conn_dict as conn
|
||||
from .connection import project_connection
|
||||
|
||||
API_ADD = 'add'
|
||||
API_UPDATE = 'update'
|
||||
@@ -83,7 +83,8 @@ class DbChangeSet:
|
||||
|
||||
|
||||
def read(name: str, sql: str) -> Row:
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(sql)
|
||||
row = cur.fetchone()
|
||||
if row == None:
|
||||
@@ -92,19 +93,22 @@ def read(name: str, sql: str) -> Row:
|
||||
|
||||
|
||||
def read_all(name: str, sql: str) -> list[Row]:
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(sql)
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
def try_read(name: str, sql: str) -> Row | None:
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(sql)
|
||||
return cur.fetchone()
|
||||
|
||||
|
||||
def write(name: str, sql: str) -> None:
|
||||
with conn[name].cursor() as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(sql)
|
||||
|
||||
|
||||
|
||||
@@ -2,7 +2,11 @@ import os
|
||||
import psycopg as pg
|
||||
from psycopg import sql
|
||||
from psycopg.rows import dict_row
|
||||
from .connection import g_conn_dict as conn
|
||||
from .connection import (
|
||||
close_connection,
|
||||
is_connection_open,
|
||||
open_connection,
|
||||
)
|
||||
from app.core.config import get_pgconn_string, get_pg_config, get_pg_password
|
||||
|
||||
# no undo/redo
|
||||
@@ -31,9 +35,7 @@ def have_project(name: str) -> bool:
|
||||
|
||||
|
||||
def copy_project(source: str, new: str) -> None:
|
||||
if source in conn:
|
||||
conn[source].close()
|
||||
del conn[source]
|
||||
close_connection(source)
|
||||
|
||||
with pg.connect(
|
||||
conninfo=get_pgconn_string(db_name="postgres"), autocommit=True
|
||||
@@ -176,17 +178,12 @@ def clean_project(excluded: list[str] = []) -> None:
|
||||
|
||||
|
||||
def open_project(name: str) -> None:
|
||||
if name not in conn:
|
||||
conn[name] = pg.connect(
|
||||
conninfo=get_pgconn_string(db_name=name), autocommit=True
|
||||
)
|
||||
open_connection(name)
|
||||
|
||||
|
||||
def is_project_open(name: str) -> bool:
|
||||
return name in conn
|
||||
return is_connection_open(name)
|
||||
|
||||
|
||||
def close_project(name: str) -> None:
|
||||
if name in conn:
|
||||
conn[name].close()
|
||||
del conn[name]
|
||||
close_connection(name)
|
||||
|
||||
+19
-11
@@ -1,5 +1,5 @@
|
||||
from psycopg.rows import dict_row, Row
|
||||
from .connection import g_conn_dict as conn
|
||||
from .connection import project_connection
|
||||
from .database import read
|
||||
from typing import Any
|
||||
|
||||
@@ -47,7 +47,8 @@ ELEMENT_TYPES : dict[str, int] = {
|
||||
}
|
||||
|
||||
def _get_from(name: str, id: str, base_type: str) -> Row | None:
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select * from {base_type} where id = '{id}'")
|
||||
return cur.fetchone()
|
||||
|
||||
@@ -125,7 +126,8 @@ def is_region(name: str, id: str) -> bool:
|
||||
|
||||
def _get_all(name: str, base_type: str) -> list[str]:
|
||||
ids : list[str] = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id from {base_type} order by id")
|
||||
for record in cur:
|
||||
ids.append(record['id'])
|
||||
@@ -138,7 +140,8 @@ def get_nodes(name: str) -> list[str]:
|
||||
# DingZQ
|
||||
def _get_nodes_by_type(name: str, type: str) -> list[str]:
|
||||
ids : list[str] = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id from {_NODE} where type = '{type}' order by id")
|
||||
for record in cur:
|
||||
ids.append(record['id'])
|
||||
@@ -147,7 +150,8 @@ def _get_nodes_by_type(name: str, type: str) -> list[str]:
|
||||
# DingZQ
|
||||
def get_nodes_id_and_type(name: str) -> dict[str, str]:
|
||||
nodes_id_and_type: dict[str, str] = {}
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id, type from {_NODE} order by id")
|
||||
for record in cur:
|
||||
nodes_id_and_type[record['id']] = record['type']
|
||||
@@ -156,7 +160,8 @@ def get_nodes_id_and_type(name: str) -> dict[str, str]:
|
||||
# DingZQ 2024-12-31
|
||||
def get_major_nodes(name: str, diameter: int) -> list[str]:
|
||||
major_nodes_set = set()
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select node1, node2 from pipes where diameter > {diameter}")
|
||||
for record in cur:
|
||||
major_nodes_set.add(record['node1'])
|
||||
@@ -183,7 +188,8 @@ def get_links(name: str) -> list[str]:
|
||||
# DingZQ
|
||||
def _get_links_by_type(name: str, type: str) -> list[str]:
|
||||
ids : list[str] = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id from {_LINK} where type = '{type}' order by id")
|
||||
for record in cur:
|
||||
ids.append(record['id'])
|
||||
@@ -192,7 +198,8 @@ def _get_links_by_type(name: str, type: str) -> list[str]:
|
||||
# DingZQ
|
||||
def get_links_id_and_type(name: str) -> dict[str, str]:
|
||||
links_id_and_type: dict[str, str] = {}
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id, type from {_LINK} order by id")
|
||||
for record in cur:
|
||||
links_id_and_type[record['id']] = record['type']
|
||||
@@ -202,7 +209,8 @@ def get_links_id_and_type(name: str) -> dict[str, str]:
|
||||
# 获取直径大于800的管道
|
||||
def get_major_pipes(name: str, diameter: int) -> list[str]:
|
||||
major_pipe_ids: list[str] = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id from pipes where diameter > {diameter} order by id")
|
||||
for record in cur:
|
||||
major_pipe_ids.append(record['id'])
|
||||
@@ -232,7 +240,8 @@ def get_regions(name: str) -> list[str]:
|
||||
return _get_all(name, _REGION)
|
||||
|
||||
def get_node_links(name: str, id: str) -> list[str]:
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
links: list[str] = []
|
||||
for p in cur.execute(f"select id from pipes where node1 = '{id}' or node2 = '{id}'").fetchall():
|
||||
links.append(p['id'])
|
||||
@@ -259,4 +268,3 @@ def get_region_type(name: str, id: str)->str:
|
||||
return type
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from .database import *
|
||||
from .connection import project_connection
|
||||
from .s0_base import get_link_nodes
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
def sql_update_coord(node: str, x: float, y: float) -> str:
|
||||
coord = f"st_geomfromtext('point({x} {y})')"
|
||||
@@ -49,7 +51,8 @@ def get_links_in_extent(name: str, x1: float, y1: float, x2: float, y2: float) -
|
||||
node_ids = set([s.split(':')[0] for s in get_nodes_in_extent(name, x1, y1, x2, y2)])
|
||||
|
||||
all_link_ids = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select id from pipes")
|
||||
for record in cur:
|
||||
all_link_ids.append(record['id'])
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from .database import *
|
||||
from .connection import project_connection
|
||||
from .s0_base import *
|
||||
from psycopg.rows import dict_row
|
||||
import json
|
||||
|
||||
def get_pipe_risk_probability_now(name: str, pipe_id: str) -> dict[str, Any]:
|
||||
@@ -28,7 +30,8 @@ def get_pipe_risk_probability(name: str, pipe_id: str) -> dict[str, Any]:
|
||||
|
||||
def get_network_pipe_risk_probability_now(name: str) -> list[dict[str, Any]]:
|
||||
pipe_risk_probability_list = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select * from pipe_risk_probability")
|
||||
for record in cur:
|
||||
#pipe_risk_probability_list.append(record)
|
||||
@@ -42,7 +45,8 @@ def get_network_pipe_risk_probability_now(name: str) -> list[dict[str, Any]]:
|
||||
|
||||
def get_pipes_risk_probability(name: str, pipe_ids: list[str]) -> list[dict[str, Any]]:
|
||||
pipe_risk_probability_list = []
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select * from pipe_risk_probability")
|
||||
for record in cur:
|
||||
if record['pipeid'] in pipe_ids:
|
||||
@@ -67,7 +71,8 @@ def get_pipe_risk_probability_geometries(name: str) -> dict[str, Any]:
|
||||
# key_endnode = '下游节点'
|
||||
key_geometry = 'geometry'
|
||||
|
||||
with conn[name].cursor(row_factory=dict_row) as cur:
|
||||
with project_connection(name) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(f"select *, ST_AsGeoJSON(geometry) AS {key_geometry} from gis_pipe")
|
||||
|
||||
for record in cur:
|
||||
|
||||
@@ -11,6 +11,7 @@ from app.infra.db.timescaledb.internal_queries import InternalQueries
|
||||
from app.services.scheme_management import (
|
||||
query_burst_location_scheme_detail,
|
||||
query_burst_location_schemes,
|
||||
query_scheme_list,
|
||||
scheme_name_exists,
|
||||
store_scheme_info,
|
||||
)
|
||||
@@ -353,9 +354,15 @@ def run_burst_location_by_network(
|
||||
}
|
||||
)
|
||||
if normalized_data_source == "simulation":
|
||||
simulation_burst_ids = _get_simulation_scheme_burst_ids(
|
||||
network=network,
|
||||
scheme_name=simulation_scheme_name,
|
||||
scheme_type=resolved_simulation_scheme_type,
|
||||
)
|
||||
payload["simulation_scheme"] = {
|
||||
"name": simulation_scheme_name,
|
||||
"type": resolved_simulation_scheme_type,
|
||||
"burst_ids": simulation_burst_ids,
|
||||
}
|
||||
if scheme_name:
|
||||
_store_burst_scheme(
|
||||
@@ -464,6 +471,30 @@ def _validate_time_window(
|
||||
return start_dt, end_dt
|
||||
|
||||
|
||||
def _get_simulation_scheme_burst_ids(
|
||||
*, network: str, scheme_name: str | None, scheme_type: str
|
||||
) -> list[str]:
|
||||
if not scheme_name:
|
||||
return []
|
||||
rows = query_scheme_list(network) or []
|
||||
for row in rows:
|
||||
if len(row) < 7:
|
||||
continue
|
||||
if row[1] != scheme_name or row[2] != scheme_type:
|
||||
continue
|
||||
detail = row[6] if isinstance(row[6], dict) else {}
|
||||
return _normalize_burst_ids(detail.get("burst_ID"))
|
||||
return []
|
||||
|
||||
|
||||
def _normalize_burst_ids(value: Any) -> list[str]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, (list, tuple, set)):
|
||||
return _dedupe_ids([str(item) for item in value])
|
||||
return _dedupe_ids([str(value)])
|
||||
|
||||
|
||||
def _align_observed_series_pair(
|
||||
*,
|
||||
ids: list[str],
|
||||
|
||||
@@ -28,7 +28,6 @@ import pytz
|
||||
import requests
|
||||
import time
|
||||
from typing import Optional, Tuple
|
||||
import app.infra.db.influxdb.api as influxdb_api
|
||||
import typing
|
||||
import psycopg
|
||||
import logging
|
||||
@@ -1260,6 +1259,9 @@ def run_simulation(
|
||||
node_result, link_result, modify_pattern_start_time, db_name=db_name
|
||||
)
|
||||
elif simulation_type.upper() == "EXTENDED":
|
||||
result_timestep_seconds = times_info.get("report_step")
|
||||
if result_timestep_seconds is None:
|
||||
raise RuntimeError("run_project output missing times.report_step")
|
||||
TimescaleInternalStorage.store_scheme_simulation(
|
||||
scheme_type,
|
||||
scheme_name,
|
||||
@@ -1267,6 +1269,7 @@ def run_simulation(
|
||||
link_result,
|
||||
modify_pattern_start_time,
|
||||
num_periods_result,
|
||||
result_timestep_seconds,
|
||||
db_name=db_name,
|
||||
)
|
||||
endtime = time.time()
|
||||
|
||||
+2
-2
@@ -212,10 +212,10 @@ if __name__ == "__main__":
|
||||
if len(sys.argv) > 1:
|
||||
if sys.argv[1] == "--clean":
|
||||
clean_mode = True
|
||||
target_directories = sys.argv[2:]
|
||||
target_directories = sys.argv[2:] or DEFAULT_TARGETS
|
||||
elif sys.argv[1] == "--delete-source":
|
||||
delete_source_mode = True
|
||||
target_directories = sys.argv[2:]
|
||||
target_directories = sys.argv[2:] or DEFAULT_TARGETS
|
||||
else:
|
||||
target_directories = sys.argv[1:]
|
||||
else:
|
||||
|
||||
@@ -108,6 +108,12 @@ def _load_simulation_module(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def _build_authenticated_client(module) -> TestClient:
|
||||
app = build_test_app(module.router, "/api/v1")
|
||||
app.dependency_overrides[module.get_current_keycloak_username] = lambda: "alice"
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_run_project_endpoint_returns_plain_text(monkeypatch):
|
||||
module = _load_simulation_module(monkeypatch)
|
||||
monkeypatch.setattr(module, "run_project", lambda network: f"report::{network}")
|
||||
@@ -339,6 +345,33 @@ def test_valve_close_endpoint_passes_scheme_name(monkeypatch):
|
||||
}
|
||||
|
||||
|
||||
def test_burst_endpoint_passes_current_username(monkeypatch):
|
||||
module = _load_simulation_module(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
def fake_burst_analysis(**kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(module, "burst_analysis", fake_burst_analysis)
|
||||
client = _build_authenticated_client(module)
|
||||
|
||||
response = client.get(
|
||||
"/api/v1/burst_analysis/",
|
||||
params={
|
||||
"network": "demo",
|
||||
"modify_pattern_start_time": "2025-01-02T03:04:05+08:00",
|
||||
"burst_ID": ["P1"],
|
||||
"burst_size": [10.0],
|
||||
"modify_total_duration": 900,
|
||||
"scheme_name": "burst_case_01",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.text == '"success"'
|
||||
assert captured["username"] == "alice"
|
||||
|
||||
|
||||
def test_flushing_endpoint_passes_required_scheme_name(monkeypatch):
|
||||
module = _load_simulation_module(monkeypatch)
|
||||
captured = {}
|
||||
@@ -348,7 +381,7 @@ def test_flushing_endpoint_passes_required_scheme_name(monkeypatch):
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr(module, "flushing_analysis", fake_flushing_analysis)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
client = _build_authenticated_client(module)
|
||||
|
||||
response = client.get(
|
||||
"/api/v1/flushing_analysis/",
|
||||
@@ -374,12 +407,41 @@ def test_flushing_endpoint_passes_required_scheme_name(monkeypatch):
|
||||
"drainage_node_ID": "N1",
|
||||
"flushing_flow": 100.0,
|
||||
"scheme_name": "flush_case_01",
|
||||
"username": "alice",
|
||||
}
|
||||
|
||||
|
||||
def test_contaminant_endpoint_passes_current_username(monkeypatch):
|
||||
module = _load_simulation_module(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
def fake_contaminant_simulation(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr(module, "contaminant_simulation", fake_contaminant_simulation)
|
||||
client = _build_authenticated_client(module)
|
||||
|
||||
response = client.get(
|
||||
"/api/v1/contaminant_simulation/",
|
||||
params={
|
||||
"network": "demo",
|
||||
"start_time": "2025-01-02T03:04:05+08:00",
|
||||
"source": "N1",
|
||||
"concentration": 10.0,
|
||||
"duration": 900,
|
||||
"scheme_name": "contaminant_case_01",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.text == "ok"
|
||||
assert captured["username"] == "alice"
|
||||
|
||||
|
||||
def test_contaminant_endpoint_requires_scheme_name(monkeypatch):
|
||||
module = _load_simulation_module(monkeypatch)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
client = _build_authenticated_client(module)
|
||||
|
||||
response = client.get(
|
||||
"/api/v1/contaminant_simulation/",
|
||||
|
||||
@@ -75,6 +75,7 @@ def _load_burst_location_module():
|
||||
scheme_management_module = types.ModuleType("app.services.scheme_management")
|
||||
scheme_management_module.query_burst_location_scheme_detail = lambda *args, **kwargs: {}
|
||||
scheme_management_module.query_burst_location_schemes = lambda *args, **kwargs: []
|
||||
scheme_management_module.query_scheme_list = lambda *args, **kwargs: []
|
||||
scheme_management_module.scheme_name_exists = lambda *args, **kwargs: False
|
||||
scheme_management_module.store_scheme_info = lambda *args, **kwargs: None
|
||||
sys.modules["app.services.scheme_management"] = scheme_management_module
|
||||
@@ -177,6 +178,21 @@ def test_run_burst_location_uses_single_timerange_with_burst_source_split(monkey
|
||||
"query_realtime_simulation_by_ids_timerange",
|
||||
staticmethod(fake_realtime_query),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"query_scheme_list",
|
||||
lambda name: [
|
||||
(
|
||||
1,
|
||||
"BurstSchemeA",
|
||||
"burst_analysis",
|
||||
"testuser",
|
||||
None,
|
||||
None,
|
||||
{"burst_ID": ["Pipe-009", "Pipe-010"]},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
@@ -194,6 +210,7 @@ def test_run_burst_location_uses_single_timerange_with_burst_source_split(monkey
|
||||
assert result["simulation_scheme"] == {
|
||||
"name": "BurstSchemeA",
|
||||
"type": "burst_analysis",
|
||||
"burst_ids": ["Pipe-009", "Pipe-010"],
|
||||
}
|
||||
assert result["pressure_samples"] == {"burst": 4, "normal": 4}
|
||||
assert result["flow_samples"] == {"burst": 4, "normal": 4}
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
import pytest
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_leakage_identifier():
|
||||
algorithms_package = types.ModuleType("app.algorithms")
|
||||
algorithms_package.__path__ = []
|
||||
utils_module = types.ModuleType("app.algorithms._utils")
|
||||
utils_module._cleanup_temp_files = lambda *args, **kwargs: None
|
||||
sys.modules["app.algorithms"] = algorithms_package
|
||||
sys.modules["app.algorithms._utils"] = utils_module
|
||||
|
||||
module_path = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "app"
|
||||
/ "algorithms"
|
||||
/ "leakage"
|
||||
/ "identifier.py"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"tests_leakage_identifier_under_test", module_path
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec and spec.loader
|
||||
spec.loader.exec_module(module)
|
||||
return module.LeakageIdentifier
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("unit", "expected"),
|
||||
[
|
||||
("m3/s", 1.0),
|
||||
("m³/s", 1.0),
|
||||
("m3/h", 3600.0),
|
||||
("m³/h", 3600.0),
|
||||
],
|
||||
)
|
||||
def test_leakage_identifier_accepts_display_flow_units(unit, expected):
|
||||
LeakageIdentifier = _load_leakage_identifier()
|
||||
|
||||
assert LeakageIdentifier._flow_from_m3s(1.0, unit) == expected
|
||||
|
||||
|
||||
def test_leakage_identifier_accepts_display_flow_units_for_input():
|
||||
LeakageIdentifier = _load_leakage_identifier()
|
||||
|
||||
assert LeakageIdentifier._flow_to_m3s(3600.0, "m³/h") == 1.0
|
||||
@@ -0,0 +1,62 @@
|
||||
import asyncio
|
||||
|
||||
from app.infra.db.postgresql.scada import ScadaInfoRepository
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self):
|
||||
self.query = None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def execute(self, query):
|
||||
self.query = query
|
||||
|
||||
async def fetchall(self):
|
||||
return [
|
||||
{
|
||||
"id": " 25470001 ",
|
||||
"type": " PRESSURE ",
|
||||
"associated_element_id": " J1 ",
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
"reliability": "0.95",
|
||||
"x_coor": "117.1",
|
||||
"y_coor": "32.9",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self):
|
||||
self.cursor_instance = _FakeCursor()
|
||||
|
||||
def cursor(self):
|
||||
return self.cursor_instance
|
||||
|
||||
|
||||
def test_get_scadas_normalizes_id_and_type():
|
||||
conn = _FakeConnection()
|
||||
|
||||
result = asyncio.run(ScadaInfoRepository.get_scadas(conn))
|
||||
|
||||
assert result == [
|
||||
{
|
||||
"id": "25470001",
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
"reliability": 0.95,
|
||||
"x": 117.1,
|
||||
"y": 32.9,
|
||||
}
|
||||
]
|
||||
assert "associated_element_id" in conn.cursor_instance.query
|
||||
assert "FROM public.scada_info" in conn.cursor_instance.query
|
||||
@@ -0,0 +1,267 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from app.api.v1.endpoints import project_data
|
||||
from app.infra.db.timescaledb import composite_queries
|
||||
|
||||
|
||||
PROJECT_SCADA = {
|
||||
"id": "fengyang-pressure-1",
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
"reliability": 1.0,
|
||||
"x": 117.1,
|
||||
"y": 32.9,
|
||||
}
|
||||
START_TIME = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||||
END_TIME = datetime(2026, 6, 2, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _cleaning_args():
|
||||
return (
|
||||
object(),
|
||||
object(),
|
||||
["fengyang-pressure-1"],
|
||||
datetime(2026, 6, 1, tzinfo=timezone.utc),
|
||||
datetime(2026, 6, 2, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
def test_clean_scada_uses_current_project_metadata(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"get_scada_field_by_id_time_range",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"fengyang-pressure-1": [
|
||||
{"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
|
||||
]
|
||||
}
|
||||
),
|
||||
)
|
||||
update_mock = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"update_scada_field",
|
||||
update_mock,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries,
|
||||
"clean_pressure_data_df_km",
|
||||
lambda frame: pd.DataFrame(
|
||||
{
|
||||
"time": frame["time"],
|
||||
"fengyang-pressure-1": frame["fengyang-pressure-1"],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
update_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_clean_scada_rejects_devices_missing_from_project_metadata(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(return_value=[{"id": "other-device", "type": "pressure"}]),
|
||||
)
|
||||
query_mock = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"get_scada_field_by_id_time_range",
|
||||
query_mock,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="缺少元数据"):
|
||||
asyncio.run(
|
||||
composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
|
||||
)
|
||||
|
||||
query_mock.assert_not_awaited()
|
||||
|
||||
|
||||
def test_clean_scada_rejects_zero_database_updates(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"get_scada_field_by_id_time_range",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"fengyang-pressure-1": [
|
||||
{"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
|
||||
]
|
||||
}
|
||||
),
|
||||
)
|
||||
update_mock = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"update_scada_field",
|
||||
update_mock,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries,
|
||||
"clean_pressure_data_df_km",
|
||||
lambda _frame: pd.DataFrame(
|
||||
{"time": [], "fengyang-pressure-1": []}
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="未产生任何数据库更新"):
|
||||
asyncio.run(
|
||||
composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
|
||||
)
|
||||
|
||||
update_mock.assert_not_awaited()
|
||||
|
||||
|
||||
def test_clean_scada_propagates_write_failures(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"get_scada_field_by_id_time_range",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"fengyang-pressure-1": [
|
||||
{"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
|
||||
]
|
||||
}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"update_scada_field",
|
||||
AsyncMock(side_effect=RuntimeError("database write failed")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries,
|
||||
"clean_pressure_data_df_km",
|
||||
lambda frame: frame,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="database write failed"):
|
||||
asyncio.run(
|
||||
composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
|
||||
)
|
||||
|
||||
|
||||
def _patch_project_scada_metadata(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(return_value=[PROJECT_SCADA.copy()]),
|
||||
)
|
||||
|
||||
|
||||
def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
|
||||
_patch_project_scada_metadata(monkeypatch)
|
||||
query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
|
||||
monkeypatch.setattr(
|
||||
composite_queries.RealtimeRepository,
|
||||
"get_node_field_by_time_range",
|
||||
query_mock,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
composite_queries.CompositeQueries.get_scada_associated_realtime_simulation_data(
|
||||
object(), object(), [PROJECT_SCADA["id"]], START_TIME, END_TIME
|
||||
)
|
||||
)
|
||||
|
||||
assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
|
||||
assert query_mock.await_args.args[1:] == (
|
||||
START_TIME,
|
||||
END_TIME,
|
||||
"J1",
|
||||
"pressure",
|
||||
)
|
||||
|
||||
|
||||
def test_scheme_scada_simulation_uses_current_project_metadata(monkeypatch):
|
||||
_patch_project_scada_metadata(monkeypatch)
|
||||
query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
|
||||
monkeypatch.setattr(
|
||||
composite_queries.SchemeRepository,
|
||||
"get_node_field_by_scheme_and_time_range",
|
||||
query_mock,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
composite_queries.CompositeQueries.get_scada_associated_scheme_simulation_data(
|
||||
object(),
|
||||
object(),
|
||||
[PROJECT_SCADA["id"]],
|
||||
START_TIME,
|
||||
END_TIME,
|
||||
"baseline",
|
||||
"scheme-1",
|
||||
)
|
||||
)
|
||||
|
||||
assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
|
||||
assert query_mock.await_args.args[5:] == ("J1", "pressure")
|
||||
|
||||
|
||||
def test_element_scada_query_uses_current_project_metadata(monkeypatch):
|
||||
_patch_project_scada_metadata(monkeypatch)
|
||||
query_mock = AsyncMock(
|
||||
return_value={PROJECT_SCADA["id"]: [{"time": START_TIME, "value": 26.5}]}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaRepository,
|
||||
"get_scada_field_by_id_time_range",
|
||||
query_mock,
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
composite_queries.CompositeQueries.get_element_associated_scada_data(
|
||||
object(), object(), "J1", START_TIME, END_TIME
|
||||
)
|
||||
)
|
||||
|
||||
assert result == {"J1": [{"time": START_TIME, "value": 26.5}]}
|
||||
|
||||
|
||||
def test_scada_info_endpoint_uses_current_project_connection(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
project_data.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(return_value=[PROJECT_SCADA.copy()]),
|
||||
)
|
||||
|
||||
result = asyncio.run(project_data.get_scada_info_with_connection(object()))
|
||||
|
||||
assert result == {"success": True, "data": [PROJECT_SCADA], "count": 1}
|
||||
@@ -0,0 +1,123 @@
|
||||
import importlib.util
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from app.services.time_api import parse_utc_time
|
||||
|
||||
|
||||
def _load_scheme_repository():
|
||||
module_path = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "app"
|
||||
/ "infra"
|
||||
/ "db"
|
||||
/ "timescaledb"
|
||||
/ "repositories"
|
||||
/ "scheme.py"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"customer_scheme_repository_for_timestep_test", module_path
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec and spec.loader
|
||||
spec.loader.exec_module(module)
|
||||
return module.SchemeRepository
|
||||
|
||||
|
||||
SchemeRepository = _load_scheme_repository()
|
||||
|
||||
|
||||
def _node_result(periods: int) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"node": "J1",
|
||||
"result": [
|
||||
{"demand": index, "head": index, "pressure": index, "quality": index}
|
||||
for index in range(periods)
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _link_result(periods: int) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"link": "P1",
|
||||
"result": [
|
||||
{
|
||||
"flow": index,
|
||||
"friction": index,
|
||||
"headloss": index,
|
||||
"quality": index,
|
||||
"reaction": index,
|
||||
"setting": index,
|
||||
"status": index,
|
||||
"velocity": index,
|
||||
}
|
||||
for index in range(periods)
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_store_scheme_simulation_uses_15_minute_report_step(monkeypatch):
|
||||
inserted: dict[str, list[dict]] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_nodes_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_links_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("links", data)),
|
||||
)
|
||||
|
||||
SchemeRepository.store_scheme_simulation_result_sync(
|
||||
conn=object(),
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="five_hour_case",
|
||||
node_result_list=_node_result(21),
|
||||
link_result_list=_link_result(21),
|
||||
result_start_time="2026-07-16T00:00:00Z",
|
||||
num_periods=21,
|
||||
result_timestep_seconds=900,
|
||||
)
|
||||
|
||||
start_time = parse_utc_time("2026-07-16T00:00:00Z")
|
||||
assert len(inserted["nodes"]) == 21
|
||||
assert inserted["nodes"][0]["time"] == start_time
|
||||
assert inserted["nodes"][-1]["time"] == start_time + timedelta(hours=5)
|
||||
assert inserted["links"][-1]["time"] == start_time + timedelta(hours=5)
|
||||
|
||||
|
||||
def test_store_scheme_simulation_uses_hourly_report_step(monkeypatch):
|
||||
inserted: dict[str, list[dict]] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_nodes_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_links_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("links", data)),
|
||||
)
|
||||
|
||||
SchemeRepository.store_scheme_simulation_result_sync(
|
||||
conn=object(),
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="hourly_case",
|
||||
node_result_list=_node_result(6),
|
||||
link_result_list=_link_result(6),
|
||||
result_start_time="2026-07-16T00:00:00Z",
|
||||
num_periods=6,
|
||||
result_timestep_seconds=3600,
|
||||
)
|
||||
|
||||
start_time = parse_utc_time("2026-07-16T00:00:00Z")
|
||||
assert [item["time"] for item in inserted["nodes"]] == [
|
||||
start_time + timedelta(hours=index) for index in range(6)
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
import pytest
|
||||
|
||||
from app.native.wndb import connection
|
||||
from app.native.wndb import database
|
||||
from app.native.wndb import project
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self, connection):
|
||||
self.connection = connection
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def execute(self, sql):
|
||||
self.connection.executed.append(sql)
|
||||
if self.connection.fail_ping and sql == "SELECT 1":
|
||||
raise connection.pg.OperationalError("server closed the connection")
|
||||
|
||||
def fetchall(self):
|
||||
return self.connection.rows
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self, rows=None, *, closed=False, fail_ping=False):
|
||||
self.rows = list(rows or [])
|
||||
self.closed = closed
|
||||
self.fail_ping = fail_ping
|
||||
self.executed = []
|
||||
self.close_calls = 0
|
||||
|
||||
def cursor(self, row_factory=None):
|
||||
if self.closed:
|
||||
raise RuntimeError("the connection is closed")
|
||||
return _FakeCursor(self)
|
||||
|
||||
def close(self):
|
||||
self.close_calls += 1
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_native_connections():
|
||||
connection.g_conn_dict.clear()
|
||||
connection._project_locks.clear()
|
||||
yield
|
||||
connection.g_conn_dict.clear()
|
||||
connection._project_locks.clear()
|
||||
|
||||
|
||||
def test_is_project_open_drops_closed_cached_connection():
|
||||
connection.g_conn_dict["fengyang"] = _FakeConnection(closed=True)
|
||||
|
||||
assert project.is_project_open("fengyang") is False
|
||||
assert "fengyang" not in connection.g_conn_dict
|
||||
|
||||
|
||||
def test_open_connection_reuses_healthy_cached_connection(monkeypatch):
|
||||
cached = _FakeConnection()
|
||||
connection.g_conn_dict["fengyang"] = cached
|
||||
|
||||
def fail_connect(*, conninfo, autocommit):
|
||||
raise AssertionError("cached connection should be reused")
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fail_connect)
|
||||
|
||||
assert connection.open_connection("fengyang") is cached
|
||||
assert cached.executed == ["SELECT 1"]
|
||||
|
||||
|
||||
def test_read_all_reopens_closed_cached_connection(monkeypatch):
|
||||
stale = _FakeConnection(closed=True)
|
||||
fresh = _FakeConnection(rows=[{"key": "DURATION", "value": "01:00:00"}])
|
||||
connection.g_conn_dict["fengyang"] = stale
|
||||
|
||||
opened = []
|
||||
|
||||
def fake_connect(*, conninfo, autocommit):
|
||||
opened.append((conninfo, autocommit))
|
||||
return fresh
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fake_connect)
|
||||
monkeypatch.setattr(
|
||||
connection, "get_pgconn_string", lambda db_name: f"dbname={db_name}"
|
||||
)
|
||||
|
||||
rows = database.read_all("fengyang", "select * from times")
|
||||
|
||||
assert rows == [{"key": "DURATION", "value": "01:00:00"}]
|
||||
assert opened == [("dbname=fengyang", True)]
|
||||
assert connection.g_conn_dict["fengyang"] is fresh
|
||||
assert fresh.executed == ["select * from times"]
|
||||
|
||||
|
||||
def test_read_all_reopens_cached_connection_when_health_check_fails(monkeypatch):
|
||||
stale = _FakeConnection(fail_ping=True)
|
||||
fresh = _FakeConnection(rows=[{"scheme_name": "base"}])
|
||||
connection.g_conn_dict["fengyang"] = stale
|
||||
|
||||
opened = []
|
||||
|
||||
def fake_connect(*, conninfo, autocommit):
|
||||
opened.append((conninfo, autocommit))
|
||||
return fresh
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fake_connect)
|
||||
monkeypatch.setattr(
|
||||
connection, "get_pgconn_string", lambda db_name: f"dbname={db_name}"
|
||||
)
|
||||
|
||||
rows = database.read_all("fengyang", "select * from scheme_list")
|
||||
|
||||
assert rows == [{"scheme_name": "base"}]
|
||||
assert stale.executed == ["SELECT 1"]
|
||||
assert stale.close_calls == 1
|
||||
assert opened == [("dbname=fengyang", True)]
|
||||
assert connection.g_conn_dict["fengyang"] is fresh
|
||||
assert fresh.executed == ["select * from scheme_list"]
|
||||
Reference in New Issue
Block a user