87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
"""Resolve authoritative metadata and project routing for each MCP request."""
|
|
|
|
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
|
|
from fastapi import HTTPException
|
|
from mcp.server.auth.provider import AccessToken
|
|
|
|
from app.auth.metadata_dependencies import get_current_metadata_user
|
|
from app.auth.project_dependencies import (
|
|
resolve_project_context,
|
|
resolve_project_business_routing,
|
|
resolve_project_simulation_routing,
|
|
get_project_pg_connection,
|
|
get_project_timescale_connection,
|
|
)
|
|
from app.infra.db.metadb.database import get_metadata_session
|
|
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository, ProjectDbRouting
|
|
from app.infra.db.project_routing import activate_project_routing
|
|
|
|
|
|
class QueryMetadataRepository(MetadataRepository):
|
|
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
|
|
|
|
|
|
@dataclass
|
|
class QueryScope:
|
|
project: ProjectContext
|
|
user: Any
|
|
connection: AsyncConnection
|
|
routing: ActiveProjectRouting
|
|
metadata: QueryMetadataRepository
|
|
|
|
@asynccontextmanager
|
|
async def simulation_routing(self) -> AsyncIterator[ActiveProjectRouting]:
|
|
"""Resolve both databases before a native calculation can modify its model."""
|
|
routing = await resolve_project_simulation_routing(self.project, self.metadata)
|
|
await self.metadata.session.commit()
|
|
with activate_project_routing(routing):
|
|
yield routing
|
|
|
|
@asynccontextmanager
|
|
async def timescale(self):
|
|
# Metadata repository here is used only to resolve a known project DB.
|
|
async with asynccontextmanager(get_project_timescale_connection)(
|
|
self.project, self.metadata
|
|
) as conn:
|
|
await self.metadata.session.commit()
|
|
yield conn
|
|
|
|
|
|
@asynccontextmanager
|
|
async def project_scope(principal: AccessToken, project_id: str):
|
|
if not principal.claims:
|
|
raise HTTPException(401, "MCP identity is missing")
|
|
async with asynccontextmanager(get_metadata_session)() as session:
|
|
metadata = QueryMetadataRepository(session)
|
|
user = await get_current_metadata_user(principal.claims, metadata)
|
|
project = await resolve_project_context(project_id, user, metadata)
|
|
routing = await resolve_project_business_routing(project, metadata)
|
|
# End the metadata read before any external query/IO waits.
|
|
await session.commit()
|
|
async with asynccontextmanager(get_project_pg_connection)(
|
|
project, metadata
|
|
) as conn:
|
|
await session.commit()
|
|
with activate_project_routing(routing):
|
|
yield QueryScope(project, user, conn, routing, metadata)
|