重构现代化 FastAPI 后端项目框架
This commit is contained in:
760
app/infra/db/timescaledb/schemas/scheme.py
Normal file
760
app/infra/db/timescaledb/schemas/scheme.py
Normal file
@@ -0,0 +1,760 @@
|
||||
from typing import List, Any, Dict
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from collections import defaultdict
|
||||
from psycopg import AsyncConnection, Connection, sql
|
||||
import globals
|
||||
|
||||
# 定义UTC+8时区
|
||||
UTC_8 = timezone(timedelta(hours=8))
|
||||
|
||||
|
||||
class SchemeRepository:
|
||||
|
||||
# --- 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,
|
||||
):
|
||||
"""
|
||||
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)
|
||||
"""
|
||||
# Convert result_start_time string to datetime if needed
|
||||
if isinstance(result_start_time, str):
|
||||
# 如果是ISO格式字符串，解析并转换为UTC+8
|
||||
if result_start_time.endswith("Z"):
|
||||
# UTC时间，转换为UTC+8
|
||||
utc_time = datetime.fromisoformat(
|
||||
result_start_time.replace("Z", "+00:00")
|
||||
)
|
||||
simulation_time = utc_time.astimezone(UTC_8)
|
||||
else:
|
||||
# 假设已经是UTC+8时间
|
||||
simulation_time = datetime.fromisoformat(result_start_time)
|
||||
if simulation_time.tzinfo is None:
|
||||
simulation_time = simulation_time.replace(tzinfo=UTC_8)
|
||||
else:
|
||||
simulation_time = result_start_time
|
||||
if simulation_time.tzinfo is None:
|
||||
simulation_time = simulation_time.replace(tzinfo=UTC_8)
|
||||
|
||||
timestep_parts = globals.hydraulic_timestep.split(":")
|
||||
timestep = timedelta(
|
||||
hours=int(timestep_parts[0]),
|
||||
minutes=int(timestep_parts[1]),
|
||||
seconds=int(timestep_parts[2]),
|
||||
)
|
||||
|
||||
# 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):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = node_result.get("result", [])[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")
|
||||
for period_index in range(num_periods):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = link_result.get("result", [])[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,
|
||||
):
|
||||
"""
|
||||
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)
|
||||
"""
|
||||
# Convert result_start_time string to datetime if needed
|
||||
if isinstance(result_start_time, str):
|
||||
# 如果是ISO格式字符串，解析并转换为UTC+8
|
||||
if result_start_time.endswith("Z"):
|
||||
# UTC时间，转换为UTC+8
|
||||
utc_time = datetime.fromisoformat(
|
||||
result_start_time.replace("Z", "+00:00")
|
||||
)
|
||||
simulation_time = utc_time.astimezone(UTC_8)
|
||||
else:
|
||||
# 假设已经是UTC+8时间
|
||||
simulation_time = datetime.fromisoformat(result_start_time)
|
||||
if simulation_time.tzinfo is None:
|
||||
simulation_time = simulation_time.replace(tzinfo=UTC_8)
|
||||
else:
|
||||
simulation_time = result_start_time
|
||||
if simulation_time.tzinfo is None:
|
||||
simulation_time = simulation_time.replace(tzinfo=UTC_8)
|
||||
|
||||
timestep_parts = globals.hydraulic_timestep.split(":")
|
||||
timestep = timedelta(
|
||||
hours=int(timestep_parts[0]),
|
||||
minutes=int(timestep_parts[1]),
|
||||
seconds=int(timestep_parts[2]),
|
||||
)
|
||||
|
||||
# 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):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = node_result.get("result", [])[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")
|
||||
for period_index in range(num_periods):
|
||||
current_time = simulation_time + (timestep * period_index)
|
||||
data = link_result.get("result", [])[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
|
||||
"""
|
||||
# Convert query_time string to datetime
|
||||
if isinstance(query_time, str):
|
||||
if query_time.endswith("Z"):
|
||||
# UTC时间，转换为UTC+8
|
||||
utc_time = datetime.fromisoformat(query_time.replace("Z", "+00:00"))
|
||||
target_time = utc_time.astimezone(UTC_8)
|
||||
else:
|
||||
# 假设已经是UTC+8时间
|
||||
target_time = datetime.fromisoformat(query_time)
|
||||
if target_time.tzinfo is None:
|
||||
target_time = target_time.replace(tzinfo=UTC_8)
|
||||
else:
|
||||
target_time = query_time
|
||||
if target_time.tzinfo is None:
|
||||
target_time = target_time.replace(tzinfo=UTC_8)
|
||||
|
||||
# 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
|
||||
"""
|
||||
# Convert query_time string to datetime
|
||||
if isinstance(query_time, str):
|
||||
if query_time.endswith("Z"):
|
||||
# UTC时间，转换为UTC+8
|
||||
utc_time = datetime.fromisoformat(query_time.replace("Z", "+00:00"))
|
||||
target_time = utc_time.astimezone(UTC_8)
|
||||
else:
|
||||
# 假设已经是UTC+8时间
|
||||
target_time = datetime.fromisoformat(query_time)
|
||||
if target_time.tzinfo is None:
|
||||
target_time = target_time.replace(tzinfo=UTC_8)
|
||||
else:
|
||||
target_time = query_time
|
||||
if target_time.tzinfo is None:
|
||||
target_time = target_time.replace(tzinfo=UTC_8)
|
||||
|
||||
# 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'")
|
||||
Reference in New Issue
Block a user