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

refactor(storage): route project DSNs and remove legacy backends

Browse Source
This commit is contained in:
jiang
2026-08-18 18:29:09 +08:00
parent b21eaffe40
commit 6b09662de6
51 changed files with 542 additions and 10951 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/v1/endpoints/cache.py
-57
View File
@@ -1,57 +0,0 @@
from fastapi import APIRouter, Query
from app.infra.cache.redis_client import redis_client
router = APIRouter()
@router.delete("/redis-keys/detail", summary="清除单个缓存键", description="根据键名清除单个Redis缓存")
async def fastapi_clear_redis_key(key: str = Query(..., description="缓存键名")):
"""
清除单个缓存键
根据指定的键名删除Redis中对应的缓存
"""
redis_client.delete(key)
return True
@router.delete("/redis-keys", summary="清除匹配的缓存键", description="根据模式清除匹配的Redis缓存键")
async def fastapi_clear_redis_keys(keys: str = Query(..., description="缓存键模式（支持通配符）")):
"""
清除匹配的缓存键
根据指定的模式删除Redis中所有匹配的缓存键
"""
# delete keys contains the key
matched_keys = redis_client.keys(f"*{keys}*")
if matched_keys:
redis_client.delete(*matched_keys)
return True
@router.delete("/all-redis", summary="清除所有缓存", description="清空整个Redis数据库的所有缓存")
async def fastapi_clear_all_redis():
"""
清除所有缓存
清空Redis数据库中的所有缓存键值对
"""
redis_client.flushdb()
return True
@router.get("/redis", summary="查询缓存键列表", description="获取Redis中所有的缓存键")
async def fastapi_query_redis():
"""
查询缓存键列表
获取Redis数据库中所有的缓存键列表
"""
# Helper to decode bytes to str for JSON response if needed,
# but original just returned keys (which might be bytes in redis-py unless decode_responses=True)
# create_redis_client usually sets decode_responses=False by default.
# We will assume user handles bytes or we should decode.
# Original just returned redis_client.keys("*")
keys = redis_client.keys("*")
# Clean output for API
return [k.decode('utf-8') if isinstance(k, bytes) else k for k in keys]
app/api/v1/endpoints/model_import.py
+22 -5
View File
@@ -17,8 +17,13 @@ from app.auth.metadata_dependencies import (
get_current_metadata_admin,
get_metadata_repository,
)
from app.auth.project_dependencies import (
ProjectContext,
resolve_project_business_routing,
)
from app.core.audit import AuditAction, log_audit_event
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
from app.infra.db.project_routing import activate_project_routing
from app.services.network_import import network_update
from app.services.tjnetwork import run_inp
@@ -118,21 +123,21 @@ async def _run_uploaded_inp(content: bytes) -> str:
return run_inp(model_name)
async def _update_from_inp(content: bytes) -> None:
async def _update_from_inp(content: bytes, project_code: str) -> None:
temp_path: Path | None = None
try:
with NamedTemporaryFile(suffix=".inp", delete=False) as temp_file:
temp_file.write(content)
temp_path = Path(temp_file.name)
network_update(str(temp_path))
network_update(str(temp_path), project_code)
finally:
if temp_path is not None:
temp_path.unlink(missing_ok=True)
async def _apply_model_update(content: bytes) -> None:
async def _apply_model_update(content: bytes, project_code: str) -> None:
try:
await _update_from_inp(content)
await _update_from_inp(content, project_code)
except Exception as exc:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -177,7 +182,19 @@ async def update_project_model(
) -> dict:
project = await _get_active_project(project_id, metadata_repo)
content, filename = await _read_upload(file)
await _apply_model_update(content)
routing = await resolve_project_business_routing(
ProjectContext(
project_id=project.id,
project_code=project.code,
user_id=current_user.id,
project_role="owner",
system_role=current_user.role,
is_superuser=current_user.is_superuser,
),
metadata_repo,
)
with activate_project_routing(routing):
await _apply_model_update(content, project.code)
await _audit_model_change(
request=request,
current_user=current_user,
app/api/v1/endpoints/network/junctions.py
-1
View File
@@ -333,7 +333,6 @@ async def fastapi_get_all_junction_properties(
list: 包含所有节点属性的列表
"""
# 缓存查询结果提高性能
# global redis_client # Redis logic removed for clean split, can be re-added if needed or imported
results = get_all_junctions(network)
return results
app/api/v1/endpoints/network/pipes.py
-1
View File
@@ -385,7 +385,6 @@ async def fastapi_get_all_pipe_properties(
包含所有管道属性的字典列表
"""
# 缓存查询结果提高性能
# global redis_client
results = get_all_pipes(network)
return results
app/api/v1/endpoints/network/pumps.py
-1
View File
@@ -177,7 +177,6 @@ async def fastapi_get_all_pump_properties(
包含所有水泵属性的字典列表
"""
# 缓存查询结果提高性能
# global redis_client
results = get_all_pumps(network)
return results
app/api/v1/endpoints/network/tanks.py
-1
View File
@@ -540,7 +540,6 @@ async def fastapi_get_all_tank_properties(
包含所有水箱属性的字典列表
"""
# 缓存查询结果提高性能
# global redis_client
results = get_all_tanks(network)
return results
app/api/v1/endpoints/network/valves.py
-1
View File
@@ -307,7 +307,6 @@ async def fastapi_get_all_valve_properties(
返回指定水网中所有阀门的完整属性列表。
"""
# 缓存查询结果提高性能
# global redis_client
results = get_all_valves(network)
return results
app/api/v1/rest_router.py
+42 -5
View File
@@ -17,7 +17,16 @@ from app.api.problem_details import ProblemDetails
from app.api.pagination import PaginatedList
from app.api.v1.router import api_router as handler_api_router
from app.auth.metadata_dependencies import get_current_metadata_user
from app.auth.project_dependencies import ProjectContext, get_project_context
from app.auth.project_dependencies import (
ProjectContext,
get_project_business_routing,
get_project_context,
get_project_simulation_routing,
)
from app.infra.db.project_routing import (
ActiveProjectRouting,
activate_project_routing,
)
T = TypeVar("T")
@@ -44,6 +53,13 @@ _PUBLIC_PARAMETER_RENAMES = {
}
_MODEL_NAME_IS_NETWORK = {"RunSimulationManuallyByDate", "PressureSensorPlacement"}
_MODEL_USERNAME_FROM_AUTH = {"PressureSensorPlacement"}
_TIMESCALE_ROUTED_ENDPOINT_MODULES = {
"app.api.v1.endpoints.burst_detection",
"app.api.v1.endpoints.burst_location",
"app.api.v1.endpoints.leakage",
"app.api.v1.endpoints.simulation",
}
_TIMESCALE_ROUTED_ENDPOINT_NAMES = {"open_project_endpoint"}
def _clean_name(name: str) -> str:
@@ -128,10 +144,14 @@ def _with_header_project_context(endpoint, route_name: str):
None,
)
injected_context_name = existing_context_parameter or "_rest_project_context"
injected_routing_name = "_rest_project_routing"
injected_user_name = "_rest_current_user"
@wraps(endpoint)
async def wrapper(*args, **kwargs):
project_routing = kwargs.pop(injected_routing_name, None)
if not isinstance(project_routing, ActiveProjectRouting):
raise RuntimeError("REST project database routing was not resolved")
project_context = kwargs.get(injected_context_name)
if not isinstance(project_context, ProjectContext):
raise RuntimeError("REST project context was not resolved")
@@ -162,10 +182,11 @@ def _with_header_project_context(endpoint, route_name: str):
kwargs[parameter_name] = original_model.model_validate(data)
if model_has_username:
kwargs.pop(injected_user_name, None)
result = endpoint(*args, **kwargs)
if inspect.isawaitable(result):
return await result
return result
with activate_project_routing(project_routing):
result = endpoint(*args, **kwargs)
if inspect.isawaitable(result):
return await result
return result
parameters = []
for name, parameter in signature.parameters.items():
@@ -181,6 +202,14 @@ def _with_header_project_context(endpoint, route_name: str):
if name in body_models:
parameter = parameter.replace(annotation=body_models[name][1])
parameters.append(parameter)
routing_dependency = (
get_project_simulation_routing
if (
endpoint.__module__ in _TIMESCALE_ROUTED_ENDPOINT_MODULES
or endpoint.__name__ in _TIMESCALE_ROUTED_ENDPOINT_NAMES
)
else get_project_business_routing
)
if not existing_context_parameter:
parameters.append(
inspect.Parameter(
@@ -190,6 +219,14 @@ def _with_header_project_context(endpoint, route_name: str):
default=Depends(get_project_context),
)
)
parameters.append(
inspect.Parameter(
injected_routing_name,
kind=inspect.Parameter.KEYWORD_ONLY,
annotation=ActiveProjectRouting,
default=Depends(routing_dependency),
)
)
if username_parameter or model_has_username:
parameters.append(
inspect.Parameter(
app/api/v1/router.py
-8
View File
@@ -7,7 +7,6 @@ from app.api.v1.endpoints import (
audit,
burst_detection,
burst_location,
cache,
extension,
geocoding,
leakage,
@@ -53,7 +52,6 @@ from app.api.v1.endpoints.timeseries import (
)
from app.auth.permissions import (
BURST_RUN,
ENVIRONMENT_MANAGE,
OPTIMIZATION_RUN,
RISK_RUN,
SCADA_CLEAN,
@@ -89,7 +87,6 @@ simulation_access = Depends(
webgis_view_access = Depends(require_permission(WEBGIS_VIEW))
simulation_run_access = Depends(require_permission(SIMULATION_RUN))
environment_manage_access = Depends(require_permission(ENVIRONMENT_MANAGE))
burst_run_access = Depends(require_permission(BURST_RUN))
risk_run_access = Depends(require_permission(RISK_RUN))
optimization_run_access = Depends(require_permission(OPTIMIZATION_RUN))
@@ -168,11 +165,6 @@ api_router.include_router(
tags=["Risk"],
dependencies=[risk_run_access],
)
api_router.include_router(
cache.router,
tags=["Cache"],
dependencies=[environment_manage_access],
)
api_router.include_router(
web_search.router,
tags=["Web Search"],
app/auth/project_dependencies.py
+57
View File
@@ -16,6 +16,7 @@ from app.infra.db.metadb.repositories.metadata_repository import (
MetadataRepository,
ProjectDbRouting,
)
from app.infra.db.project_routing import ActiveProjectRouting
DB_ROLE_BIZ_DATA = "biz_data"
DB_ROLE_IOT_DATA = "iot_data"
@@ -99,6 +100,62 @@ async def get_project_context(
return await resolve_project_context(x_project_id, current_user, metadata_repo)
async def get_project_business_routing(
ctx: ProjectContext = Depends(get_project_context),
metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> ActiveProjectRouting:
return await resolve_project_business_routing(ctx, metadata_repo)
async def resolve_project_business_routing(
ctx: ProjectContext,
metadata_repo: MetadataRepository,
) -> ActiveProjectRouting:
business = await _get_project_routing(
metadata_repo,
ctx.project_id,
DB_ROLE_BIZ_DATA,
DB_TYPE_POSTGRES,
"PostgreSQL",
)
return ActiveProjectRouting(
project_code=ctx.project_code,
business_dsn=business.dsn,
)
async def get_project_simulation_routing(
ctx: ProjectContext = Depends(get_project_context),
metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> ActiveProjectRouting:
return await resolve_project_simulation_routing(ctx, metadata_repo)
async def resolve_project_simulation_routing(
ctx: ProjectContext,
metadata_repo: MetadataRepository,
) -> ActiveProjectRouting:
business = await _get_project_routing(
metadata_repo,
ctx.project_id,
DB_ROLE_BIZ_DATA,
DB_TYPE_POSTGRES,
"PostgreSQL",
)
timescale = await _get_project_routing(
metadata_repo,
ctx.project_id,
DB_ROLE_IOT_DATA,
DB_TYPE_TIMESCALE,
"TimescaleDB",
)
return ActiveProjectRouting(
project_code=ctx.project_code,
business_dsn=business.dsn,
timescale_dsn=timescale.dsn,
)
async def _get_project_routing(
metadata_repo: MetadataRepository,
project_id: UUID,
app/core/config.py
-6
View File
@@ -26,12 +26,6 @@ class Settings(BaseSettings):
TIMESCALEDB_DB_PORT: str = "5433"
TIMESCALEDB_DB_USER: str = "postgres"
TIMESCALEDB_DB_PASSWORD: str = "password"
# InfluxDB
INFLUXDB_URL: str = "http://localhost:8086"
INFLUXDB_TOKEN: str = "token"
INFLUXDB_ORG: str = "org"
INFLUXDB_BUCKET: str = "bucket"
# Metadata Database Config (PostgreSQL)
METADATA_DB_NAME: str = "system_hub"
METADATA_DB_HOST: str = "localhost"
app/infra/cache/__init__.py
Vendored
View File
app/infra/cache/redis_client.py
Vendored
-19
View File
@@ -1,19 +0,0 @@
import redis
import msgpack
from datetime import datetime
from typing import Any
# Initialize Redis connection
redis_client = redis.Redis(host="127.0.0.1", port=6379, db=0)
def encode_datetime(obj: Any) -> Any:
"""Serialize datetime objects to dictionary format."""
if isinstance(obj, datetime):
return {"__datetime__": True, "as_str": obj.strftime("%Y%m%dT%H:%M:%S.%f")}
return obj
def decode_datetime(obj: Any) -> Any:
"""Deserialize dictionary format to datetime objects."""
if "__datetime__" in obj:
return datetime.strptime(obj["as_str"], "%Y%m%dT%H:%M:%S.%f")
return obj
app/infra/db/influxdb/__init__.py
View File
app/infra/db/influxdb/api.py
-4964
View File
File diff suppressed because it is too large Load Diff
app/infra/db/influxdb/info.py
-5
View File
@@ -1,5 +0,0 @@
from app.core.config import settings
url = settings.INFLUXDB_URL
token = settings.INFLUXDB_TOKEN
org = settings.INFLUXDB_ORG
app/infra/db/influxdb/query.py
-33
View File
@@ -1,33 +0,0 @@
from influxdb_client import InfluxDBClient, Point, WriteOptions
from influxdb_client.client.query_api import QueryApi
import influxdb_info
# 配置 InfluxDB 连接
url = influxdb_info.url
token = influxdb_info.token
org = influxdb_info.org
bucket = "SCADA_data"
# 创建 InfluxDB 客户端
client = InfluxDBClient(url=url, token=token, org=org)
# 创建查询 API 对象
query_api = client.query_api()
# 构建查询语句
query = f'''
from(bucket: "{bucket}")
|> range(start: -1h)
'''
# 执行查询
result = query_api.query(query)
print(result)
# 处理查询结果
for table in result:
for record in table.records:
print(f"Time: {record.get_time()}, Value: {record.get_value()}, Measurement: {record.get_measurement()}, Field: {record.get_field()}")
# 关闭客户端连接
client.close()
app/infra/db/metadb/repositories/metadata_repository.py
+6 -3
View File
@@ -20,14 +20,17 @@ def _normalize_postgres_dsn(dsn: str) -> str:
scheme, rest = dsn.split("://", 1)
if scheme not in ("postgresql", "postgres", "postgresql+psycopg"):
return dsn
if scheme == "postgresql+psycopg":
scheme = "postgresql"
normalized_dsn = f"{scheme}://{rest}"
if "@" not in rest:
return dsn
return normalized_dsn
userinfo, hostinfo = rest.rsplit("@", 1)
if ":" not in userinfo:
return dsn
return normalized_dsn
username, password = userinfo.split(":", 1)
if "@" not in password:
return dsn
return normalized_dsn
password = password.replace("@", "%40")
return f"{scheme}://{username}:{password}@{hostinfo}"
app/infra/db/postgresql/database.py
+11 -3
View File
@@ -3,7 +3,7 @@ from contextlib import asynccontextmanager
from typing import AsyncGenerator, Dict, Optional
import psycopg_pool
from psycopg.rows import dict_row
import app.core.config as postgresql_info
from app.infra.db.project_routing import get_project_pgconn_string
# Configure logging
logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ 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."""
@@ -21,9 +22,10 @@ class Database:
# Get connection string, handling default case where target_db_name might be None
if target_db_name:
conn_string = postgresql_info.get_pgconn_string(db_name=target_db_name)
conn_string = get_project_pgconn_string(db_name=target_db_name)
else:
conn_string = postgresql_info.get_pgconn_string()
conn_string = get_project_pgconn_string()
self.conninfo = conn_string
try:
self.pool = psycopg_pool.AsyncConnectionPool(
@@ -75,6 +77,12 @@ async def get_database_instance(db_name: Optional[str] = None) -> Database:
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)
app/infra/db/project_routing.py
+65
View File
@@ -0,0 +1,65 @@
from __future__ import annotations
from contextlib import contextmanager
from contextvars import ContextVar, Token
from dataclasses import dataclass
from typing import Iterator
from psycopg.conninfo import make_conninfo
from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string
@dataclass(frozen=True)
class ActiveProjectRouting:
project_code: str
business_dsn: str
timescale_dsn: str | None = None
_active_project_routing: ContextVar[ActiveProjectRouting | None] = ContextVar(
"active_project_routing",
default=None,
)
def get_active_project_routing() -> ActiveProjectRouting | None:
return _active_project_routing.get()
@contextmanager
def activate_project_routing(
routing: ActiveProjectRouting,
) -> Iterator[ActiveProjectRouting]:
token: Token[ActiveProjectRouting | None] = _active_project_routing.set(routing)
try:
yield routing
finally:
_active_project_routing.reset(token)
def _dsn_for_database(dsn: str, database_name: str) -> str:
return make_conninfo(dsn, dbname=database_name)
def get_project_pgconn_string(db_name: str | None = None) -> str:
routing = get_active_project_routing()
if routing is None:
return get_pgconn_string(db_name=db_name)
if db_name is None or db_name == routing.project_code:
return routing.business_dsn
return _dsn_for_database(routing.business_dsn, db_name)
def get_project_timescale_pgconn_string(db_name: str | None = None) -> str:
routing = get_active_project_routing()
if routing is None:
return get_timescaledb_pgconn_string(db_name=db_name)
if routing.timescale_dsn is None:
raise RuntimeError(
f"TimescaleDB routing is not configured for project {routing.project_code}"
)
# Legacy simulation code used to derive the Timescale database name from
# the project code. Project-scoped requests must instead use the complete
# iot_data DSN selected by metadata routing.
return routing.timescale_dsn
app/infra/db/timescaledb/database.py
+13 -5
View File
@@ -3,7 +3,7 @@ from contextlib import asynccontextmanager
from typing import AsyncGenerator, Dict, Optional
import psycopg_pool
from psycopg.rows import dict_row
from app.core.config import get_timescaledb_pgconn_string
from app.infra.db.project_routing import get_project_timescale_pgconn_string
# Configure logging
logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ 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."""
@@ -21,9 +22,10 @@ class Database:
# Get connection string, handling default case where target_db_name might be None
if target_db_name:
conn_string = get_timescaledb_pgconn_string(db_name=target_db_name)
conn_string = get_project_timescale_pgconn_string(db_name=target_db_name)
else:
conn_string = get_timescaledb_pgconn_string()
conn_string = get_project_timescale_pgconn_string()
self.conninfo = conn_string
try:
self.pool = psycopg_pool.AsyncConnectionPool(
@@ -54,8 +56,8 @@ class Database:
"""Get the TimescaleDB connection string."""
target_db_name = db_name or self.db_name
if target_db_name:
return get_timescaledb_pgconn_string(db_name=target_db_name)
return get_timescaledb_pgconn_string()
return get_project_timescale_pgconn_string(db_name=target_db_name)
return get_project_timescale_pgconn_string()
@asynccontextmanager
async def get_connection(self) -> AsyncGenerator:
@@ -84,6 +86,12 @@ async def get_database_instance(db_name: Optional[str] = None) -> Database:
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)
app/infra/db/timescaledb/internal_queries.py
+13 -13
View File
@@ -6,7 +6,7 @@ import psycopg
from psycopg import sql
from psycopg.rows import dict_row
import time
from app.core.config import get_timescaledb_pgconn_string
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.repositories.realtime import RealtimeRepository
from app.infra.db.timescaledb.repositories.scada import ScadaRepository
@@ -26,9 +26,9 @@ class InternalStorage:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
RealtimeRepository.store_realtime_simulation_result_sync(
@@ -58,9 +58,9 @@ class InternalStorage:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
SchemeRepository.store_scheme_simulation_result_sync(
@@ -99,9 +99,9 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
rows = ScadaRepository.get_scada_by_ids_time_range_sync(
@@ -140,9 +140,9 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
rows = ScadaRepository.get_scada_by_ids_time_range_sync(
@@ -185,9 +185,9 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
return ScadaRepository.get_latest_scada_time_sync(
@@ -286,9 +286,9 @@ class InternalQueries:
for attempt in range(max_retries):
try:
conn_string = (
get_timescaledb_pgconn_string(db_name=db_name)
get_project_timescale_pgconn_string(db_name=db_name)
if db_name
else get_timescaledb_pgconn_string()
else get_project_timescale_pgconn_string()
)
with psycopg.Connection.connect(conn_string) as conn:
with conn.cursor(row_factory=dict_row) as cur:
app/native/wndb/connection.py
+17 -5
View File
@@ -4,9 +4,10 @@ from threading import RLock
import psycopg as pg
from app.core.config import get_pgconn_string
from app.infra.db.project_routing import get_project_pgconn_string
g_conn_dict: dict[str, pg.Connection] = {}
g_conninfo_dict: dict[str, str] = {}
_registry_lock = RLock()
_project_locks: dict[str, RLock] = {}
@@ -42,14 +43,18 @@ def _get_project_lock(name: str) -> RLock:
def open_connection(name: str) -> pg.Connection:
with _get_project_lock(name):
conninfo = get_project_pgconn_string(db_name=name)
connection = g_conn_dict.get(name)
if connection is None or not _is_healthy(connection):
if (
connection is None
or g_conninfo_dict.get(name) != conninfo
or not _is_healthy(connection)
):
if connection is not None:
_close_connection(connection)
connection = pg.connect(
conninfo=get_pgconn_string(db_name=name), autocommit=True
)
connection = pg.connect(conninfo=conninfo, autocommit=True)
g_conn_dict[name] = connection
g_conninfo_dict[name] = conninfo
return connection
@@ -60,6 +65,12 @@ def is_connection_open(name: str) -> bool:
return False
if not _is_healthy(connection):
del g_conn_dict[name]
g_conninfo_dict.pop(name, None)
_close_connection(connection)
return False
if g_conninfo_dict.get(name) != get_project_pgconn_string(db_name=name):
del g_conn_dict[name]
g_conninfo_dict.pop(name, None)
_close_connection(connection)
return False
return True
@@ -68,6 +79,7 @@ def is_connection_open(name: str) -> bool:
def close_connection(name: str) -> None:
with _get_project_lock(name):
connection = g_conn_dict.pop(name, None)
g_conninfo_dict.pop(name, None)
if connection is not None:
_close_connection(connection)
app/native/wndb/project.py
+11 -8
View File
@@ -7,7 +7,8 @@ from .connection import (
is_connection_open,
open_connection,
)
from app.core.config import get_pgconn_string, get_pg_config, get_pg_password
from app.core.config import get_pg_config, get_pg_password
from app.infra.db.project_routing import get_project_pgconn_string
# no undo/redo
@@ -16,7 +17,7 @@ _server_databases = ["template0", "template1", "postgres", "project"]
def list_project() -> list[str]:
ps = []
with pg.connect(conninfo=get_pgconn_string(), autocommit=True) as conn:
with pg.connect(conninfo=get_project_pgconn_string(), autocommit=True) as conn:
with conn.cursor(row_factory=dict_row) as cur:
for p in cur.execute(
f"select datname from pg_database where datname <> 'postgres' and datname <> 'template0' and datname <> 'template1' and datname <> 'project'"
@@ -27,7 +28,7 @@ def list_project() -> list[str]:
def have_project(name: str) -> bool:
with pg.connect(
conninfo=get_pgconn_string(db_name="postgres"), autocommit=True
conninfo=get_project_pgconn_string(db_name="postgres"), autocommit=True
) as conn:
with conn.cursor() as cur:
cur.execute("select 1 from pg_database where datname = %s", (name,))
@@ -38,7 +39,7 @@ def copy_project(source: str, new: str) -> None:
close_connection(source)
with pg.connect(
conninfo=get_pgconn_string(db_name="postgres"), autocommit=True
conninfo=get_project_pgconn_string(db_name="postgres"), autocommit=True
) as admin_conn:
with admin_conn.cursor() as cur:
cur.execute(
@@ -131,7 +132,9 @@ class CopyProjectEx:
connection.commit()
def __call__(self, source: str, new_db: str, excluded_tables: [str] = None) -> None:
source_connection = pg.connect(conninfo=get_pgconn_string(), autocommit=True)
source_connection = pg.connect(
conninfo=get_project_pgconn_string(), autocommit=True
)
self.create_database(source_connection, new_db)
@@ -140,7 +143,7 @@ class CopyProjectEx:
source_connection.close()
new_db_connection = pg.connect(
conninfo=get_pgconn_string(db_name=new_db), autocommit=True
conninfo=get_project_pgconn_string(db_name=new_db), autocommit=True
)
self.init_operation_table(new_db_connection, excluded_tables)
new_db_connection.close()
@@ -151,7 +154,7 @@ def create_project(name: str) -> None:
def delete_project(name: str) -> None:
with pg.connect(conninfo=get_pgconn_string(), autocommit=True) as conn:
with pg.connect(conninfo=get_project_pgconn_string(), autocommit=True) as conn:
with conn.cursor() as cur:
cur.execute(
f"select pg_terminate_backend(pid) from pg_stat_activity where datname = '{name}'"
@@ -161,7 +164,7 @@ def delete_project(name: str) -> None:
def clean_project(excluded: list[str] = []) -> None:
projects = list_project()
with pg.connect(conninfo=get_pgconn_string(), autocommit=True) as conn:
with pg.connect(conninfo=get_project_pgconn_string(), autocommit=True) as conn:
with conn.cursor(row_factory=dict_row) as cur:
row = cur.execute(f"select current_database()").fetchone()
if row != None:
app/services/globals.py
-1
View File
@@ -23,7 +23,6 @@ non_realtime_region_patterns = {} # 基于source_outflow_region进行区分
realtime_region_pipe_flow_and_demand_id = {} # 基于source_outflow_region搜索该分区中的实时pipe_flow和demand的api_query_id，后续用region的流量 - 实时流量计的流量
realtime_region_pipe_flow_and_demand_patterns = {} # 基于source_outflow_region搜索该分区中的实时pipe_flow和demand的associated_pattern，后续用region的流量 - 实时流量计的流量
# ---------------------------------------------------------
# influxdb_api.py中的全局变量
# 全局变量，用于存储不同类型的realtime api_query_id
reservoir_liquid_level_realtime_ids = []
tank_liquid_level_realtime_ids = []
app/services/network_import.py
+6 -7
View File
@@ -5,8 +5,7 @@ import chardet
import psycopg
from psycopg import sql
import app.services.project_info as project_info
from app.core.config import get_pgconn_string
from app.infra.db.project_routing import get_project_pgconn_string
from app.services.tjnetwork import read_inp
@@ -15,13 +14,14 @@ from app.services.tjnetwork import read_inp
############################################################
def network_update(file_path: str) -> None:
def network_update(file_path: str, project_code: str) -> None:
"""
更新pg数据库中的inp文件
:param file_path: inp文件
:param project_code: 元数据项目代码
:return:
"""
read_inp("szh", file_path)
read_inp(project_code, file_path)
csv_path = "./history_pattern_flow.csv"
@@ -51,8 +51,7 @@ def network_update(file_path: str) -> None:
if os.path.exists(csv_path):
print(f"history_patterns_flows文件存在，开始处理...")
# 连接到 PostgreSQL 数据库（这里是数据库 "bb"）
with psycopg.connect(f"dbname={project_info.name} host=127.0.0.1") as conn:
with psycopg.connect(get_project_pgconn_string(project_code)) as conn:
with conn.cursor() as cur:
with open(csv_path, newline="", encoding="utf-8-sig") as csvfile:
reader = csv.DictReader(csvfile)
@@ -92,7 +91,7 @@ def submit_scada_info(name: str, coord_id: str) -> None:
print(f"检测到的文件编码：{file_encoding}")
try:
# 动态替换数据库名称
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
# 连接到 PostgreSQL 数据库（这里是数据库 "bb"）
with psycopg.connect(conn_string) as conn:
app/services/scheme_management.py
+15 -15
View File
@@ -7,7 +7,7 @@ import pandas as pd
import psycopg
from sqlalchemy import create_engine
from app.core.config import get_pgconn_string
from app.infra.db.project_routing import get_project_pgconn_string
from app.services.time_api import parse_utc_time
@@ -20,7 +20,7 @@ def scheme_name_exists(name: str, scheme_name: str) -> bool:
:return: 如果存在返回 True，否则返回 False
"""
try:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
cur.execute(
@@ -57,7 +57,7 @@ def store_scheme_info(
:return:
"""
try:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
sql = """
@@ -93,7 +93,7 @@ def delete_scheme_info(name: str, scheme_name: str) -> None:
:param scheme_name: 要删除的方案名称
"""
try:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
# 使用参数化查询删除方案记录
@@ -121,7 +121,7 @@ def query_scheme_list(
"""
try:
# 动态替换数据库名称
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
# 连接到 PostgreSQL 数据库（这里是数据库 "bb"）
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -203,7 +203,7 @@ def query_scheme_detail(
scheme_type,
)
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
if scheme_type:
@@ -255,7 +255,7 @@ def store_leakage_identify_result(
run_status: str = "completed",
error_message: str | None = None,
) -> None:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
cur.execute(
@@ -299,7 +299,7 @@ def query_leakage_identify_schemes(
scheme_type: str = "dma_leak_identification",
query_date: date | None = None,
) -> list[dict]:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
if query_date is None:
@@ -343,7 +343,7 @@ def query_leakage_identify_schemes(
def query_leakage_identify_scheme_detail(name: str, scheme_name: str) -> dict:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
cur.execute(
@@ -400,7 +400,7 @@ def query_burst_location_schemes(
scheme_type: str = "burst_location",
query_date: date | None = None,
) -> list[dict]:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
if query_date is None:
@@ -444,7 +444,7 @@ def query_burst_location_schemes(
def query_burst_location_scheme_detail(name: str, scheme_name: str) -> dict:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
cur.execute(
@@ -479,7 +479,7 @@ def query_burst_detection_schemes(
scheme_type: str = "burst_detection",
query_date: date | None = None,
) -> list[dict]:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
if query_date is None:
@@ -523,7 +523,7 @@ def query_burst_detection_schemes(
def query_burst_detection_scheme_detail(name: str, scheme_name: str) -> dict:
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
cur.execute(
@@ -564,7 +564,7 @@ def upload_shp_to_pg(name: str, table_name: str, role: str, shp_file_path: str):
"""
try:
# 动态连接到指定的数据库
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
with psycopg.connect(conn_string) as conn:
# 读取 Shapefile 文件
gdf = gpd.read_file(shp_file_path)
@@ -604,7 +604,7 @@ def submit_risk_probability_result(name: str, result_file_path: str) -> None:
try:
# 动态替换数据库名称
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
# 连接到 PostgreSQL 数据库
with psycopg.connect(conn_string) as conn:
app/services/simulation.py
+10 -11
View File
@@ -28,14 +28,13 @@ import pytz
import requests
import time
from typing import Optional, Tuple
import app.infra.db.influxdb.api as influxdb_api
import typing
import psycopg
import logging
import app.services.globals as globals
import app.services.project_info as project_info
from app.services.time_api import parse_beijing_time, parse_clock_duration_seconds
from app.core.config import get_pgconn_string
from app.infra.db.project_routing import get_project_pgconn_string
from app.infra.db.timescaledb.internal_queries import (
InternalQueries as TimescaleInternalQueries,
)
@@ -55,7 +54,7 @@ def query_corresponding_element_id_and_query_id(name: str) -> None:
:return:
"""
# 连接数据库
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -100,7 +99,7 @@ def query_corresponding_pattern_id_and_query_id(name: str) -> None:
:return:
"""
# 连接数据库
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -140,7 +139,7 @@ def query_non_realtime_region(name: str) -> dict:
"""
source_outflow_regions = [] # 用于存储所有 region（包含重复的）
# 构建连接字符串
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
# 连接到数据库
with psycopg.connect(conn_string) as conn:
@@ -217,7 +216,7 @@ def query_non_realtime_region_patterns(
region_tuple_to_key = {
frozenset(ids): region for region, ids in globals.source_outflow_region.items()
}
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -305,7 +304,7 @@ def query_realtime_region_pipe_flow_and_demand_id(
region_tuple_to_key = {
frozenset(ids): region for region, ids in globals.source_outflow_region.items()
}
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -377,7 +376,7 @@ def query_pipe_flow_region_patterns(
:param column_prefix: 需要提取的列的前缀
:return: pipe_flow_region_patterns 字典
"""
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -440,7 +439,7 @@ def query_SCADA_ID_corresponding_info(name: str, SCADA_ID: str) -> dict:
:param SCADA_ID: SCADA设备的ID
:return: 包含associated_element_id和api_query_id的字典
"""
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
# 使用 psycopg.connect 创建连接
with psycopg.connect(conn_string) as conn:
@@ -496,7 +495,7 @@ def get_source_outflow_region_id(
"No associated_source_outflow_id found in source_outflow_region."
)
return globals.source_outflow_region_id
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur:
@@ -553,7 +552,7 @@ def get_realtime_region_patterns(
globals.realtime_region_pipe_flow_and_demand_patterns = {
region: [] for region in globals.realtime_region_pipe_flow_and_demand_id.keys()
}
conn_string = get_pgconn_string(db_name=name)
conn_string = get_project_pgconn_string(db_name=name)
try:
with psycopg.connect(conn_string) as conn:
with conn.cursor() as cur: