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

refactor(db)!: finalize pooled WNDB v2 migration

Browse Source
This commit is contained in:
jiang
2026-08-27 17:26:22 +08:00
parent fa188af0b1
commit b74799a39d
105 changed files with 4988 additions and 5565 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/infra/db/dynamic_manager.py
+250 -181
View File
@@ -1,40 +1,29 @@
import asyncio
import logging
from collections import OrderedDict
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Dict
from uuid import UUID
from psycopg import AsyncConnection
from psycopg_pool import AsyncConnectionPool
from psycopg.rows import dict_row
from sqlalchemy.engine.url import make_url
from sqlalchemy.ext.asyncio import (
AsyncEngine,
AsyncSession,
async_sessionmaker,
create_async_engine,
)
from app.core.config import settings
logger = logging.getLogger(__name__)
_check_async_connection = AsyncConnectionPool.check_connection
@dataclass(frozen=True)
class PgEngineEntry:
engine: AsyncEngine
sessionmaker: async_sessionmaker[AsyncSession]
connection_url: str
pool_min_size: int
pool_max_size: int
@dataclass(frozen=True)
@dataclass
class PoolEntry:
pool: AsyncConnectionPool
connection_url: str
pool_min_size: int
pool_max_size: int
borrow_count: int = 0
@dataclass(frozen=True)
@@ -45,186 +34,203 @@ class CacheKey:
class ProjectConnectionManager:
def __init__(self) -> None:
self._pg_cache: Dict[CacheKey, PgEngineEntry] = OrderedDict()
self._ts_cache: Dict[CacheKey, PoolEntry] = OrderedDict()
self._pg_raw_cache: Dict[CacheKey, PoolEntry] = OrderedDict()
self._pg_lock = asyncio.Lock()
self._retired_ts: list[tuple[CacheKey, PoolEntry]] = []
self._retired_pg: list[tuple[CacheKey, PoolEntry]] = []
self._ts_lock = asyncio.Lock()
self._pg_raw_lock = asyncio.Lock()
def _normalize_pg_url(self, url: str) -> str:
parsed = make_url(url)
if parsed.drivername in {"postgresql", "postgres"}:
parsed = parsed.set(drivername="postgresql+psycopg")
return parsed.render_as_string(hide_password=False)
async def get_pg_sessionmaker(
async def _get_timescale_pool_locked(
self,
project_id: UUID,
db_role: str,
connection_url: str,
pool_min_size: int,
pool_max_size: int,
) -> async_sessionmaker[AsyncSession]:
async with self._pg_lock:
normalized_url = self._normalize_pg_url(connection_url)
pool_min_size = max(1, pool_min_size)
pool_max_size = max(pool_min_size, pool_max_size)
key = CacheKey(project_id=project_id, db_role=db_role)
entry = self._pg_cache.get(key)
if entry:
if (
entry.connection_url == normalized_url
and entry.pool_min_size == pool_min_size
and entry.pool_max_size == pool_max_size
):
self._pg_cache.move_to_end(key)
return entry.sessionmaker
await entry.engine.dispose()
logger.info(
"Rebuilding PostgreSQL engine for project %s (%s) due to config change",
project_id,
db_role,
)
self._pg_cache.pop(key, None)
engine = create_async_engine(
normalized_url,
pool_size=pool_min_size,
max_overflow=max(0, pool_max_size - pool_min_size),
pool_pre_ping=True,
)
sessionmaker = async_sessionmaker(engine, expire_on_commit=False)
self._pg_cache[key] = PgEngineEntry(
engine=engine,
sessionmaker=sessionmaker,
connection_url=normalized_url,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
)
await self._evict_pg_if_needed()
) -> tuple[CacheKey, AsyncConnectionPool]:
pool_min_size = max(0, pool_min_size)
pool_max_size = max(1, pool_min_size, pool_max_size)
key = CacheKey(project_id=project_id, db_role=db_role)
entry = self._ts_cache.get(key)
if entry:
if (
entry.connection_url == connection_url
and entry.pool_min_size == pool_min_size
and entry.pool_max_size == pool_max_size
):
self._ts_cache.move_to_end(key)
return key, entry.pool
logger.info(
"Created PostgreSQL engine for project %s (%s)", project_id, db_role
"Rebuilding TimescaleDB pool for project %s (%s) due to config change",
project_id,
db_role,
)
return sessionmaker
async def get_timescale_pool(
self,
project_id: UUID,
db_role: str,
connection_url: str,
pool_min_size: int,
pool_max_size: int,
) -> AsyncConnectionPool:
async with self._ts_lock:
pool_min_size = max(1, pool_min_size)
pool_max_size = max(pool_min_size, pool_max_size)
key = CacheKey(project_id=project_id, db_role=db_role)
entry = self._ts_cache.get(key)
if entry:
if (
entry.connection_url == connection_url
and entry.pool_min_size == pool_min_size
and entry.pool_max_size == pool_max_size
):
self._ts_cache.move_to_end(key)
return entry.pool
pool = AsyncConnectionPool(
conninfo=connection_url,
min_size=pool_min_size,
max_size=pool_max_size,
open=False,
kwargs={"row_factory": dict_row},
check=_check_async_connection,
)
await pool.open()
if entry is not None:
if entry.borrow_count:
self._retired_ts.append((key, entry))
else:
await entry.pool.close()
logger.info(
"Rebuilding TimescaleDB pool for project %s (%s) due to config change",
project_id,
db_role,
)
self._ts_cache.pop(key, None)
self._ts_cache[key] = PoolEntry(
pool=pool,
connection_url=connection_url,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
)
logger.info("Created TimescaleDB pool for project %s (%s)", project_id, db_role)
return key, pool
pool = AsyncConnectionPool(
conninfo=connection_url,
min_size=pool_min_size,
max_size=pool_max_size,
open=False,
kwargs={"row_factory": dict_row},
)
await pool.open()
self._ts_cache[key] = PoolEntry(
pool=pool,
connection_url=connection_url,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
)
await self._evict_ts_if_needed()
logger.info(
"Created TimescaleDB pool for project %s (%s)", project_id, db_role
)
return pool
async def get_pg_pool(
async def _get_pg_pool_locked(
self,
project_id: UUID,
db_role: str,
connection_url: str,
pool_min_size: int,
pool_max_size: int,
) -> AsyncConnectionPool:
) -> tuple[CacheKey, AsyncConnectionPool]:
pool_min_size = max(0, pool_min_size)
pool_max_size = max(1, pool_min_size, pool_max_size)
key = CacheKey(project_id=project_id, db_role=db_role)
entry = self._pg_raw_cache.get(key)
if entry:
if (
entry.connection_url == connection_url
and entry.pool_min_size == pool_min_size
and entry.pool_max_size == pool_max_size
):
self._pg_raw_cache.move_to_end(key)
return key, entry.pool
logger.info(
"Rebuilding PostgreSQL pool for project %s (%s) due to config change",
project_id,
db_role,
)
pool = AsyncConnectionPool(
conninfo=connection_url,
min_size=pool_min_size,
max_size=pool_max_size,
open=False,
kwargs={"row_factory": dict_row},
check=_check_async_connection,
)
await pool.open()
if entry is not None:
if entry.borrow_count:
self._retired_pg.append((key, entry))
else:
await entry.pool.close()
self._pg_raw_cache[key] = PoolEntry(
pool=pool,
connection_url=connection_url,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
)
logger.info("Created PostgreSQL pool for project %s (%s)", project_id, db_role)
return key, pool
@asynccontextmanager
async def pg_connection(
self,
project_id: UUID,
db_role: str,
connection_url: str,
pool_min_size: int,
pool_max_size: int,
) -> AsyncIterator[AsyncConnection]:
async with self._pg_raw_lock:
pool_min_size = max(1, pool_min_size)
pool_max_size = max(pool_min_size, pool_max_size)
key = CacheKey(project_id=project_id, db_role=db_role)
entry = self._pg_raw_cache.get(key)
if entry:
if (
entry.connection_url == connection_url
and entry.pool_min_size == pool_min_size
and entry.pool_max_size == pool_max_size
):
self._pg_raw_cache.move_to_end(key)
return entry.pool
await entry.pool.close()
logger.info(
"Rebuilding PostgreSQL pool for project %s (%s) due to config change",
project_id,
db_role,
)
self._pg_raw_cache.pop(key, None)
pool = AsyncConnectionPool(
conninfo=connection_url,
min_size=pool_min_size,
max_size=pool_max_size,
open=False,
kwargs={"row_factory": dict_row},
)
await pool.open()
self._pg_raw_cache[key] = PoolEntry(
pool=pool,
connection_url=connection_url,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
key, pool = await self._get_pg_pool_locked(
project_id,
db_role,
connection_url,
pool_min_size,
pool_max_size,
)
borrowed_entry = self._pg_raw_cache[key]
borrowed_entry.borrow_count += 1
await self._evict_pg_raw_if_needed()
logger.info(
"Created PostgreSQL pool for project %s (%s)", project_id, db_role
)
return pool
async def _evict_pg_if_needed(self) -> None:
while len(self._pg_cache) > settings.PROJECT_PG_CACHE_SIZE:
key, entry = self._pg_cache.popitem(last=False)
await entry.engine.dispose()
logger.info(
"Evicted PostgreSQL engine for project %s (%s)",
key.project_id,
key.db_role,
)
try:
async with pool.connection() as conn:
yield conn
finally:
async with self._pg_raw_lock:
borrowed_entry.borrow_count -= 1
if borrowed_entry.borrow_count == 0 and borrowed_entry.pool is not (
self._pg_raw_cache.get(key).pool
if key in self._pg_raw_cache
else None
):
self._retired_pg = [
item for item in self._retired_pg if item[1] is not borrowed_entry
]
await borrowed_entry.pool.close()
await self._evict_pg_raw_if_needed()
async def _evict_ts_if_needed(self) -> None:
@asynccontextmanager
async def timescale_connection(
self,
project_id: UUID,
db_role: str,
connection_url: str,
pool_min_size: int,
pool_max_size: int,
) -> AsyncIterator[AsyncConnection]:
async with self._ts_lock:
key, pool = await self._get_timescale_pool_locked(
project_id,
db_role,
connection_url,
pool_min_size,
pool_max_size,
)
borrowed_entry = self._ts_cache[key]
borrowed_entry.borrow_count += 1
await self._evict_ts_if_needed()
try:
async with pool.connection() as conn:
yield conn
finally:
async with self._ts_lock:
borrowed_entry.borrow_count -= 1
if borrowed_entry.borrow_count == 0 and borrowed_entry.pool is not (
self._ts_cache.get(key).pool
if key in self._ts_cache
else None
):
self._retired_ts = [
item for item in self._retired_ts if item[1] is not borrowed_entry
]
await borrowed_entry.pool.close()
await self._evict_ts_if_needed()
async def _evict_ts_if_needed(
self, protected_key: CacheKey | None = None
) -> None:
while len(self._ts_cache) > settings.PROJECT_TS_CACHE_SIZE:
key, entry = self._ts_cache.popitem(last=False)
idle = next(
(
(key, entry)
for key, entry in self._ts_cache.items()
if entry.borrow_count == 0 and key != protected_key
),
None,
)
if idle is None:
return
key, entry = idle
self._ts_cache.pop(key)
await entry.pool.close()
logger.info(
"Evicted TimescaleDB pool for project %s (%s)",
@@ -232,9 +238,22 @@ class ProjectConnectionManager:
key.db_role,
)
async def _evict_pg_raw_if_needed(self) -> None:
async def _evict_pg_raw_if_needed(
self, protected_key: CacheKey | None = None
) -> None:
while len(self._pg_raw_cache) > settings.PROJECT_PG_CACHE_SIZE:
key, entry = self._pg_raw_cache.popitem(last=False)
idle = next(
(
(key, entry)
for key, entry in self._pg_raw_cache.items()
if entry.borrow_count == 0 and key != protected_key
),
None,
)
if idle is None:
return
key, entry = idle
self._pg_raw_cache.pop(key)
await entry.pool.close()
logger.info(
"Evicted PostgreSQL pool for project %s (%s)",
@@ -242,17 +261,61 @@ class ProjectConnectionManager:
key.db_role,
)
async def close_all(self) -> None:
async with self._pg_lock:
for key, entry in list(self._pg_cache.items()):
await entry.engine.dispose()
logger.info(
"Closed PostgreSQL engine for project %s (%s)",
key.project_id,
key.db_role,
)
self._pg_cache.clear()
async def close_project(
self, project_id: UUID, db_role: str | None = None
) -> bool:
"""Close this worker's idle pools for a project.
Returns ``False`` without interrupting requests when any matching pool is
currently borrowed. Callers may retry after those requests complete.
"""
closed = True
for cache, lock, label in (
(self._ts_cache, self._ts_lock, "TimescaleDB"),
(self._pg_raw_cache, self._pg_raw_lock, "PostgreSQL"),
):
async with lock:
keys = [
key
for key in cache
if key.project_id == project_id
and (db_role is None or key.db_role == db_role)
]
for key in keys:
entry = cache[key]
if entry.borrow_count:
closed = False
continue
cache.pop(key)
await entry.pool.close()
logger.info(
"Closed %s pool for project %s (%s)",
label,
key.project_id,
key.db_role,
)
for retired, lock in (
(self._retired_ts, self._ts_lock),
(self._retired_pg, self._pg_raw_lock),
):
async with lock:
matches = [
item
for item in retired
if item[0].project_id == project_id
and (db_role is None or item[0].db_role == db_role)
]
if any(entry.borrow_count for _key, entry in matches):
closed = False
for item in matches:
key, entry = item
if entry.borrow_count:
continue
retired.remove(item)
await entry.pool.close()
return closed
async def close_all(self) -> None:
async with self._ts_lock:
for key, entry in list(self._ts_cache.items()):
await entry.pool.close()
@@ -262,6 +325,9 @@ class ProjectConnectionManager:
key.db_role,
)
self._ts_cache.clear()
for _key, entry in self._retired_ts:
await entry.pool.close()
self._retired_ts.clear()
async with self._pg_raw_lock:
for key, entry in list(self._pg_raw_cache.items()):
@@ -272,6 +338,9 @@ class ProjectConnectionManager:
key.db_role,
)
self._pg_raw_cache.clear()
for _key, entry in self._retired_pg:
await entry.pool.close()
self._retired_pg.clear()
project_connection_manager = ProjectConnectionManager()
app/infra/db/project_routing.py
+30 -2
View File
@@ -5,9 +5,9 @@ from contextvars import ContextVar, Token
from dataclasses import dataclass
from typing import Iterator
from psycopg.conninfo import make_conninfo
from psycopg.conninfo import conninfo_to_dict, make_conninfo
from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string
from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string, settings
@dataclass(frozen=True)
@@ -16,6 +16,17 @@ class ActiveProjectRouting:
business_dsn: str
timescale_dsn: str | None = None
@property
def business_database_name(self) -> str:
"""Return the physical BizDB name selected by metadata routing."""
database_name = conninfo_to_dict(self.business_dsn).get("dbname")
if not database_name:
raise RuntimeError(
f"Business database routing for project {self.project_code!r} "
"does not contain a database name"
)
return database_name
_active_project_routing: ContextVar[ActiveProjectRouting | None] = ContextVar(
"active_project_routing",
@@ -42,6 +53,23 @@ def _dsn_for_database(dsn: str, database_name: str) -> str:
return make_conninfo(dsn, dbname=database_name)
def get_project_database_name(name: str) -> str:
"""Resolve a logical project code to its routed physical BizDB name."""
routing = get_active_project_routing()
if routing is not None and name == routing.project_code:
return routing.business_database_name
return name
def get_project_template_database_name(name: str | None = None) -> str:
"""Return the configured immutable template for the WNDB schema version.
The template belongs to the database schema version, not to an individual
logical project or its temporary physical database name.
"""
return settings.WNDB_TEMPLATE_DB_NAME
def get_project_pgconn_string(db_name: str | None = None) -> str:
routing = get_active_project_routing()
if routing is None:
app/infra/db/timescaledb/composite_queries.py
+111 -83
View File
@@ -56,38 +56,40 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
link_devices: dict[str, str] = {}
node_devices: dict[str, str] = {}
for device_id in device_ids:
target_scada = scada_by_id.get(device_id)
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
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 RealtimeRepository.get_link_field_by_time_range(
timescale_conn, start_time, end_time, element_id, "flow"
)
if scada_type in {"pipe_flow", "flow"}:
link_devices[device_id] = target_scada["link_id"]
elif scada_type == "pressure":
# 查询 node 模拟数据
res = await RealtimeRepository.get_node_field_by_time_range(
timescale_conn, start_time, end_time, element_id, "pressure"
)
node_devices[device_id] = target_scada["node_id"]
else:
raise ValueError(f"Unknown SCADA type: {scada_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["scada_id"] = device_id
result[device_id] = res
return result
link_series = await RealtimeRepository.get_link_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(link_devices.values())), "flow",
)
node_series = await RealtimeRepository.get_node_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(node_devices.values())), "pressure",
)
return {
device_id: [
{**item, "scada_id": device_id}
for item in (
link_series.get(element_id, [])
if device_id in link_devices
else node_series.get(element_id, [])
)
]
for device_id, element_id in (link_devices | node_devices).items()
}
@staticmethod
async def get_scada_associated_analysis_simulation_data(
@@ -117,38 +119,41 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
link_devices: dict[str, str] = {}
node_devices: dict[str, str] = {}
for device_id in device_ids:
target_scada = scada_by_id.get(device_id)
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
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 AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, element_id, start_time, end_time, "flow"
)
if scada_type in {"pipe_flow", "flow"}:
link_devices[device_id] = target_scada["link_id"]
elif scada_type == "pressure":
# 查询 node 模拟数据
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, element_id, start_time, end_time, "pressure"
)
node_devices[device_id] = target_scada["node_id"]
else:
raise ValueError(f"Unknown SCADA type: {scada_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["scada_id"] = device_id
result[device_id] = res
return result
link_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "link",
list(dict.fromkeys(link_devices.values())), start_time, end_time, "flow",
)
node_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "node",
list(dict.fromkeys(node_devices.values())), start_time, end_time, "pressure",
)
return {
device_id: [
{**item, "scada_id": device_id}
for item in (
link_series.get(element_id, [])
if device_id in link_devices
else node_series.get(element_id, [])
)
]
for device_id, element_id in (link_devices | node_devices).items()
}
@staticmethod
async def get_realtime_simulation_data(
@@ -175,26 +180,33 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
pipe_ids: list[str] = []
junction_ids: list[str] = []
for feature_id, feature_type in feature_infos:
if feature_type.lower() == "pipe":
# 查询 link 模拟数据
res = await RealtimeRepository.get_link_field_by_time_range(
timescale_conn, start_time, end_time, feature_id, "flow"
)
pipe_ids.append(feature_id)
elif feature_type.lower() == "junction":
# 查询 node 模拟数据
res = await RealtimeRepository.get_node_field_by_time_range(
timescale_conn, start_time, end_time, feature_id, "pressure"
)
junction_ids.append(feature_id)
else:
raise ValueError(f"Unknown type: {feature_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["feature_id"] = feature_id
result[feature_id] = res
return result
link_series = await RealtimeRepository.get_link_fields_by_ids_time_range(
timescale_conn, start_time, end_time, list(dict.fromkeys(pipe_ids)), "flow"
)
node_series = await RealtimeRepository.get_node_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(junction_ids)), "pressure",
)
return {
feature_id: [
{**item, "feature_id": feature_id}
for item in (
link_series.get(feature_id, [])
if feature_type.lower() == "pipe"
else node_series.get(feature_id, [])
)
]
for feature_id, feature_type in feature_infos
}
@staticmethod
async def get_analysis_simulation_data(
@@ -223,25 +235,34 @@ class CompositeQueries:
Raises:
ValueError: 当类型无效时
"""
result = {}
pipe_ids: list[str] = []
junction_ids: list[str] = []
for feature_id, feature_type in feature_infos:
if feature_type.lower() == "pipe":
# 查询 link 模拟数据
res = await AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, feature_id, start_time, end_time, "flow"
)
pipe_ids.append(feature_id)
elif feature_type.lower() == "junction":
# 查询 node 模拟数据
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, feature_id, start_time, end_time, "pressure"
)
junction_ids.append(feature_id)
else:
raise ValueError(f"Unknown type: {feature_type}")
# 添加 feature_id 到每个数据项
for item in res:
item["feature_id"] = feature_id
result[feature_id] = res
return result
link_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "link", list(dict.fromkeys(pipe_ids)),
start_time, end_time, "flow",
)
node_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "node", list(dict.fromkeys(junction_ids)),
start_time, end_time, "pressure",
)
return {
feature_id: [
{**item, "feature_id": feature_id}
for item in (
link_series.get(feature_id, [])
if feature_type.lower() == "pipe"
else node_series.get(feature_id, [])
)
]
for feature_id, feature_type in feature_infos
}
@staticmethod
async def get_element_associated_scada_data(
@@ -399,7 +420,7 @@ class CompositeQueries:
if scada_by_id[device_id]["device_type"] in {"pipe_flow", "flow"}
]
updated_rows = 0
cleaned_rows: list[tuple[datetime, str, float | None]] = []
for grouped_ids, cleaning_function in (
(pressure_ids, clean_pressure_data_df_km),
(flow_ids, clean_flow_data_df_kf),
@@ -422,18 +443,25 @@ class CompositeQueries:
if isinstance(time_value, datetime)
else datetime.fromisoformat(str(time_value))
)
await ScadaRepository.update_scada_field(
timescale_conn,
time_dt,
device_id,
"cleaned_value",
value,
cleaned_rows.append(
(
time_dt,
device_id,
None if pd.isna(value) else float(value),
)
)
updated_rows += 1
if updated_rows == 0:
if not cleaned_rows:
raise ValueError("SCADA 数据清洗未产生任何数据库更新")
updated_rows = await ScadaRepository.update_scada_field_batch(
timescale_conn,
cleaned_rows,
"cleaned_value",
)
if updated_rows == 0:
raise ValueError("SCADA 清洗结果未匹配任何已有监测数据")
return "success"
@staticmethod
app/infra/db/timescaledb/internal_queries.py
+3 -2
View File
@@ -259,7 +259,7 @@ class InternalQueries:
query = sql.SQL(
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE run_id = %s AND time >= %s AND time <= %s "
"AND btrim({}::text) = ANY(%s)"
"AND btrim({}::text) = ANY(%s) ORDER BY id, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
@@ -274,7 +274,8 @@ class InternalQueries:
else:
query = sql.SQL(
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE time >= %s AND time <= %s AND btrim({}::text) = ANY(%s)"
"WHERE time >= %s AND time <= %s "
"AND btrim({}::text) = ANY(%s) ORDER BY id, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
app/infra/db/timescaledb/repositories/analysis.py
+47
View File
@@ -175,6 +175,53 @@ class AnalysisResultsRepository:
await cur.execute(query, (run_id, link_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_series_by_ids(
conn: AsyncConnection,
run_id: UUID,
element_type: str,
element_ids: list[str],
start_time: datetime,
end_time: datetime,
field: str,
) -> dict[str, list[dict[str, Any]]]:
if element_type == "node":
table_name, id_column, valid_fields = (
"node_results", "node_id", AnalysisResultsRepository.NODE_FIELDS
)
elif element_type == "link":
table_name, id_column, valid_fields = (
"link_results", "link_id", AnalysisResultsRepository.LINK_FIELDS
)
else:
raise ValueError(f"invalid analysis element type: {element_type}")
if field not in valid_fields:
raise ValueError(f"invalid {element_type} result field: {field}")
result: dict[str, list[dict[str, Any]]] = {
element_id: [] for element_id in element_ids
}
if not element_ids:
return result
query = sql.SQL(
"SELECT {} AS element_id, time, {} AS value FROM analysis.{} "
"WHERE run_id = %s AND {} = ANY(%s) AND time BETWEEN %s AND %s "
"ORDER BY {}, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
sql.Identifier(table_name),
sql.Identifier(id_column),
sql.Identifier(id_column),
)
async with conn.cursor() as cur:
await cur.execute(query, (run_id, element_ids, start_time, end_time))
for row in await cur.fetchall():
result.setdefault(str(row["element_id"]), []).append(
{"time": row["time"], "value": row["value"]}
)
return result
@staticmethod
async def get_values_at_time(
conn: AsyncConnection,
app/infra/db/timescaledb/repositories/realtime.py
+120 -20
View File
@@ -7,6 +7,19 @@ from app.services.time_api import parse_utc_time
class RealtimeRepository:
@staticmethod
def _batch_time(data: List[dict]) -> datetime:
"""Return one normalized timestamp shared by every row in a snapshot."""
if not data:
raise ValueError("Realtime batch must not be empty")
times = {
parse_utc_time(item["time"], field_name="time")
for item in data
}
if len(times) != 1:
raise ValueError("Realtime batch must contain exactly one timestamp")
return times.pop()
# --- Link Simulation ---
@staticmethod
@@ -15,8 +28,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
async with conn.transaction():
@@ -38,7 +50,7 @@ class RealtimeRepository:
for item in data:
await copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("flow"),
item.get("friction"),
@@ -57,8 +69,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
with conn.transaction():
@@ -80,7 +91,7 @@ class RealtimeRepository:
for item in data:
copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("flow"),
item.get("friction"),
@@ -99,7 +110,8 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
"AND link_id = %s ORDER BY time",
(start_time, end_time, link_id),
)
return await cur.fetchall()
@@ -112,7 +124,8 @@ 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_results WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
"ORDER BY time, link_id",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -140,7 +153,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s"
"SELECT time, {} FROM realtime.link_results WHERE time >= %s "
"AND time <= %s AND link_id = %s ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -150,6 +164,36 @@ class RealtimeRepository:
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_link_fields_by_ids_time_range(
conn: AsyncConnection,
start_time: datetime,
end_time: datetime,
link_ids: list[str],
field: str,
) -> dict[str, list[dict[str, Any]]]:
valid_fields = {
"flow", "friction", "headloss", "quality", "reaction",
"setting", "status", "velocity",
}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
result = {link_id: [] for link_id in link_ids}
if not link_ids:
return result
query = sql.SQL(
"SELECT link_id, time, {} FROM realtime.link_results "
"WHERE time BETWEEN %s AND %s AND link_id = ANY(%s) "
"ORDER BY link_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (start_time, end_time, link_ids))
for row in await cur.fetchall():
result.setdefault(str(row["link_id"]), []).append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return result
@staticmethod
async def get_links_field_by_time_range(
conn: AsyncConnection,
@@ -172,7 +216,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT link_id, time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s"
"SELECT link_id, time, {} FROM realtime.link_results "
"WHERE time >= %s AND time <= %s ORDER BY link_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -230,8 +275,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
async with conn.transaction():
@@ -253,7 +297,7 @@ class RealtimeRepository:
for item in data:
await copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("actual_demand"),
item.get("total_head"),
@@ -267,8 +311,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
with conn.transaction():
@@ -290,7 +333,7 @@ class RealtimeRepository:
for item in data:
copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("actual_demand"),
item.get("total_head"),
@@ -305,7 +348,8 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
"AND node_id = %s ORDER BY time",
(start_time, end_time, node_id),
)
return await cur.fetchall()
@@ -318,7 +362,8 @@ 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_results WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
"ORDER BY time, node_id",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -336,7 +381,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s"
"SELECT time, {} FROM realtime.node_results WHERE time >= %s "
"AND time <= %s AND node_id = %s ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -346,6 +392,33 @@ class RealtimeRepository:
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_node_fields_by_ids_time_range(
conn: AsyncConnection,
start_time: datetime,
end_time: datetime,
node_ids: list[str],
field: str,
) -> dict[str, list[dict[str, Any]]]:
valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
result = {node_id: [] for node_id in node_ids}
if not node_ids:
return result
query = sql.SQL(
"SELECT node_id, time, {} FROM realtime.node_results "
"WHERE time BETWEEN %s AND %s AND node_id = ANY(%s) "
"ORDER BY node_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (start_time, end_time, node_ids))
for row in await cur.fetchall():
result.setdefault(str(row["node_id"]), []).append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return result
@staticmethod
async def get_nodes_field_by_time_range(
conn: AsyncConnection, start_time: datetime, end_time: datetime, field: str
@@ -355,7 +428,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT node_id, time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s"
"SELECT node_id, time, {} FROM realtime.node_results "
"WHERE time >= %s AND time <= %s ORDER BY node_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -459,6 +533,19 @@ class RealtimeRepository:
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(simulation_time,),
)
await cur.execute(
"DELETE FROM realtime.node_results WHERE time = %s",
(simulation_time,),
)
await cur.execute(
"DELETE FROM realtime.link_results WHERE time = %s",
(simulation_time,),
)
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
@@ -525,6 +612,19 @@ class RealtimeRepository:
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(simulation_time,),
)
cur.execute(
"DELETE FROM realtime.node_results WHERE time = %s",
(simulation_time,),
)
cur.execute(
"DELETE FROM realtime.link_results WHERE time = %s",
(simulation_time,),
)
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
app/infra/db/timescaledb/repositories/scada.py
+32 -3
View File
@@ -35,7 +35,8 @@ class ScadaRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scada.measurements 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 ORDER BY device_id, time",
(device_ids, start_time, end_time),
)
return await cur.fetchall()
@@ -49,7 +50,8 @@ class ScadaRepository:
) -> List[dict]:
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"SELECT * FROM scada.measurements 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 ORDER BY device_id, time",
(device_ids, start_time, end_time),
)
return cur.fetchall()
@@ -88,7 +90,9 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT device_id, time, {} FROM scada.measurements 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) "
"ORDER BY device_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -122,6 +126,31 @@ class ScadaRepository:
if cur.rowcount == 0:
await cur.execute(insert_query, (time, device_id, value))
@staticmethod
async def update_scada_field_batch(
conn: AsyncConnection,
rows: list[tuple[datetime, str, float | None]],
field: str,
) -> int:
"""Update existing SCADA samples in one set-based statement."""
valid_fields = {"monitored_value", "cleaned_value"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
if not rows:
return 0
query = sql.SQL(
"UPDATE scada.measurements AS measurement SET {} = batch.value "
"FROM unnest(%s::timestamptz[], %s::text[], %s::double precision[]) "
"AS batch(time, device_id, value) "
"WHERE measurement.time = batch.time "
"AND measurement.device_id = batch.device_id"
).format(sql.Identifier(field))
times, device_ids, values = zip(*rows)
async with conn.cursor() as cur:
await cur.execute(query, (list(times), list(device_ids), list(values)))
return cur.rowcount
@staticmethod
async def delete_scada_by_id_time_range(
conn: AsyncConnection, device_id: str, start_time: datetime, end_time: datetime
app/infra/db/timescaledb/sync_pool.py
+2
View File
@@ -11,6 +11,7 @@ from app.core.config import settings
from app.infra.db.project_routing import get_project_timescale_pgconn_string
_check_connection = ConnectionPool.check_connection
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
@@ -49,6 +50,7 @@ def get_timescale_pool(db_name: str) -> ConnectionPool:
min_size=settings.PROJECT_TS_POOL_MIN_SIZE,
max_size=settings.PROJECT_TS_POOL_MAX_SIZE,
kwargs={"row_factory": dict_row},
check=_check_connection,
open=True,
)
_pools[db_name] = pool