OrgTJWater/TJWaterServerBinary
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(db)!: adopt project-routed pooled databases

Browse Source 
Reorganize WNDB by responsibility and remove legacy scheme endpoints.\n\nRoute analysis and time-series access through project pools, preserve transactional realtime replacement, and refresh GIS materialized views after writes.\n\nAdd database architecture documentation, live pooling coverage, API contract updates, and executable container verification.\n\nBREAKING CHANGE: legacy scheme APIs and flat app.native.wndb module imports are removed.
This commit is contained in:
jiang
2026-08-25 18:35:05 +08:00
parent fdbcc5c033
commit fa188af0b1
181 changed files with 8446 additions and 33546 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/infra/db/postgresql/analysis.py
+52
View File
@@ -0,0 +1,52 @@
from uuid import UUID
from psycopg import AsyncConnection
class AnalysisRepository:
@staticmethod
async def list_runs(conn: AsyncConnection) -> list[dict]:
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT run_id, name, run_type, created_by, created_at,
started_at, status, parameters
FROM analysis.runs
ORDER BY created_at DESC, run_id
"""
)
return await cur.fetchall()
@staticmethod
async def get_run(conn: AsyncConnection, run_id: UUID) -> dict | None:
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT run_id, name, run_type, created_by, created_at,
started_at, status, parameters
FROM analysis.runs
WHERE run_id = %s
""",
(run_id,),
)
return await cur.fetchone()
@staticmethod
async def list_results(
conn: AsyncConnection, run_id: UUID, result_type: str | None = None
) -> list[dict]:
query = """
SELECT result_id, run_id, result_type, node_id, link_id,
payload, created_at
FROM analysis.results
WHERE run_id = %s
"""
params: tuple[UUID] | tuple[UUID, str] = (run_id,)
if result_type is not None:
query += " AND result_type = %s"
params = (run_id, result_type)
query += " ORDER BY created_at, result_id"
async with conn.cursor() as cur:
await cur.execute(query, params)
return await cur.fetchall()
app/infra/db/postgresql/database.py
-123
View File
@@ -1,123 +0,0 @@
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Dict, Optional
import psycopg_pool
from psycopg.rows import dict_row
from app.infra.db.project_routing import get_project_pgconn_string
# Configure logging
logger = logging.getLogger(__name__)
class Database:
def __init__(self, db_name=None):
self.pool = None
self.db_name = db_name
self.conninfo = None
def init_pool(self, db_name=None):
"""Initialize the connection pool."""
# Use provided db_name, or the one from constructor, or default from config
target_db_name = db_name or self.db_name
# Get connection string, handling default case where target_db_name might be None
if target_db_name:
conn_string = get_project_pgconn_string(db_name=target_db_name)
else:
conn_string = get_project_pgconn_string()
self.conninfo = conn_string
try:
self.pool = psycopg_pool.AsyncConnectionPool(
conninfo=conn_string,
min_size=5,
max_size=20,
open=False, # Don't open immediately, wait for startup
kwargs={"row_factory": dict_row}, # Return rows as dictionaries
)
logger.info(f"PostgreSQL connection pool initialized for database: {target_db_name or 'default'}")
except Exception as e:
logger.error(f"Failed to initialize postgresql connection pool: {e}")
raise
async def open(self):
if self.pool:
await self.pool.open()
async def close(self):
"""Close the connection pool."""
if self.pool:
await self.pool.close()
logger.info("PostgreSQL connection pool closed.")
@asynccontextmanager
async def get_connection(self) -> AsyncGenerator:
"""Get a connection from the pool."""
if not self.pool:
raise Exception("Database pool is not initialized.")
async with self.pool.connection() as conn:
yield conn
# 默认数据库实例
db = Database()
# 缓存不同数据库的实例 - 避免重复创建连接池
_database_instances: Dict[str, Database] = {}
def create_database_instance(db_name):
"""Create a new Database instance for a specific database."""
return Database(db_name=db_name)
async def get_database_instance(db_name: Optional[str] = None) -> Database:
"""Get or create a database instance for the specified database name."""
if not db_name:
return db # 返回默认数据库实例
expected_conninfo = get_project_pgconn_string(db_name=db_name)
existing = _database_instances.get(db_name)
if existing is not None and existing.conninfo != expected_conninfo:
await existing.close()
del _database_instances[db_name]
if db_name not in _database_instances:
# 创建新的数据库实例
instance = create_database_instance(db_name)
instance.init_pool()
await instance.open()
_database_instances[db_name] = instance
logger.info(f"Created new database instance for: {db_name}")
return _database_instances[db_name]
async def get_db_connection():
"""Dependency for FastAPI to get a database connection."""
async with db.get_connection() as conn:
yield conn
async def get_database_connection(db_name: Optional[str] = None):
"""
FastAPI dependency to get database connection with optional database name.
使用方法: conn: AsyncConnection = Depends(lambda: get_database_connection("your_db_name"))
或在路由函数中: conn: AsyncConnection = Depends(get_database_connection)
"""
instance = await get_database_instance(db_name)
async with instance.get_connection() as conn:
yield conn
async def cleanup_database_instances():
"""Clean up all database instances (call this on application shutdown)."""
for db_name, instance in _database_instances.items():
await instance.close()
logger.info(f"Closed database instance for: {db_name}")
_database_instances.clear()
# 关闭默认数据库
await db.close()
logger.info("All database instances cleaned up.")
app/infra/db/postgresql/scada.py
+20 -15
View File
@@ -11,6 +11,10 @@ def _optional_float(value: Any) -> float | None:
return float(value) if value is not None else None
def _optional_int(value: Any) -> int | None:
return int(value) if value is not None else None
class ScadaInfoRepository:
"""Read SCADA metadata from the current project's business database."""
@@ -19,33 +23,34 @@ class ScadaInfoRepository:
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT id,
type,
associated_element_id,
SELECT id AS device_id,
device_type,
node_id,
link_id,
api_query_id,
transmission_mode,
transmission_frequency,
reliability,
x_coor,
y_coor
FROM public.scada_info
x,
y
FROM gis.scada_devices
ORDER BY id
"""
)
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"],
"device_id": str(record["device_id"]).strip(),
"device_type": str(record["device_type"]).strip().lower(),
"node_id": _optional_text(record["node_id"]),
"link_id": _optional_text(record["link_id"]),
"api_query_id": _optional_text(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"]),
"reliability": _optional_int(record["reliability"]),
"x": _optional_float(record["x"]),
"y": _optional_float(record["y"]),
}
for record in records
]
app/infra/db/postgresql/scada_assets.py
+66
View File
@@ -0,0 +1,66 @@
from typing import Any
from app.native.wndb.core.database import read_all, try_read
def get_scada_info_schema(name: str) -> dict[str, dict[str, Any]]:
return {
"device_id": {"type": "str", "optional": False, "readonly": True},
"device_type": {"type": "str", "optional": False, "readonly": True},
"node_id": {"type": "str", "optional": True, "readonly": True},
"link_id": {"type": "str", "optional": True, "readonly": True},
"api_query_id": {"type": "str", "optional": True, "readonly": True},
"transmission_mode": {"type": "str", "optional": False, "readonly": True},
"transmission_frequency": {"type": "str", "optional": False, "readonly": True},
"reliability": {"type": "int", "optional": False, "readonly": True},
"x": {"type": "float", "optional": True, "readonly": True},
"y": {"type": "float", "optional": True, "readonly": True},
}
_SELECT = """
SELECT device_id, device_type, node_id, link_id, api_query_id,
transmission_mode, transmission_frequency, reliability,
x, y
FROM asset.scada_devices
"""
_SELECT_MATERIALIZED = """
SELECT id AS device_id, device_type, node_id, link_id, api_query_id,
transmission_mode, transmission_frequency, reliability,
x, y
FROM gis.scada_devices
"""
def _device(row: dict[str, Any]) -> dict[str, Any]:
return {
"device_id": str(row["device_id"]),
"device_type": str(row["device_type"]),
"node_id": str(row["node_id"]) if row["node_id"] is not None else None,
"link_id": str(row["link_id"]) if row["link_id"] is not None else None,
"api_query_id": (
str(row["api_query_id"]) if row["api_query_id"] is not None else None
),
"transmission_mode": str(row["transmission_mode"]),
"transmission_frequency": str(row["transmission_frequency"]),
"reliability": int(row["reliability"]),
"x": float(row["x"]) if row["x"] is not None else None,
"y": float(row["y"]) if row["y"] is not None else None,
}
def get_scada_info(name: str, device_id: str) -> dict[str, Any]:
row = try_read(
name,
_SELECT + " WHERE device_id = %s",
(device_id,),
)
return _device(row) if row else {}
def get_all_scada_info(name: str) -> list[dict[str, Any]]:
return [
_device(row)
for row in read_all(name, _SELECT_MATERIALIZED + " ORDER BY device_id")
]
app/infra/db/postgresql/scheme.py
-104
View File
@@ -1,104 +0,0 @@
from typing import List, Optional, Any
from psycopg import AsyncConnection
class SchemeRepository:
@staticmethod
async def get_schemes(conn: AsyncConnection) -> List[dict]:
"""
查询pg数据库中, scheme_list 的所有记录
:param conn: 异步数据库连接
:return: 包含所有记录的列表, 每条记录为一个字典
"""
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT scheme_id, scheme_name, scheme_type, username, create_time, scheme_start_time, scheme_detail
FROM public.scheme_list
"""
)
records = await cur.fetchall()
scheme_list = []
for record in records:
scheme_list.append(
{
"scheme_id": record["scheme_id"],
"scheme_name": record["scheme_name"],
"scheme_type": record["scheme_type"],
"username": record["username"],
"create_time": record["create_time"],
"scheme_start_time": record["scheme_start_time"],
"scheme_detail": record["scheme_detail"],
}
)
return scheme_list
@staticmethod
async def get_burst_locate_results(conn: AsyncConnection) -> List[dict]:
"""
查询pg数据库中, burst_locate_result 的所有记录
:param conn: 异步数据库连接
:return: 包含所有记录的列表, 每条记录为一个字典
"""
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT id, type, burst_incident, leakage, detect_time, locate_result
FROM public.burst_locate_result
"""
)
records = await cur.fetchall()
results = []
for record in records:
results.append(
{
"id": record["id"],
"type": record["type"],
"burst_incident": record["burst_incident"],
"leakage": record["leakage"],
"detect_time": record["detect_time"],
"locate_result": record["locate_result"],
}
)
return results
@staticmethod
async def get_burst_locate_result_by_incident(
conn: AsyncConnection, burst_incident: str
) -> List[dict]:
"""
根据 burst_incident 查询爆管定位结果
:param conn: 异步数据库连接
:param burst_incident: 爆管事件标识
:return: 包含匹配记录的列表
"""
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT id, type, burst_incident, leakage, detect_time, locate_result
FROM public.burst_locate_result
WHERE burst_incident = %s
""",
(burst_incident,),
)
records = await cur.fetchall()
results = []
for record in records:
results.append(
{
"id": record["id"],
"type": record["type"],
"burst_incident": record["burst_incident"],
"leakage": record["leakage"],
"detect_time": record["detect_time"],
"locate_result": record["locate_result"],
}
)
return results
app/infra/db/postgresql/sensor_placement.py
+189
View File
@@ -0,0 +1,189 @@
from typing import Any
from uuid import UUID, uuid4
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb
from app.native.wndb.core.connection import project_connection
RUN_TYPE = "sensor_placement"
RESULT_TYPE = "sensor_placement"
def _placement_row(row: dict[str, Any]) -> dict[str, Any]:
payload = row.get("payload") if isinstance(row.get("payload"), dict) else {}
locations = [str(item) for item in payload.get("sensor_locations", [])]
return {
"run_id": row["run_id"],
"name": row["name"],
"sensor_count": len(locations),
"min_diameter": int(payload.get("minimum_diameter", 0)),
"created_by": row["created_by"],
"created_at": row["created_at"],
"status": row["status"],
"sensor_locations": locations,
}
def get_all_sensor_placements(name: str) -> list[dict[str, Any]]:
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"""
SELECT r.run_id, r.name, r.created_by, r.created_at, r.status,
result.payload
FROM analysis.runs AS r
JOIN LATERAL (
SELECT payload
FROM analysis.results
WHERE run_id = r.run_id AND result_type = %s
ORDER BY created_at DESC, result_id DESC
LIMIT 1
) AS result ON true
WHERE r.run_type = %s
ORDER BY r.created_at DESC, r.run_id
""",
(RESULT_TYPE, RUN_TYPE),
)
return [_placement_row(row) for row in cur.fetchall()]
def create_sensor_placement(
name: str,
*,
run_name: str,
min_diameter: int,
created_by: str,
sensor_locations: list[str],
) -> dict[str, Any]:
run_id = uuid4()
payload = {
"sensor_number": len(sensor_locations),
"minimum_diameter": min_diameter,
"sensor_locations": sensor_locations,
}
with project_connection(name) as conn, conn.transaction():
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"""
INSERT INTO analysis.runs
(run_id, name, run_type, created_by, started_at, status, parameters)
VALUES (%s, %s, %s, %s, now(), 'completed', '{}'::jsonb)
RETURNING run_id, name, created_by, created_at, status
""",
(run_id, run_name, RUN_TYPE, created_by),
)
created = cur.fetchone()
cur.execute(
"""
INSERT INTO analysis.results (run_id, result_type, payload)
VALUES (%s, %s, %s)
""",
(run_id, RESULT_TYPE, Jsonb(payload)),
)
if created is None:
raise RuntimeError("监测点优化运行写入失败")
return _placement_row(dict(created) | {"payload": payload})
def get_sensor_placement(name: str, run_id: UUID) -> dict[str, Any] | None:
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"""
SELECT r.run_id, r.name, r.created_by, r.created_at, r.status,
result.payload
FROM analysis.runs AS r
JOIN LATERAL (
SELECT payload
FROM analysis.results
WHERE run_id = r.run_id AND result_type = %s
ORDER BY created_at DESC, result_id DESC
LIMIT 1
) AS result ON true
WHERE r.run_id = %s AND r.run_type = %s
""",
(RESULT_TYPE, run_id, RUN_TYPE),
)
row = cur.fetchone()
return _placement_row(row) if row else None
def get_sensor_placement_nodes(
name: str,
node_ids: list[str],
) -> list[dict[str, Any]]:
if not node_ids:
return []
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"""
WITH incident_pipe_diameters AS (
SELECT node_id, MAX(diameter) AS max_pipe_diameter
FROM (
SELECT l.start_node_id AS node_id, p.diameter
FROM network.pipes AS p
JOIN network.links AS l ON l.id = p.link_id
WHERE l.start_node_id = ANY(%s)
UNION ALL
SELECT l.end_node_id AS node_id, p.diameter
FROM network.pipes AS p
JOIN network.links AS l ON l.id = p.link_id
WHERE l.end_node_id = ANY(%s)
) AS incident_pipes
GROUP BY node_id
)
SELECT j.node_id,
ipd.max_pipe_diameter,
j.elevation,
ST_X(g.geom) AS project_x,
ST_Y(g.geom) AS project_y,
ST_X(ST_Transform(g.geom, 3857)) AS map_x,
ST_Y(ST_Transform(g.geom, 3857)) AS map_y
FROM network.junctions AS j
JOIN gis.node_geometries AS g ON g.node_id = j.node_id
LEFT JOIN incident_pipe_diameters AS ipd ON ipd.node_id = j.node_id
WHERE j.node_id = ANY(%s)
ORDER BY j.node_id
""",
(node_ids, node_ids, node_ids),
)
return list(cur.fetchall())
def update_sensor_placement(
name: str,
run_id: UUID,
*,
expected_sensor_locations: list[str],
sensor_locations: list[str],
) -> dict[str, Any] | None:
with project_connection(name) as conn, conn.transaction():
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"""
SELECT result_id, payload
FROM analysis.results
WHERE run_id = %s AND result_type = %s
ORDER BY created_at DESC, result_id DESC
LIMIT 1
FOR UPDATE
""",
(run_id, RESULT_TYPE),
)
result = cur.fetchone()
if result is None:
return None
payload = result["payload"] if isinstance(result["payload"], dict) else {}
current = [str(item) for item in payload.get("sensor_locations", [])]
if current != expected_sensor_locations:
return None
payload = {
**payload,
"sensor_number": len(sensor_locations),
"sensor_locations": sensor_locations,
}
cur.execute(
"UPDATE analysis.results SET payload = %s WHERE result_id = %s",
(Jsonb(payload), result["result_id"]),
)
return get_sensor_placement(name, run_id)
app/infra/db/timescaledb/__init__.py
+1 -2
View File
@@ -1,2 +1 @@
from .database import *
from .composite_queries import CompositeQueries
from .composite_queries import CompositeQueries
app/infra/db/timescaledb/composite_queries.py
+38 -56
View File
@@ -5,15 +5,16 @@ from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
from psycopg import AsyncConnection
from uuid import UUID
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
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.analysis import AnalysisResultsRepository
from app.infra.db.timescaledb.repositories.scada import ScadaRepository
from app.native.wndb.model.pipes import get_pipes_by_property
class CompositeQueries:
@@ -26,7 +27,7 @@ class CompositeQueries:
postgres_conn: AsyncConnection,
) -> Dict[str, Dict[str, Any]]:
scadas = await ScadaInfoRepository.get_scadas(postgres_conn)
return {scada["id"]: scada for scada in scadas}
return {scada["device_id"]: scada for scada in scadas}
@staticmethod
async def get_scada_associated_realtime_simulation_data(
@@ -63,8 +64,12 @@ class CompositeQueries:
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
element_id = target_scada["associated_element_id"]
scada_type = target_scada["type"]
scada_type = target_scada["device_type"]
element_id = (
target_scada["link_id"]
if scada_type in {"pipe_flow", "flow"}
else target_scada["node_id"]
)
if scada_type == "pipe_flow":
# 查询 link 模拟数据
@@ -85,17 +90,16 @@ class CompositeQueries:
return result
@staticmethod
async def get_scada_associated_scheme_simulation_data(
async def get_scada_associated_analysis_simulation_data(
timescale_conn: AsyncConnection,
postgres_conn: AsyncConnection,
device_ids: List[str],
start_time: datetime,
end_time: datetime,
scheme_type: str,
scheme_name: str,
run_id: UUID,
) -> Dict[str, List[Dict[str, Any]]]:
"""
获取 SCADA 关联的 link/node scheme 模拟值
获取 SCADA 关联的 link/node 分析模拟值
根据传入的 SCADA device_ids，找到关联的 link/node，
并根据对应的 type，查询对应的模拟数据
@@ -121,30 +125,22 @@ class CompositeQueries:
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
element_id = target_scada["associated_element_id"]
scada_type = target_scada["type"]
scada_type = target_scada["device_type"]
element_id = (
target_scada["link_id"]
if scada_type in {"pipe_flow", "flow"}
else target_scada["node_id"]
)
if scada_type == "pipe_flow":
# 查询 link 模拟数据
res = await SchemeRepository.get_link_field_by_scheme_and_time_range(
timescale_conn,
scheme_type,
scheme_name,
start_time,
end_time,
element_id,
"flow",
res = await AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, element_id, start_time, end_time, "flow"
)
elif scada_type == "pressure":
# 查询 node 模拟数据
res = await SchemeRepository.get_node_field_by_scheme_and_time_range(
timescale_conn,
scheme_type,
scheme_name,
start_time,
end_time,
element_id,
"pressure",
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, element_id, start_time, end_time, "pressure"
)
else:
raise ValueError(f"Unknown SCADA type: {scada_type}")
@@ -201,16 +197,15 @@ class CompositeQueries:
return result
@staticmethod
async def get_scheme_simulation_data(
async def get_analysis_simulation_data(
timescale_conn: AsyncConnection,
feature_infos: List[Tuple[str, str]],
start_time: datetime,
end_time: datetime,
scheme_type: str,
scheme_name: str,
run_id: UUID,
) -> Dict[str, List[Dict[str, Any]]]:
"""
获取 link/node scheme 模拟值
获取 link/node 分析模拟值
根据传入的 feature_infos，找到关联的 link/node，
并根据对应的 type，查询对应的模拟数据
@@ -220,8 +215,7 @@ class CompositeQueries:
feature_infos: 传入的 feature 信息列表，包含 (element_id, type)
start_time: 开始时间
end_time: 结束时间
scheme_type: 工况类型
scheme_name: 工况名称
run_id: 分析运行 ID
Returns:
模拟数据字典，以 feature_id 为键，值为数据列表，每个数据包含 time, value 和 feature_id
@@ -233,25 +227,13 @@ class CompositeQueries:
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,
scheme_type,
scheme_name,
start_time,
end_time,
feature_id,
"flow",
res = await AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, feature_id, start_time, end_time, "flow"
)
elif feature_type.lower() == "junction":
# 查询 node 模拟数据
res = await SchemeRepository.get_node_field_by_scheme_and_time_range(
timescale_conn,
scheme_type,
scheme_name,
start_time,
end_time,
feature_id,
"pressure",
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, feature_id, start_time, end_time, "pressure"
)
else:
raise ValueError(f"Unknown type: {feature_type}")
@@ -296,7 +278,7 @@ class CompositeQueries:
(
scada
for scada in scada_by_id.values()
if scada["associated_element_id"] == element_id
if (scada.get("node_id") or scada.get("link_id")) == element_id
),
None,
)
@@ -304,7 +286,7 @@ class CompositeQueries:
if not associated_scada:
return None
device_id = associated_scada["id"]
device_id = associated_scada["device_id"]
data_field = "cleaned_value" if use_cleaned else "monitored_value"
@@ -358,7 +340,7 @@ class CompositeQueries:
unsupported_ids = [
device_id
for device_id in device_ids
if scada_by_id[device_id]["type"] not in supported_types
if scada_by_id[device_id]["device_type"] not in supported_types
]
if unsupported_ids:
raise ValueError(
@@ -368,7 +350,7 @@ class CompositeQueries:
device_ids = [
device_id
for device_id, info in scada_by_id.items()
if info["type"] in supported_types
if info["device_type"] in supported_types
]
if not device_ids:
@@ -409,12 +391,12 @@ class CompositeQueries:
pressure_ids = [
device_id
for device_id in df.columns
if scada_by_id[device_id]["type"] == "pressure"
if scada_by_id[device_id]["device_type"] == "pressure"
]
flow_ids = [
device_id
for device_id in df.columns
if scada_by_id[device_id]["type"] in {"pipe_flow", "flow"}
if scada_by_id[device_id]["device_type"] in {"pipe_flow", "flow"}
]
updated_rows = 0
@@ -497,7 +479,7 @@ class CompositeQueries:
# 批量查询这些管道的详细信息
fields = ["id", "diameter", "node1", "node2"]
all_links = wndb.get_pipes_by_property(network_name, fields=fields)
all_links = get_pipes_by_property(network_name, fields=fields)
# 转换为字典以快速查找
links_dict = {link["id"]: link for link in all_links}
app/infra/db/timescaledb/database.py
-132
View File
@@ -1,132 +0,0 @@
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Dict, Optional
import psycopg_pool
from psycopg.rows import dict_row
from app.infra.db.project_routing import get_project_timescale_pgconn_string
# Configure logging
logger = logging.getLogger(__name__)
class Database:
def __init__(self, db_name=None):
self.pool = None
self.db_name = db_name
self.conninfo = None
def init_pool(self, db_name=None):
"""Initialize the connection pool."""
# Use provided db_name, or the one from constructor, or default from config
target_db_name = db_name or self.db_name
# Get connection string, handling default case where target_db_name might be None
if target_db_name:
conn_string = get_project_timescale_pgconn_string(db_name=target_db_name)
else:
conn_string = get_project_timescale_pgconn_string()
self.conninfo = conn_string
try:
self.pool = psycopg_pool.AsyncConnectionPool(
conninfo=conn_string,
min_size=5,
max_size=20,
open=False, # Don't open immediately, wait for startup
kwargs={"row_factory": dict_row}, # Return rows as dictionaries
)
logger.info(
f"TimescaleDB connection pool initialized for database: {target_db_name or 'default'}"
)
except Exception as e:
logger.error(f"Failed to initialize TimescaleDB connection pool: {e}")
raise
async def open(self):
if self.pool:
await self.pool.open()
async def close(self):
"""Close the connection pool."""
if self.pool:
await self.pool.close()
logger.info("TimescaleDB connection pool closed.")
def get_pgconn_string(self, db_name=None):
"""Get the TimescaleDB connection string."""
target_db_name = db_name or self.db_name
if target_db_name:
return get_project_timescale_pgconn_string(db_name=target_db_name)
return get_project_timescale_pgconn_string()
@asynccontextmanager
async def get_connection(self) -> AsyncGenerator:
"""Get a connection from the pool."""
if not self.pool:
raise Exception("Database pool is not initialized.")
async with self.pool.connection() as conn:
yield conn
# 默认数据库实例
db = Database()
# 缓存不同数据库的实例 - 避免重复创建连接池
_database_instances: Dict[str, Database] = {}
def create_database_instance(db_name):
"""Create a new Database instance for a specific database."""
return Database(db_name=db_name)
async def get_database_instance(db_name: Optional[str] = None) -> Database:
"""Get or create a database instance for the specified database name."""
if not db_name:
return db # 返回默认数据库实例
expected_conninfo = get_project_timescale_pgconn_string(db_name=db_name)
existing = _database_instances.get(db_name)
if existing is not None and existing.conninfo != expected_conninfo:
await existing.close()
del _database_instances[db_name]
if db_name not in _database_instances:
# 创建新的数据库实例
instance = create_database_instance(db_name)
instance.init_pool()
await instance.open()
_database_instances[db_name] = instance
logger.info(f"Created new database instance for: {db_name}")
return _database_instances[db_name]
async def get_db_connection():
"""Dependency for FastAPI to get a database connection."""
async with db.get_connection() as conn:
yield conn
async def get_database_connection(db_name: Optional[str] = None):
"""
FastAPI dependency to get database connection with optional database name.
使用方法: conn: AsyncConnection = Depends(lambda: get_database_connection("your_db_name"))
或在路由函数中: conn: AsyncConnection = Depends(get_database_connection)
"""
instance = await get_database_instance(db_name)
async with instance.get_connection() as conn:
yield conn
async def cleanup_database_instances():
"""Clean up all database instances (call this on application shutdown)."""
for db_name, instance in _database_instances.items():
await instance.close()
logger.info(f"Closed database instance for: {db_name}")
_database_instances.clear()
# 关闭默认数据库
await db.close()
logger.info("All database instances cleaned up.")
app/infra/db/timescaledb/internal_queries.py
+41 -80
View File
@@ -2,12 +2,12 @@ from typing import List
from fastapi.logger import logger
from datetime import datetime, timedelta
import psycopg
from psycopg import sql
from psycopg.rows import dict_row
import time
from app.infra.db.project_routing import get_project_timescale_pgconn_string
from app.infra.db.timescaledb.repositories.scheme import SchemeRepository
from app.infra.db.timescaledb.sync_pool import timescale_connection
from app.infra.db.timescaledb.repositories.analysis import AnalysisResultsRepository
from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository
from app.infra.db.timescaledb.repositories.scada import ScadaRepository
from app.services.time_api import parse_utc_time
@@ -25,12 +25,7 @@ class InternalStorage:
"""存储实时模拟结果"""
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with timescale_connection(db_name) as conn:
RealtimeRepository.store_realtime_simulation_result_sync(
conn, node_result_list, link_result_list, result_start_time
)
@@ -43,9 +38,8 @@ class InternalStorage:
raise # 达到最大重试次数后抛出异常
@staticmethod
def store_scheme_simulation(
scheme_type: str,
scheme_name: str,
def store_analysis_simulation(
run_id,
node_result_list: List[dict],
link_result_list: List[dict],
result_start_time: str,
@@ -54,24 +48,16 @@ class InternalStorage:
db_name: str = None,
max_retries: int = 3,
):
"""存储方案模拟结果"""
"""Store immutable simulation results for one analysis run."""
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
SchemeRepository.store_scheme_simulation_result_sync(
conn,
scheme_type,
scheme_name,
node_result_list,
link_result_list,
result_start_time,
num_periods,
result_timestep_seconds,
with timescale_connection(db_name) as conn:
node_rows, link_rows = AnalysisResultsRepository.prepare_simulation_rows(
node_result_list, link_result_list, result_start_time,
num_periods, result_timestep_seconds or 3600,
)
AnalysisResultsRepository.store_results_sync(
conn, run_id, node_rows, link_rows
)
break # 成功
except Exception as e:
@@ -98,12 +84,7 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with timescale_connection(db_name) as conn:
rows = ScadaRepository.get_scada_by_ids_time_range_sync(
conn, device_ids, start_time, end_time
)
@@ -139,12 +120,7 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with timescale_connection(db_name) as conn:
rows = ScadaRepository.get_scada_by_ids_time_range_sync(
conn, device_ids, start_dt, end_dt
)
@@ -184,12 +160,7 @@ class InternalQueries:
)
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with timescale_connection(db_name) as conn:
return ScadaRepository.get_latest_scada_time_sync(
conn,
device_ids,
@@ -224,20 +195,19 @@ class InternalQueries:
)
@staticmethod
def query_scheme_simulation_by_ids_timerange(
def query_analysis_simulation_by_ids_timerange(
element_ids: List[str],
start_time: str | datetime,
end_time: str | datetime,
element_type: str,
field: str,
scheme_type: str,
scheme_name: str,
run_id,
db_name: str = None,
max_retries: int = 3,
) -> dict[str, list[dict]]:
"""查询方案模拟结果，返回 {id: [{time, value}, ...]}。"""
"""Query one analysis run, returning {id: [{time, value}, ...]}."""
return InternalQueries._query_simulation_by_ids_timerange(
schema_name="scheme",
schema_name="analysis",
element_ids=element_ids,
start_time=start_time,
end_time=end_time,
@@ -245,8 +215,7 @@ class InternalQueries:
field=field,
db_name=db_name,
max_retries=max_retries,
scheme_type=scheme_type,
scheme_name=scheme_name,
run_id=run_id,
)
@staticmethod
@@ -260,8 +229,7 @@ class InternalQueries:
field: str,
db_name: str = None,
max_retries: int = 3,
scheme_type: str | None = None,
scheme_name: str | None = None,
run_id=None,
) -> dict[str, list[dict]]:
normalized_element_ids = list(
dict.fromkeys(
@@ -275,51 +243,44 @@ class InternalQueries:
start_dt = parse_utc_time(start_time, field_name="start_time")
end_dt = parse_utc_time(end_time, field_name="end_time")
table_name, valid_fields = InternalQueries._resolve_simulation_table(element_type)
table_name, id_column, valid_fields = InternalQueries._resolve_simulation_table(element_type)
if field not in valid_fields:
raise ValueError(f"Invalid field for {element_type}: {field}")
if schema_name not in {"realtime", "scheme"}:
if schema_name not in {"realtime", "analysis"}:
raise ValueError(f"Unsupported schema_name: {schema_name}")
if schema_name == "scheme" and (not scheme_type or not scheme_name):
raise ValueError("scheme 查询必须提供 scheme_type 和 scheme_name。")
if schema_name == "analysis" and run_id is None:
raise ValueError("analysis query requires run_id")
for attempt in range(max_retries):
try:
conn_string = (
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with timescale_connection(db_name) as conn:
with conn.cursor(row_factory=dict_row) as cur:
if schema_name == "scheme":
if schema_name == "analysis":
query = sql.SQL(
"SELECT btrim(id::text) AS id, time, {} FROM {}.{} "
"WHERE scheme_type = %s AND scheme_name = %s "
"AND time >= %s AND time <= %s AND btrim(id::text) = ANY(%s)"
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE run_id = %s AND time >= %s AND time <= %s "
"AND btrim({}::text) = ANY(%s)"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
sql.Identifier(schema_name),
sql.Identifier(table_name),
sql.Identifier(id_column),
)
cur.execute(
query,
(
scheme_type,
scheme_name,
start_dt,
end_dt,
normalized_element_ids,
),
(run_id, start_dt, end_dt, normalized_element_ids),
)
else:
query = sql.SQL(
"SELECT btrim(id::text) AS id, time, {} FROM {}.{} "
"WHERE time >= %s AND time <= %s AND btrim(id::text) = ANY(%s)"
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE time >= %s AND time <= %s AND btrim({}::text) = ANY(%s)"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
sql.Identifier(schema_name),
sql.Identifier(table_name),
sql.Identifier(id_column),
)
cur.execute(query, (start_dt, end_dt, normalized_element_ids))
rows = cur.fetchall()
@@ -342,12 +303,12 @@ class InternalQueries:
raise
@staticmethod
def _resolve_simulation_table(element_type: str) -> tuple[str, set[str]]:
def _resolve_simulation_table(element_type: str) -> tuple[str, str, set[str]]:
normalized_type = element_type.lower()
if normalized_type == "node":
return "node_simulation", {"actual_demand", "total_head", "pressure", "quality"}
return "node_results", "node_id", {"actual_demand", "total_head", "pressure", "quality"}
if normalized_type == "link":
return "link_simulation", {
return "link_results", "link_id", {
"flow",
"friction",
"headloss",
app/infra/db/timescaledb/repositories/analysis.py
+268
View File
@@ -0,0 +1,268 @@
from datetime import datetime, timedelta
from typing import Any
from uuid import UUID
from psycopg import AsyncConnection, Connection, sql
from app.services.time_api import parse_utc_time
class AnalysisResultsRepository:
NODE_FIELDS = {"actual_demand", "total_head", "pressure", "quality"}
LINK_FIELDS = {
"flow",
"friction",
"headloss",
"quality",
"reaction",
"setting",
"status",
"velocity",
}
@staticmethod
def prepare_simulation_rows(
node_results: list[dict[str, Any]],
link_results: list[dict[str, Any]],
result_start_time: str,
num_periods: int,
result_timestep_seconds: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
start_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
timestep = timedelta(seconds=result_timestep_seconds)
node_rows: list[dict[str, Any]] = []
for node_result in node_results:
for period_index, values in enumerate(
node_result.get("result", [])[:num_periods]
):
node_rows.append(
{
"time": start_time + timestep * period_index,
"node_id": node_result["node"],
"actual_demand": values.get("demand"),
"total_head": values.get("head"),
"pressure": values.get("pressure"),
"quality": values.get("quality"),
}
)
link_rows: list[dict[str, Any]] = []
for link_result in link_results:
for period_index, values in enumerate(
link_result.get("result", [])[:num_periods]
):
link_rows.append(
{
"time": start_time + timestep * period_index,
"link_id": link_result["link"],
**{field: values.get(field) for field in AnalysisResultsRepository.LINK_FIELDS},
}
)
return node_rows, link_rows
@staticmethod
async def store_results(
conn: AsyncConnection,
run_id: UUID,
node_rows: list[dict[str, Any]],
link_rows: list[dict[str, Any]],
) -> None:
async with conn.transaction(), conn.cursor() as cur:
await AnalysisResultsRepository._lock_run(cur, run_id)
await AnalysisResultsRepository._assert_run_is_empty(cur, run_id)
if node_rows:
async with cur.copy(
"COPY analysis.node_results "
"(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
"FROM STDIN"
) as copy:
for row in node_rows:
await copy.write_row(
(
row["time"],
run_id,
row["node_id"],
row.get("actual_demand"),
row.get("total_head"),
row.get("pressure"),
row.get("quality"),
)
)
if link_rows:
async with cur.copy(
"COPY analysis.link_results "
"(time, run_id, link_id, flow, friction, headloss, quality, "
"reaction, setting, status, velocity) FROM STDIN"
) as copy:
for row in link_rows:
await copy.write_row(
(
row["time"],
run_id,
row["link_id"],
row.get("flow"),
row.get("friction"),
row.get("headloss"),
row.get("quality"),
row.get("reaction"),
row.get("setting"),
row.get("status"),
row.get("velocity"),
)
)
@staticmethod
async def _assert_run_is_empty(cur, run_id: UUID) -> None:
await cur.execute(
"""
SELECT EXISTS (
SELECT 1 FROM analysis.node_results WHERE run_id = %s
UNION ALL
SELECT 1 FROM analysis.link_results WHERE run_id = %s
) AS exists
""",
(run_id, run_id),
)
row = await cur.fetchone()
if row and row["exists"]:
raise ValueError(f"analysis results already exist for run {run_id}")
@staticmethod
async def _lock_run(cur, run_id: UUID) -> None:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
(run_id,),
)
@staticmethod
async def get_node_series(
conn: AsyncConnection,
run_id: UUID,
node_id: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> list[dict[str, Any]]:
if field not in AnalysisResultsRepository.NODE_FIELDS:
raise ValueError(f"invalid node result field: {field}")
query = sql.SQL(
"SELECT time, {} AS value FROM analysis.node_results "
"WHERE run_id = %s AND node_id = %s AND time BETWEEN %s AND %s "
"ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (run_id, node_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_link_series(
conn: AsyncConnection,
run_id: UUID,
link_id: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> list[dict[str, Any]]:
if field not in AnalysisResultsRepository.LINK_FIELDS:
raise ValueError(f"invalid link result field: {field}")
query = sql.SQL(
"SELECT time, {} AS value FROM analysis.link_results "
"WHERE run_id = %s AND link_id = %s AND time BETWEEN %s AND %s "
"ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (run_id, link_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_values_at_time(
conn: AsyncConnection,
run_id: UUID,
element_type: str,
result_time: datetime,
field: str,
) -> dict[str, Any]:
if element_type == "node":
table, id_column, fields = (
"node_results",
"node_id",
AnalysisResultsRepository.NODE_FIELDS,
)
elif element_type == "link":
table, id_column, fields = (
"link_results",
"link_id",
AnalysisResultsRepository.LINK_FIELDS,
)
else:
raise ValueError("element_type must be node or link")
if field not in fields:
raise ValueError(f"invalid {element_type} result field: {field}")
query = sql.SQL(
"SELECT {id_column}, {field} AS value FROM analysis.{table} "
"WHERE run_id = %s AND time = %s ORDER BY {id_column}"
).format(
id_column=sql.Identifier(id_column),
field=sql.Identifier(field),
table=sql.Identifier(table),
)
async with conn.cursor() as cur:
await cur.execute(query, (run_id, result_time))
return {row[id_column]: row["value"] for row in await cur.fetchall()}
@staticmethod
def store_results_sync(
conn: Connection,
run_id: UUID,
node_rows: list[dict[str, Any]],
link_rows: list[dict[str, Any]],
) -> None:
with conn.transaction(), conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
(run_id,),
)
cur.execute(
"""
SELECT EXISTS (
SELECT 1 FROM analysis.node_results WHERE run_id = %s
UNION ALL
SELECT 1 FROM analysis.link_results WHERE run_id = %s
) AS exists
""",
(run_id, run_id),
)
row = cur.fetchone()
if row and row["exists"]:
raise ValueError(f"analysis results already exist for run {run_id}")
if node_rows:
with cur.copy(
"COPY analysis.node_results "
"(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
"FROM STDIN"
) as copy:
for item in node_rows:
copy.write_row(
(
item["time"], run_id, item["node_id"],
item.get("actual_demand"), item.get("total_head"),
item.get("pressure"), item.get("quality"),
)
)
if link_rows:
with cur.copy(
"COPY analysis.link_results "
"(time, run_id, link_id, flow, friction, headloss, quality, "
"reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in link_rows:
copy.write_row(
(
item["time"], run_id, item["link_id"],
item.get("flow"), item.get("friction"),
item.get("headloss"), item.get("quality"),
item.get("reaction"), item.get("setting"),
item.get("status"), item.get("velocity"),
)
)
app/infra/db/timescaledb/repositories/realtime.py
+56 -34
View File
@@ -11,7 +11,7 @@ class RealtimeRepository:
@staticmethod
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
"""Batch insert for realtime.link_simulation using DELETE then COPY for performance."""
"""Batch insert for realtime.link_results using DELETE then COPY."""
if not data:
return
@@ -21,15 +21,19 @@ class RealtimeRepository:
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
await cur.execute(
"DELETE FROM realtime.link_simulation WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
"COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -49,7 +53,7 @@ class RealtimeRepository:
@staticmethod
def insert_links_batch_sync(conn: Connection, data: List[dict]):
"""Batch insert for realtime.link_simulation using DELETE then COPY for performance (sync version)."""
"""Synchronous batch insert for realtime.link_results."""
if not data:
return
@@ -59,15 +63,19 @@ class RealtimeRepository:
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
cur.execute(
"DELETE FROM realtime.link_simulation WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
"COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
@@ -91,7 +99,7 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s",
(start_time, end_time, link_id),
)
return await cur.fetchall()
@@ -104,7 +112,7 @@ class RealtimeRepository:
normalized_end_time = parse_utc_time(end_time, field_name="end_time")
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -132,7 +140,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s"
"SELECT time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -164,7 +172,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s"
"SELECT link_id, time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -172,7 +180,7 @@ class RealtimeRepository:
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
result[row["link_id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@@ -199,7 +207,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE realtime.link_simulation SET {} = %s WHERE time = %s AND id = %s"
"UPDATE realtime.link_results SET {} = %s WHERE time = %s AND link_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -211,7 +219,7 @@ class RealtimeRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
"DELETE FROM realtime.link_results WHERE time >= %s AND time <= %s",
(start_time, end_time),
)
@@ -228,15 +236,19 @@ class RealtimeRepository:
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
await cur.execute(
"DELETE FROM realtime.node_simulation WHERE time = %s",
"DELETE FROM realtime.node_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
"COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -261,15 +273,19 @@ class RealtimeRepository:
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
cur.execute(
"DELETE FROM realtime.node_simulation WHERE time = %s",
"DELETE FROM realtime.node_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
"COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
@@ -289,7 +305,7 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.node_simulation WHERE time >= %s AND time <= %s AND id = %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s",
(start_time, end_time, node_id),
)
return await cur.fetchall()
@@ -302,7 +318,7 @@ class RealtimeRepository:
normalized_end_time = parse_utc_time(end_time, field_name="end_time")
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -320,7 +336,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s AND id = %s"
"SELECT time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -339,7 +355,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s"
"SELECT node_id, time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -347,7 +363,7 @@ class RealtimeRepository:
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
result[row["node_id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@@ -365,7 +381,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE realtime.node_simulation SET {} = %s WHERE time = %s AND id = %s"
"UPDATE realtime.node_results SET {} = %s WHERE time = %s AND node_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -377,7 +393,7 @@ class RealtimeRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
"DELETE FROM realtime.node_results WHERE time >= %s AND time <= %s",
(start_time, end_time),
)
@@ -439,12 +455,15 @@ class RealtimeRepository:
}
)
# Insert data using batch methods
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
# Keep node and link replacement atomic. The batch helpers use nested
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
async with conn.transaction():
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
if link_data:
await RealtimeRepository.insert_links_batch(conn, link_data)
if link_data:
await RealtimeRepository.insert_links_batch(conn, link_data)
@staticmethod
def store_realtime_simulation_result_sync(
@@ -502,12 +521,15 @@ class RealtimeRepository:
}
)
# Insert data using batch methods
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
# Keep node and link replacement atomic. The batch helpers use nested
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
with conn.transaction():
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
if link_data:
RealtimeRepository.insert_links_batch_sync(conn, link_data)
if link_data:
RealtimeRepository.insert_links_batch_sync(conn, link_data)
@staticmethod
async def query_all_record_by_time_property(
app/infra/db/timescaledb/repositories/scada.py
+9 -9
View File
@@ -14,7 +14,7 @@ class ScadaRepository:
async with conn.cursor() as cur:
async with cur.copy(
"COPY scada.scada_data (time, device_id, monitored_value, cleaned_value) FROM STDIN"
"COPY scada.measurements (time, device_id, monitored_value, cleaned_value) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -35,7 +35,7 @@ class ScadaRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scada.scada_data WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
"SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
(device_ids, start_time, end_time),
)
return await cur.fetchall()
@@ -49,7 +49,7 @@ class ScadaRepository:
) -> List[dict]:
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"SELECT * FROM scada.scada_data WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
"SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
(device_ids, start_time, end_time),
)
return cur.fetchall()
@@ -63,12 +63,12 @@ class ScadaRepository:
with conn.cursor(row_factory=dict_row) as cur:
if before_time is None:
cur.execute(
"SELECT max(time) AS time FROM scada.scada_data WHERE device_id = ANY(%s)",
"SELECT max(time) AS time FROM scada.measurements WHERE device_id = ANY(%s)",
(device_ids,),
)
else:
cur.execute(
"SELECT max(time) AS time FROM scada.scada_data "
"SELECT max(time) AS time FROM scada.measurements "
"WHERE device_id = ANY(%s) AND time <= %s",
(device_ids, before_time),
)
@@ -88,7 +88,7 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT device_id, time, {} FROM scada.scada_data WHERE time >= %s AND time <= %s AND device_id = ANY(%s)"
"SELECT device_id, time, {} FROM scada.measurements WHERE time >= %s AND time <= %s AND device_id = ANY(%s)"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -111,10 +111,10 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
update_query = sql.SQL(
"UPDATE scada.scada_data SET {} = %s WHERE time = %s AND device_id = %s"
"UPDATE scada.measurements SET {} = %s WHERE time = %s AND device_id = %s"
).format(sql.Identifier(field))
insert_query = sql.SQL(
"INSERT INTO scada.scada_data (time, device_id, {}) VALUES (%s, %s, %s)"
"INSERT INTO scada.measurements (time, device_id, {}) VALUES (%s, %s, %s)"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -128,6 +128,6 @@ class ScadaRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scada.scada_data WHERE device_id = %s AND time >= %s AND time <= %s",
"DELETE FROM scada.measurements WHERE device_id = %s AND time >= %s AND time <= %s",
(device_id, start_time, end_time),
)
app/infra/db/timescaledb/repositories/scheme.py
-710
View File
@@ -1,710 +0,0 @@
from typing import List, Any, Dict
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_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 ---
@staticmethod
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
"""Batch insert for scheme.link_simulation using DELETE then COPY for performance."""
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
await cur.execute(
"DELETE FROM scheme.link_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY scheme.link_simulation (time, scheme_type, scheme_name, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("flow"),
item.get("friction"),
item.get("headloss"),
item.get("quality"),
item.get("reaction"),
item.get("setting"),
item.get("status"),
item.get("velocity"),
)
)
@staticmethod
def insert_links_batch_sync(conn: Connection, data: List[dict]):
"""Batch insert for scheme.link_simulation using DELETE then COPY for performance (sync version)."""
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
cur.execute(
"DELETE FROM scheme.link_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY scheme.link_simulation (time, scheme_type, scheme_name, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("flow"),
item.get("friction"),
item.get("headloss"),
item.get("quality"),
item.get("reaction"),
item.get("setting"),
item.get("status"),
item.get("velocity"),
)
)
@staticmethod
async def get_link_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
link_id: str,
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s",
(scheme_type, scheme_name, start_time, end_time, link_id),
)
return await cur.fetchall()
@staticmethod
async def get_links_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
return await cur.fetchall()
@staticmethod
async def get_link_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
link_id: str,
field: str,
) -> List[Dict[str, Any]]:
# Validate field name to prevent SQL injection
valid_fields = {
"flow",
"friction",
"headloss",
"quality",
"reaction",
"setting",
"status",
"velocity",
}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(
query, (scheme_type, scheme_name, start_time, end_time, link_id)
)
rows = await cur.fetchall()
return [
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_links_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> dict:
# Validate field name to prevent SQL injection
valid_fields = {
"flow",
"friction",
"headloss",
"quality",
"reaction",
"setting",
"status",
"velocity",
}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (scheme_type, scheme_name, start_time, end_time))
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@staticmethod
async def update_link_field(
conn: AsyncConnection,
time: datetime,
scheme_type: str,
scheme_name: str,
link_id: str,
field: str,
value: Any,
):
valid_fields = {
"flow",
"friction",
"headloss",
"quality",
"reaction",
"setting",
"status",
"velocity",
}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE scheme.link_simulation SET {} = %s WHERE time = %s AND scheme_type = %s AND scheme_name = %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (value, time, scheme_type, scheme_name, link_id))
@staticmethod
async def delete_links_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
# --- Node Simulation ---
@staticmethod
async def insert_nodes_batch(conn: AsyncConnection, data: List[dict]):
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
await cur.execute(
"DELETE FROM scheme.node_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY scheme.node_simulation (time, scheme_type, scheme_name, id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("actual_demand"),
item.get("total_head"),
item.get("pressure"),
item.get("quality"),
)
)
@staticmethod
def insert_nodes_batch_sync(conn: Connection, data: List[dict]):
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
cur.execute(
"DELETE FROM scheme.node_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY scheme.node_simulation (time, scheme_type, scheme_name, id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("actual_demand"),
item.get("total_head"),
item.get("pressure"),
item.get("quality"),
)
)
@staticmethod
async def get_node_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
node_id: str,
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s",
(scheme_type, scheme_name, start_time, end_time, node_id),
)
return await cur.fetchall()
@staticmethod
async def get_nodes_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
return await cur.fetchall()
@staticmethod
async def get_node_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
node_id: str,
field: str,
) -> List[Dict[str, Any]]:
# Validate field name to prevent SQL injection
valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(
query, (scheme_type, scheme_name, start_time, end_time, node_id)
)
rows = await cur.fetchall()
return [
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_nodes_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> dict:
# Validate field name to prevent SQL injection
valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (scheme_type, scheme_name, start_time, end_time))
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@staticmethod
async def update_node_field(
conn: AsyncConnection,
time: datetime,
scheme_type: str,
scheme_name: str,
node_id: str,
field: str,
value: Any,
):
valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE scheme.node_simulation SET {} = %s WHERE time = %s AND scheme_type = %s AND scheme_name = %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (value, time, scheme_type, scheme_name, node_id))
@staticmethod
async def delete_nodes_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
# --- 复合查询 ---
@staticmethod
async def store_scheme_simulation_result(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
node_result_list: List[Dict[str, any]],
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.
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: Scheme name
node_result_list: List of node simulation results
link_result_list: List of link simulation results
result_start_time: Start time for the results (ISO format string)
"""
simulation_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
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")
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 = result_rows[period_index]
node_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": node_id,
"actual_demand": data.get("demand"),
"total_head": data.get("head"),
"pressure": data.get("pressure"),
"quality": data.get("quality"),
}
)
# Prepare link data for batch insert
link_data = []
for link_result in link_result_list:
link_id = link_result.get("link")
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 = result_rows[period_index]
link_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": link_id,
"flow": data.get("flow"),
"friction": data.get("friction"),
"headloss": data.get("headloss"),
"quality": data.get("quality"),
"reaction": data.get("reaction"),
"setting": data.get("setting"),
"status": data.get("status"),
"velocity": data.get("velocity"),
}
)
# Insert data using batch methods
if node_data:
await SchemeRepository.insert_nodes_batch(conn, node_data)
if link_data:
await SchemeRepository.insert_links_batch(conn, link_data)
@staticmethod
def store_scheme_simulation_result_sync(
conn: Connection,
scheme_type: str,
scheme_name: str,
node_result_list: List[Dict[str, any]],
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).
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: Scheme name
node_result_list: List of node simulation results
link_result_list: List of link simulation results
result_start_time: Start time for the results (ISO format string)
"""
simulation_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
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")
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 = result_rows[period_index]
node_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": node_id,
"actual_demand": data.get("demand"),
"total_head": data.get("head"),
"pressure": data.get("pressure"),
"quality": data.get("quality"),
}
)
# Prepare link data for batch insert
link_data = []
for link_result in link_result_list:
link_id = link_result.get("link")
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 = result_rows[period_index]
link_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": link_id,
"flow": data.get("flow"),
"friction": data.get("friction"),
"headloss": data.get("headloss"),
"quality": data.get("quality"),
"reaction": data.get("reaction"),
"setting": data.get("setting"),
"status": data.get("status"),
"velocity": data.get("velocity"),
}
)
# Insert data using batch methods
if node_data:
SchemeRepository.insert_nodes_batch_sync(conn, node_data)
if link_data:
SchemeRepository.insert_links_batch_sync(conn, link_data)
@staticmethod
async def query_all_record_by_scheme_time_property(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
query_time: str,
type: str,
property: str,
) -> list:
"""
Query all records by scheme, time and property from TimescaleDB.
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: Scheme name
query_time: Time to query (ISO format string)
type: Type of data ("node" or "link")
property: Property/field to query
Returns:
List of records matching the criteria
"""
target_time = parse_utc_time(query_time, field_name="query_time")
# Create time range: query_time ± 1 second
start_time = target_time - timedelta(seconds=1)
end_time = target_time + timedelta(seconds=1)
# Query based on type
if type.lower() == "node":
data = await SchemeRepository.get_nodes_field_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, property
)
elif type.lower() == "link":
data = await SchemeRepository.get_links_field_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, property
)
else:
raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")
# Format the results
# Format the results
result = []
for id, items in data.items():
for item in items:
result.append({"ID": id, "value": item["value"]})
return result
@staticmethod
async def query_scheme_simulation_result_by_id_time(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
id: str,
type: str,
query_time: str,
) -> list[dict]:
"""
Query scheme simulation results by id and time from TimescaleDB.
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: Scheme name
id: The id of the node or link
type: Type of data ("node" or "link")
query_time: Time to query (ISO format string)
Returns:
List of records matching the criteria
"""
target_time = parse_utc_time(query_time, field_name="query_time")
# Create time range: query_time ± 1 second
start_time = target_time - timedelta(seconds=1)
end_time = target_time + timedelta(seconds=1)
# Query based on type
if type.lower() == "node":
return await SchemeRepository.get_node_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, id
)
elif type.lower() == "link":
return await SchemeRepository.get_link_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, id
)
else:
raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")
app/infra/db/timescaledb/sync_pool.py
+95
View File
@@ -0,0 +1,95 @@
from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager
from threading import RLock
from psycopg import Connection
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool
from app.core.config import settings
from app.infra.db.project_routing import get_project_timescale_pgconn_string
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
_lock = RLock()
def _evict_idle_pools(*, protected: str | None = None) -> None:
limit = max(1, settings.PROJECT_TS_CACHE_SIZE)
while len(_pools) > limit:
candidate = next(
(key for key in _pools if key != protected and _pool_borrows.get(key, 0) == 0),
None,
)
if candidate is None:
return
pool = _pools.pop(candidate)
_pool_conninfo.pop(candidate, None)
_pool_borrows.pop(candidate, None)
if not pool.closed:
pool.close()
def get_timescale_pool(db_name: str) -> ConnectionPool:
conninfo = get_project_timescale_pgconn_string(db_name=db_name)
with _lock:
pool = _pools.get(db_name)
if pool is not None and _pool_conninfo.get(db_name) == conninfo and not pool.closed:
_pools.move_to_end(db_name)
return pool
if pool is not None and not pool.closed:
if _pool_borrows.get(db_name, 0):
raise RuntimeError(f"Cannot replace active TimescaleDB pool {db_name!r}")
pool.close()
pool = ConnectionPool(
conninfo=conninfo,
min_size=settings.PROJECT_TS_POOL_MIN_SIZE,
max_size=settings.PROJECT_TS_POOL_MAX_SIZE,
kwargs={"row_factory": dict_row},
open=True,
)
_pools[db_name] = pool
_pool_conninfo[db_name] = conninfo
_pool_borrows.setdefault(db_name, 0)
_evict_idle_pools(protected=db_name)
return pool
@contextmanager
def timescale_connection(db_name: str) -> Iterator[Connection]:
with _lock:
pool = get_timescale_pool(db_name)
_pool_borrows[db_name] = _pool_borrows.get(db_name, 0) + 1
try:
with pool.connection() as conn:
yield conn
finally:
with _lock:
_pool_borrows[db_name] -= 1
_evict_idle_pools()
def close_timescale_pool(db_name: str) -> None:
with _lock:
if _pool_borrows.get(db_name, 0):
raise RuntimeError(f"Cannot close active TimescaleDB pool {db_name!r}")
pool = _pools.pop(db_name, None)
_pool_conninfo.pop(db_name, None)
_pool_borrows.pop(db_name, None)
if pool is not None and not pool.closed:
pool.close()
def close_all_timescale_pools() -> None:
"""Close every synchronous TimescaleDB pool."""
with _lock:
pools = list(_pools.values())
_pools.clear()
_pool_conninfo.clear()
_pool_borrows.clear()
for pool in pools:
if not pool.closed:
pool.close()