From e47c61c121ed34d161c945075e9c6c03f7ec655d Mon Sep 17 00:00:00 2001 From: Huarch Date: Mon, 9 Nov 2026 18:49:32 +0800 Subject: [PATCH] refactor: reuse routing and remove unused analysis helpers --- app/mcp/scope.py | 14 +++++- app/services/burst_location.py | 19 ------- app/services/dma_leakage_estimation.py | 68 -------------------------- scripts/run_operational_worker.py | 30 ++++++------ tests/unit/test_mcp_scope.py | 52 ++++++++++++++++++++ tests/unit/test_operational_worker.py | 27 ++++++++++ 6 files changed, 107 insertions(+), 103 deletions(-) create mode 100644 tests/unit/test_mcp_scope.py diff --git a/app/mcp/scope.py b/app/mcp/scope.py index acf302c..69a6617 100644 --- a/app/mcp/scope.py +++ b/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 diff --git a/app/services/burst_location.py b/app/services/burst_location.py index 29f05e5..6eac442 100644 --- a/app/services/burst_location.py +++ b/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, diff --git a/app/services/dma_leakage_estimation.py b/app/services/dma_leakage_estimation.py index b350212..1802663 100644 --- a/app/services/dma_leakage_estimation.py +++ b/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], diff --git a/scripts/run_operational_worker.py b/scripts/run_operational_worker.py index 99b70f5..dbb3ee0 100644 --- a/scripts/run_operational_worker.py +++ b/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} diff --git a/tests/unit/test_mcp_scope.py b/tests/unit/test_mcp_scope.py new file mode 100644 index 0000000..0fb4ae0 --- /dev/null +++ b/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)] diff --git a/tests/unit/test_operational_worker.py b/tests/unit/test_operational_worker.py index 8a5640c..9e14365 100644 --- a/tests/unit/test_operational_worker.py +++ b/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"