refactor: reuse routing and remove unused analysis helpers

app/mcp/scope.py

@@ -3,6 +3,8 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Any, AsyncIterator
+from uuid import UUID
+from sqlalchemy.ext.asyncio import AsyncSession
 from psycopg import AsyncConnection
 from app.auth.project_dependencies import ProjectContext
 from app.infra.db.project_routing import ActiveProjectRouting
@@ -18,16 +20,24 @@ from app.auth.project_dependencies import (
     get_project_timescale_connection,
 )
 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
 
 
 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)
         # Existing connection generators obtain routing before awaiting a pool.
         # End that implicit SELECT transaction here, before the network wait.
         await self.session.commit()
+        self._routing[key] = routing
         return routing
 
 

app/services/burst_location.py

@@ -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(
     *,
     start_value: datetime | str | None,

app/services/dma_leakage_estimation.py

@@ -188,74 +188,6 @@ def _build_area_map_by_topology(
     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(
     network: str,
     sensor_nodes: list[str],

scripts/run_operational_worker.py

@@ -90,25 +90,27 @@ class CredentialProvider:
         self.settings = settings
         self._token = ""
         self._expires = 0.0
+        self._refresh_lock = asyncio.Lock()
 
     async def headers(self, client: httpx.AsyncClient) -> dict[str, str]:
         if self.settings.bearer:
             token = self.settings.bearer.removeprefix("Bearer ")
         else:
-            if time.monotonic() >= self._expires:
-                response = await client.post(self.settings.token_url, data={
-                    "grant_type": "client_credentials", "client_id": self.settings.client_id,
-                    "client_secret": self.settings.client_secret,
-                })
-                response.raise_for_status()
-                payload = response.json()
-                self._token = payload["access_token"]
-                if not isinstance(self._token, str) or not self._token:
-                    raise ValueError("identity provider returned no access token")
-                # Product auth also restricts token age, independently of the
-                # provider's exp. Do not cache a long-lived OAuth token for hours.
-                lifetime = min(float(payload.get("expires_in", 60)), self.settings.token_refresh_seconds)
-                self._expires = time.monotonic() + max(1, lifetime - 15)
+            async with self._refresh_lock:
+                if time.monotonic() >= self._expires:
+                    response = await client.post(self.settings.token_url, data={
+                        "grant_type": "client_credentials", "client_id": self.settings.client_id,
+                        "client_secret": self.settings.client_secret,
+                    })
+                    response.raise_for_status()
+                    payload = response.json()
+                    self._token = payload["access_token"]
+                    if not isinstance(self._token, str) or not self._token:
+                        raise ValueError("identity provider returned no access token")
+                    # Product auth also restricts token age, independently of the
+                    # provider's exp. Do not cache a long-lived OAuth token for hours.
+                    lifetime = min(float(payload.get("expires_in", 60)), self.settings.token_refresh_seconds)
+                    self._expires = time.monotonic() + max(1, lifetime - 15)
             token = self._token
         return {"Authorization": f"Bearer {token}", "X-Project-Id": self.settings.project_id}
 

tests/unit/test_mcp_scope.py

@@ -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)]

tests/unit/test_operational_worker.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 from uuid import uuid4
@@ -8,6 +9,32 @@ import pytest
 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
 def anyio_backend():
     return "asyncio"