新增模拟数据源支持，重构爆管定位逻辑
This commit is contained in:
@@ -3,6 +3,8 @@ 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.timescaledb.schemas.scheme import SchemeRepository
|
||||
from app.infra.db.timescaledb.schemas.realtime import RealtimeRepository
|
||||
@@ -170,3 +172,159 @@ class InternalQueries:
|
||||
time.sleep(1)
|
||||
else:
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def query_realtime_simulation_by_ids_timerange(
|
||||
element_ids: List[str],
|
||||
start_time: str | datetime,
|
||||
end_time: str | datetime,
|
||||
element_type: str,
|
||||
field: str,
|
||||
db_name: str = None,
|
||||
max_retries: int = 3,
|
||||
) -> dict[str, list[dict]]:
|
||||
"""查询实时模拟结果，返回 {id: [{time, value}, ...]}。"""
|
||||
return InternalQueries._query_simulation_by_ids_timerange(
|
||||
schema_name="realtime",
|
||||
element_ids=element_ids,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
element_type=element_type,
|
||||
field=field,
|
||||
db_name=db_name,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def query_scheme_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,
|
||||
db_name: str = None,
|
||||
max_retries: int = 3,
|
||||
) -> dict[str, list[dict]]:
|
||||
"""查询方案模拟结果，返回 {id: [{time, value}, ...]}。"""
|
||||
return InternalQueries._query_simulation_by_ids_timerange(
|
||||
schema_name="scheme",
|
||||
element_ids=element_ids,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
element_type=element_type,
|
||||
field=field,
|
||||
db_name=db_name,
|
||||
max_retries=max_retries,
|
||||
scheme_type=scheme_type,
|
||||
scheme_name=scheme_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _query_simulation_by_ids_timerange(
|
||||
*,
|
||||
schema_name: str,
|
||||
element_ids: List[str],
|
||||
start_time: str | datetime,
|
||||
end_time: str | datetime,
|
||||
element_type: str,
|
||||
field: str,
|
||||
db_name: str = None,
|
||||
max_retries: int = 3,
|
||||
scheme_type: str | None = None,
|
||||
scheme_name: str | None = None,
|
||||
) -> dict[str, list[dict]]:
|
||||
if not element_ids:
|
||||
return {}
|
||||
|
||||
start_dt = (
|
||||
datetime.fromisoformat(start_time)
|
||||
if isinstance(start_time, str)
|
||||
else start_time
|
||||
)
|
||||
end_dt = (
|
||||
datetime.fromisoformat(end_time) if isinstance(end_time, str) else end_time
|
||||
)
|
||||
table_name, 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"}:
|
||||
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。")
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
conn_string = (
|
||||
timescaledb_info.get_pgconn_string(db_name=db_name)
|
||||
if db_name
|
||||
else timescaledb_info.get_pgconn_string()
|
||||
)
|
||||
with psycopg.Connection.connect(conn_string) as conn:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
if schema_name == "scheme":
|
||||
query = sql.SQL(
|
||||
"SELECT id, time, {} FROM {}.{} "
|
||||
"WHERE scheme_type = %s AND scheme_name = %s "
|
||||
"AND time >= %s AND time <= %s AND id = ANY(%s)"
|
||||
).format(
|
||||
sql.Identifier(field),
|
||||
sql.Identifier(schema_name),
|
||||
sql.Identifier(table_name),
|
||||
)
|
||||
cur.execute(
|
||||
query,
|
||||
(
|
||||
scheme_type,
|
||||
scheme_name,
|
||||
start_dt,
|
||||
end_dt,
|
||||
element_ids,
|
||||
),
|
||||
)
|
||||
else:
|
||||
query = sql.SQL(
|
||||
"SELECT id, time, {} FROM {}.{} "
|
||||
"WHERE time >= %s AND time <= %s AND id = ANY(%s)"
|
||||
).format(
|
||||
sql.Identifier(field),
|
||||
sql.Identifier(schema_name),
|
||||
sql.Identifier(table_name),
|
||||
)
|
||||
cur.execute(query, (start_dt, end_dt, element_ids))
|
||||
rows = cur.fetchall()
|
||||
result: dict[str, list[dict]] = {
|
||||
element_id: [] for element_id in element_ids
|
||||
}
|
||||
for row in rows:
|
||||
result.setdefault(row["id"], []).append(
|
||||
{"time": row["time"].isoformat(), "value": row[field]}
|
||||
)
|
||||
for element_id in result:
|
||||
result[element_id].sort(key=lambda item: item["time"])
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"查询尝试 {attempt + 1} 失败: {e}")
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(1)
|
||||
else:
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _resolve_simulation_table(element_type: str) -> tuple[str, set[str]]:
|
||||
normalized_type = element_type.lower()
|
||||
if normalized_type == "node":
|
||||
return "node_simulation", {"actual_demand", "total_head", "pressure", "quality"}
|
||||
if normalized_type == "link":
|
||||
return "link_simulation", {
|
||||
"flow",
|
||||
"friction",
|
||||
"headloss",
|
||||
"quality",
|
||||
"reaction",
|
||||
"setting",
|
||||
"status",
|
||||
"velocity",
|
||||
}
|
||||
raise ValueError(f"Unsupported element_type: {element_type}")
|
||||
|
||||
Reference in New Issue
Block a user