refactor: reuse routing and remove unused analysis helpers
Generic Container CI/CD / test-build-publish (push) Successful in 4m3s
Server CI/CD v2 / build-test-publish-and-deploy (push) Successful in 4m3s

This commit is contained in:
2026-11-09 18:49:32 +08:00
parent 6a4478bfe4
commit e47c61c121
6 changed files with 107 additions and 103 deletions
+12 -2
View File
@@ -3,6 +3,8 @@
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, AsyncIterator from typing import Any, AsyncIterator
from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
from psycopg import AsyncConnection from psycopg import AsyncConnection
from app.auth.project_dependencies import ProjectContext from app.auth.project_dependencies import ProjectContext
from app.infra.db.project_routing import ActiveProjectRouting from app.infra.db.project_routing import ActiveProjectRouting
@@ -18,16 +20,24 @@ from app.auth.project_dependencies import (
get_project_timescale_connection, get_project_timescale_connection,
) )
from app.infra.db.metadb.database import get_metadata_session from app.infra.db.metadb.database import get_metadata_session
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository, ProjectDbRouting
from app.infra.db.project_routing import activate_project_routing from app.infra.db.project_routing import activate_project_routing
class QueryMetadataRepository(MetadataRepository): class QueryMetadataRepository(MetadataRepository):
async def get_project_db_routing(self, project_id, role): def __init__(self, session: AsyncSession):
super().__init__(session)
self._routing: dict[tuple[UUID, str], ProjectDbRouting | None] = {}
async def get_project_db_routing(self, project_id: UUID, role: str) -> ProjectDbRouting | None:
key = (project_id, role)
if key in self._routing:
return self._routing[key]
routing = await super().get_project_db_routing(project_id, role) routing = await super().get_project_db_routing(project_id, role)
# Existing connection generators obtain routing before awaiting a pool. # Existing connection generators obtain routing before awaiting a pool.
# End that implicit SELECT transaction here, before the network wait. # End that implicit SELECT transaction here, before the network wait.
await self.session.commit() await self.session.commit()
self._routing[key] = routing
return routing return routing
-19
View File
@@ -405,25 +405,6 @@ def _store_burst_scheme(
) )
def _validate_scada_windows(
*,
scada_burst_start: datetime | str | None,
scada_burst_end: datetime | str | None,
) -> tuple[datetime, datetime]:
values = [scada_burst_start, scada_burst_end]
if any(v is None for v in values):
raise ValueError(
"使用后端 SCADA 查询时,必须同时提供 scada_burst_start/scada_burst_end。"
)
burst_start_dt = _to_datetime(scada_burst_start)
burst_end_dt = _to_datetime(scada_burst_end)
if burst_start_dt >= burst_end_dt:
raise ValueError(
"爆管时段 SCADA 时间窗非法:scada_burst_start 必须早于 scada_burst_end。"
)
return burst_start_dt, burst_end_dt
def _validate_time_window( def _validate_time_window(
*, *,
start_value: datetime | str | None, start_value: datetime | str | None,
-68
View File
@@ -188,74 +188,6 @@ def _build_area_map_by_topology(
return area_map, areas, node_coords return area_map, areas, node_coords
def _build_area_node_map(area_map: dict[str, str]) -> dict[str, list[str]]:
area_node_map: dict[str, list[str]] = {}
for node_id, area_id in area_map.items():
area_node_map.setdefault(area_id, []).append(node_id)
for area_id in list(area_node_map.keys()):
area_node_map[area_id] = sorted(area_node_map[area_id])
return area_node_map
def _build_node_visual_payload(
area_map: dict[str, str],
node_coords: dict[str, dict[str, float]],
rows: list[dict[str, Any]],
) -> dict[str, Any]:
area_leakage_map = _build_area_leakage_map(rows)
max_leakage = max(area_leakage_map.values(), default=0.0)
features: list[dict[str, Any]] = []
for node_id, area_id in area_map.items():
coord = node_coords.get(node_id)
if not coord:
continue
leakage_flow = float(area_leakage_map.get(area_id, 0.0))
leakage_level = _classify_leakage_level(leakage_flow, max_leakage)
features.append(
{
"type": "Feature",
"properties": {
"node_id": node_id,
"area_id": area_id,
"leakage_flow_m3_per_s": leakage_flow,
"leakage_level": leakage_level,
},
"geometry": {
"type": "Point",
"coordinates": [float(coord["x"]), float(coord["y"])],
},
}
)
return {"type": "FeatureCollection", "features": features}
def _build_area_leakage_map(rows: list[dict[str, Any]]) -> dict[str, float]:
area_leakage_map: dict[str, float] = {}
for row in rows:
area_id = str(row.get("Area", "")).strip()
if not area_id:
continue
area_leakage_map[area_id] = float(row.get("LeakageFlow_m3_per_s", 0.0))
return area_leakage_map
def _classify_leakage_level(leakage_flow: float, max_leakage: float) -> str:
if max_leakage <= 0:
return "normal"
ratio = leakage_flow / max_leakage
if ratio >= 0.75:
return "high"
if ratio >= 0.4:
return "medium"
if ratio > 0:
return "low"
return "normal"
def _build_drawing_payload(node_visual_payload: dict[str, Any]) -> dict[str, Any]:
return node_visual_payload
def _build_observed_pressure_from_scada( def _build_observed_pressure_from_scada(
network: str, network: str,
sensor_nodes: list[str], sensor_nodes: list[str],
+2
View File
@@ -90,11 +90,13 @@ class CredentialProvider:
self.settings = settings self.settings = settings
self._token = "" self._token = ""
self._expires = 0.0 self._expires = 0.0
self._refresh_lock = asyncio.Lock()
async def headers(self, client: httpx.AsyncClient) -> dict[str, str]: async def headers(self, client: httpx.AsyncClient) -> dict[str, str]:
if self.settings.bearer: if self.settings.bearer:
token = self.settings.bearer.removeprefix("Bearer ") token = self.settings.bearer.removeprefix("Bearer ")
else: else:
async with self._refresh_lock:
if time.monotonic() >= self._expires: if time.monotonic() >= self._expires:
response = await client.post(self.settings.token_url, data={ response = await client.post(self.settings.token_url, data={
"grant_type": "client_credentials", "client_id": self.settings.client_id, "grant_type": "client_credentials", "client_id": self.settings.client_id,
+52
View File
@@ -0,0 +1,52 @@
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4
import pytest
from app.auth.project_dependencies import ProjectContext
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository, ProjectDbRouting
from app.infra.db.project_routing import get_active_project_routing
from app.mcp import scope
@pytest.mark.anyio
async def test_mcp_scope_reads_routing_once_and_refreshes_it_on_the_next_request(monkeypatch):
project = ProjectContext(uuid4(), "fixture", uuid4(), "member")
session = SimpleNamespace(commit=AsyncMock())
connections = []
current_dsn = ["dbname=first"]
async def metadata_session():
yield session
async def read_routing(_repo, project_id, role):
assert project_id == project.project_id and role == "biz_data"
return ProjectDbRouting(project_id, role, "postgresql", current_dsn[0], 1, 4)
@asynccontextmanager
async def connection(project_id, role, dsn, minimum, maximum):
assert session.commit.await_count > 0
connections.append((project_id, role, dsn, minimum, maximum))
yield object()
read = AsyncMock(side_effect=read_routing)
monkeypatch.setattr(scope, "get_metadata_session", metadata_session)
monkeypatch.setattr(scope, "get_current_metadata_user", AsyncMock(return_value=object()))
monkeypatch.setattr(scope, "resolve_project_context", AsyncMock(return_value=project))
# Patch at the repository boundary so both real routing resolvers are exercised.
async def repository_read(self, project_id, role):
return await read(self, project_id, role)
monkeypatch.setattr(MetadataRepository, "get_project_db_routing", repository_read)
monkeypatch.setattr("app.auth.project_dependencies.project_connection_manager.pg_connection", connection)
for dsn in ("dbname=first", "dbname=second"):
current_dsn[0] = dsn
async with scope.project_scope(SimpleNamespace(claims={"sub": "fixture"}), str(project.project_id)) as value:
assert value.routing.business_dsn == dsn
assert get_active_project_routing() is value.routing
assert get_active_project_routing() is None
assert read.await_count == 2
assert [row[2:] for row in connections] == [("dbname=first", 1, 4), ("dbname=second", 1, 4)]
+27
View File
@@ -1,3 +1,4 @@
import asyncio
import json import json
import os import os
from uuid import uuid4 from uuid import uuid4
@@ -8,6 +9,32 @@ import pytest
from scripts.run_operational_worker import CredentialProvider, OperationalWorker, WorkerSettings from scripts.run_operational_worker import CredentialProvider, OperationalWorker, WorkerSettings
@pytest.mark.anyio
async def test_parallel_worker_requests_share_cold_and_expired_token_refresh(monkeypatch):
clock = [100.0]
requests = []
monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
async def handler(request):
requests.append(request)
await asyncio.sleep(0)
return httpx.Response(200, json={"access_token": f"TEST_ONLY_{len(requests)}", "expires_in": 60})
provider = CredentialProvider(WorkerSettings(
"http://server.test", "http://agent.test", str(uuid4()),
token_url="https://identity.test/token", client_id="fixture", client_secret="TEST_ONLY",
))
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
first = await asyncio.gather(provider.headers(client), provider.headers(client))
assert first[0] == first[1]
assert len(requests) == 1
clock[0] = 146.0
second = await asyncio.gather(provider.headers(client), provider.headers(client))
assert second[0] == second[1]
assert second[0] != first[0]
assert len(requests) == 2
@pytest.fixture @pytest.fixture
def anyio_backend(): def anyio_backend():
return "asyncio" return "asyncio"