refactor: reuse routing and remove unused analysis helpers
This commit is contained in:
+12
-2
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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],
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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)]
|
||||||
@@ -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"
|
||||||
|
|||||||
Reference in New Issue
Block a user