Tools reference
This commit is contained in:
Chen, Li
@@ -13,16 +13,25 @@ class DatabaseConfig(BaseSettings):
|
||||
port: int = 3306
|
||||
user: str = "root"
|
||||
password: str = ""
|
||||
database: str = ""
|
||||
pool_min: int = 1
|
||||
pool_max: int = 10
|
||||
schema: str = ""
|
||||
connect_timeout: int = 30
|
||||
|
||||
pool_size: int = 5
|
||||
max_overflow: int = 10
|
||||
pool_timeout: int = 30
|
||||
pool_recycle: int = 1800
|
||||
pool_pre_ping: bool = True
|
||||
# Physical connections opened eagerly right after an engine is created
|
||||
warmup_size: int = 1
|
||||
max_total_connections: int = 20
|
||||
health_check_interval: int = 30
|
||||
|
||||
|
||||
class QueryConfig(BaseSettings):
|
||||
model_config = SettingsConfigDict(env_prefix="QUERY_", env_file=_ENV_FILE, extra="ignore")
|
||||
|
||||
max_rows: int = 1000
|
||||
max_bytes: int = 4_000_000
|
||||
timeout_seconds: int = 30
|
||||
allow_write: bool = False
|
||||
|
||||
|
||||
+543
-84
@@ -1,110 +1,569 @@
|
||||
"""Async MySQL connection manager with connection pool.
|
||||
|
||||
Design mirrors DorisConnectionManager: a single shared pool, acquired via
|
||||
an async context manager, with an explicit ``close()`` for clean shutdown.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
from typing import Any, Sequence
|
||||
|
||||
import aiomysql
|
||||
from aiomysql import Pool
|
||||
from sqlalchemy import URL, text
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
|
||||
|
||||
from mini_db_mcp.config import Config, DatabaseConfig
|
||||
from mini_db_mcp.security import PermissionPolicy, SqlGuard, SqlValidationError, returns_rows
|
||||
|
||||
from .config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Only these SQL keywords are permitted when write access is disabled
|
||||
_READ_ONLY_KEYWORDS = frozenset(
|
||||
{"SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "WITH"}
|
||||
)
|
||||
_FIRST_KEYWORD_RE = re.compile(r"--[^\n]*|/\*.*?\*/", re.DOTALL)
|
||||
|
||||
HEALTH_CHECK_SQL = "SELECT 1"
|
||||
GLOBAL_PRINCIPAL = "__global__"
|
||||
DEFAULT_SESSION = "__anonymous__"
|
||||
|
||||
|
||||
def _first_sql_keyword(sql: str) -> str:
|
||||
stripped = _FIRST_KEYWORD_RE.sub("", sql).strip()
|
||||
return stripped.split(None, 1)[0].upper() if stripped else ""
|
||||
class CredentialKind(str, Enum):
|
||||
GLOBAL = "global"
|
||||
|
||||
|
||||
class PoolUnavailableError(RuntimeError):
|
||||
"""Raised when the configured host cannot serve the request."""
|
||||
|
||||
|
||||
class StaleConnectionError(RuntimeError):
|
||||
"""Raised when a leased connection outlives its pool generation."""
|
||||
|
||||
|
||||
class CredentialError(RuntimeError):
|
||||
"""Base for credential resolution failures.
|
||||
|
||||
Deliberately not an OSError subclass: the leasing path turns OSError into a
|
||||
pool invalidation, which must not be triggered by an auth problem.
|
||||
"""
|
||||
|
||||
|
||||
class UnknownSessionError(CredentialError):
|
||||
"""Raised when a session has no credential bound to it."""
|
||||
|
||||
|
||||
class CredentialExpiredError(CredentialError):
|
||||
"""Raised when the resolved credential is past its expiry."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PoolKey:
|
||||
kind: CredentialKind
|
||||
principal: str
|
||||
database: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.kind.value}:{self.principal}:{self.database or '-'}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Credential:
|
||||
"""Identity used to build and isolate a pool."""
|
||||
|
||||
kind: CredentialKind
|
||||
principal: str
|
||||
username: str
|
||||
secret: str
|
||||
database: str | None = None
|
||||
expires_at: float | None = None
|
||||
|
||||
@property
|
||||
def pool_key(self) -> PoolKey:
|
||||
return PoolKey(self.kind, self.principal, self.database or "")
|
||||
|
||||
def is_expired(self, now: float | None = None) -> bool:
|
||||
return self.expires_at is not None and (now or time.time()) >= self.expires_at
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineEntry:
|
||||
key: PoolKey
|
||||
engine: AsyncEngine
|
||||
generation: int
|
||||
secret_fingerprint: str
|
||||
created_at: float
|
||||
last_used_at: float
|
||||
last_healthy_at: float
|
||||
expires_at: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
columns: list[str]
|
||||
rows: list[dict[str, Any]]
|
||||
row_count: int
|
||||
truncated: bool = False
|
||||
session_id: str = DEFAULT_SESSION
|
||||
elapsed_ms: float = 0.0
|
||||
|
||||
|
||||
def bind_parameters(
|
||||
sql: str, params: Mapping[str, Any] | Sequence[Any] | None
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""Normalise ``?``/``%s`` placeholders to the named binds SQLAlchemy wants."""
|
||||
if params is None:
|
||||
return sql, {}
|
||||
if isinstance(params, Mapping):
|
||||
return sql, dict(params)
|
||||
|
||||
values = tuple(params)
|
||||
out: list[str] = []
|
||||
bound: dict[str, Any] = {}
|
||||
index = 0
|
||||
cursor = 0
|
||||
end = len(sql)
|
||||
while cursor < end:
|
||||
char = sql[cursor]
|
||||
if char in "'\"`":
|
||||
cursor = _copy_quoted(sql, cursor, out)
|
||||
continue
|
||||
if sql.startswith(("--", "#"), cursor):
|
||||
stop = sql.find("\n", cursor)
|
||||
stop = end if stop < 0 else stop
|
||||
out.append(sql[cursor:stop])
|
||||
cursor = stop
|
||||
continue
|
||||
if sql.startswith("/*", cursor):
|
||||
stop = sql.find("*/", cursor)
|
||||
stop = end if stop < 0 else stop + 2
|
||||
out.append(sql[cursor:stop])
|
||||
cursor = stop
|
||||
continue
|
||||
if char == "?" or sql.startswith("%s", cursor):
|
||||
if index >= len(values):
|
||||
raise SqlValidationError("More placeholders than parameters.")
|
||||
name = f"p{index}"
|
||||
bound[name] = values[index]
|
||||
out.append(f":{name}")
|
||||
index += 1
|
||||
cursor += 1 if char == "?" else 2
|
||||
continue
|
||||
out.append(char)
|
||||
cursor += 1
|
||||
|
||||
if index != len(values):
|
||||
raise SqlValidationError(
|
||||
f"Expected {index} parameters, received {len(values)}."
|
||||
)
|
||||
return "".join(out), bound
|
||||
|
||||
|
||||
def _copy_quoted(sql: str, start: int, out: list[str]) -> int:
|
||||
quote = sql[start]
|
||||
out.append(quote)
|
||||
cursor = start + 1
|
||||
end = len(sql)
|
||||
while cursor < end:
|
||||
char = sql[cursor]
|
||||
if char == "\\" and quote != "`":
|
||||
out.append(sql[cursor : cursor + 2])
|
||||
cursor += 2
|
||||
continue
|
||||
out.append(char)
|
||||
cursor += 1
|
||||
if char == quote:
|
||||
if cursor < end and sql[cursor] == quote:
|
||||
out.append(quote)
|
||||
cursor += 1
|
||||
continue
|
||||
break
|
||||
return cursor
|
||||
|
||||
|
||||
def _row_bytes(row: Mapping[str, Any]) -> int:
|
||||
return sum(len(str(value).encode("utf-8", "replace")) for value in row.values())
|
||||
|
||||
|
||||
def _fingerprint(secret: str) -> str:
|
||||
return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""Holds the aiomysql connection pool and enforces write policy."""
|
||||
"""Owns one AsyncEngine per (credential kind, principal, database).
|
||||
|
||||
def __init__(self, config: Config) -> None:
|
||||
Every lease is tagged with the pool generation so a rebuilt pool can never
|
||||
hand back a connection from the previous generation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Config,
|
||||
*,
|
||||
guard: SqlGuard | None = None,
|
||||
policy: PermissionPolicy | None = None,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self._pool: Pool | None = None
|
||||
self.db: DatabaseConfig = config.database
|
||||
self.guard = guard or SqlGuard(
|
||||
allow_write=config.query.allow_write,
|
||||
max_rows=config.query.max_rows
|
||||
)
|
||||
self.policy = policy or PermissionPolicy(allow_write=config.query.allow_write)
|
||||
self._entries: dict[PoolKey, EngineEntry] = {}
|
||||
self._generations: dict[PoolKey, int] = {}
|
||||
self._locks: dict[PoolKey, asyncio.Lock] = {}
|
||||
self._sessions: dict[str, Credential] = {}
|
||||
self._registry_lock = asyncio.Lock()
|
||||
self._semaphore = asyncio.Semaphore(self.db.max_total_connections)
|
||||
self._health_task: asyncio.Task[None] | None = None
|
||||
self._closed = False
|
||||
|
||||
|
||||
def global_credential(self) -> Credential:
|
||||
return Credential(
|
||||
kind=CredentialKind.GLOBAL,
|
||||
principal=GLOBAL_PRINCIPAL,
|
||||
username=self.db.user,
|
||||
secret=self.db.password,
|
||||
database=self.db.schema or None,
|
||||
)
|
||||
|
||||
|
||||
async def start(self) -> None:
|
||||
db = self.config.database
|
||||
self._pool = await aiomysql.create_pool(
|
||||
host=db.host,
|
||||
port=db.port,
|
||||
user=db.user,
|
||||
password=db.password,
|
||||
db=db.database or None,
|
||||
minsize=db.pool_min,
|
||||
maxsize=db.pool_max,
|
||||
connect_timeout=db.connect_timeout,
|
||||
autocommit=True,
|
||||
charset="utf8mb4",
|
||||
)
|
||||
logger.info(
|
||||
"Connection pool started: %s:%d db=%s",
|
||||
db.host,
|
||||
db.port,
|
||||
db.database or "(none)",
|
||||
)
|
||||
await self._get_entry(self.global_credential())
|
||||
if self.db.health_check_interval > 0:
|
||||
self._health_task = asyncio.create_task(self._health_loop())
|
||||
|
||||
|
||||
def bind_session(self, session_id: str, credential: Credential) -> None:
|
||||
self._sessions[session_id] = credential
|
||||
|
||||
|
||||
def unbind_session(self, session_id: str) -> None:
|
||||
self._sessions.pop(session_id, None)
|
||||
|
||||
|
||||
def _effective_credential(
|
||||
self, session_id: str, auth_context: Credential | None
|
||||
) -> Credential:
|
||||
credential = auth_context or self._sessions.get(session_id)
|
||||
if credential is None:
|
||||
raise UnknownSessionError(
|
||||
f"Session {session_id} is not bound to a credential."
|
||||
)
|
||||
if credential.is_expired():
|
||||
if self._sessions.get(session_id) is credential:
|
||||
self.unbind_session(session_id)
|
||||
raise CredentialExpiredError(
|
||||
f"Credential for session {session_id} has expired."
|
||||
)
|
||||
return credential
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._pool is not None:
|
||||
self._pool.close()
|
||||
await self._pool.wait_closed()
|
||||
self._pool = None
|
||||
logger.info("Connection pool closed")
|
||||
self._closed = True
|
||||
if self._health_task is not None:
|
||||
self._health_task.cancel()
|
||||
try:
|
||||
await self._health_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._health_task = None
|
||||
async with self._registry_lock:
|
||||
entries = list(self._entries.values())
|
||||
self._entries.clear()
|
||||
for entry in entries:
|
||||
await self._dispose(entry)
|
||||
|
||||
async def __aenter__(self) -> "ConnectionManager":
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_exc: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
def _build_url(self, credential: Credential) -> URL:
|
||||
return URL.create(
|
||||
"mysql+aiomysql",
|
||||
username=credential.username,
|
||||
password=credential.secret,
|
||||
host=self.db.url,
|
||||
port=self.db.port,
|
||||
database=credential.database or None,
|
||||
)
|
||||
|
||||
def _create_engine(self, credential: Credential) -> AsyncEngine:
|
||||
return create_async_engine(
|
||||
self._build_url(credential),
|
||||
pool_size=self.db.pool_size,
|
||||
max_overflow=self.db.max_overflow,
|
||||
pool_timeout=self.db.pool_timeout,
|
||||
pool_recycle=self.db.pool_recycle,
|
||||
pool_pre_ping=self.db.pool_pre_ping,
|
||||
# Lazy by design: pool_size is capacity, physical connections open on demand.
|
||||
connect_args={"connect_timeout": self.db.connect_timeout},
|
||||
)
|
||||
|
||||
async def _build_entry(self, credential: Credential) -> EngineEntry:
|
||||
key = credential.pool_key
|
||||
engine = self._create_engine(credential)
|
||||
try:
|
||||
async with engine.connect() as conn:
|
||||
await conn.execute(text(HEALTH_CHECK_SQL))
|
||||
except (SQLAlchemyError, OSError, asyncio.TimeoutError) as exc:
|
||||
await engine.dispose()
|
||||
raise PoolUnavailableError(
|
||||
f"Database {self.db.url} did not accept connections for {key}: {exc}"
|
||||
) from exc
|
||||
|
||||
generation = self._generations.get(key, 0) + 1
|
||||
self._generations[key] = generation
|
||||
now = time.time()
|
||||
logger.info("Pool ready key=%s node=%s generation=%s", key, self.db.url, generation)
|
||||
return EngineEntry(
|
||||
key=key,
|
||||
engine=engine,
|
||||
generation=generation,
|
||||
secret_fingerprint=_fingerprint(credential.secret),
|
||||
created_at=now,
|
||||
last_used_at=now,
|
||||
last_healthy_at=now,
|
||||
expires_at=credential.expires_at,
|
||||
)
|
||||
|
||||
async def _warmup(self, entry: EngineEntry) -> None:
|
||||
count = min(self.db.warmup_size, self.db.pool_size)
|
||||
if count <= 0:
|
||||
return
|
||||
|
||||
async def _open() -> None:
|
||||
async with entry.engine.connect() as conn:
|
||||
await conn.execute(text(HEALTH_CHECK_SQL))
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(_open() for _ in range(count)), return_exceptions=True
|
||||
)
|
||||
for result in results:
|
||||
if isinstance(result, BaseException):
|
||||
logger.warning("Warmup connection failed for %s: %s", entry.key, result)
|
||||
|
||||
def _lock_for(self, key: PoolKey) -> asyncio.Lock:
|
||||
lock = self._locks.get(key)
|
||||
if lock is None:
|
||||
lock = self._locks.setdefault(key, asyncio.Lock())
|
||||
return lock
|
||||
|
||||
async def _get_entry(self, credential: Credential) -> EngineEntry:
|
||||
if self._closed:
|
||||
raise PoolUnavailableError("ConnectionManager is closed.")
|
||||
if credential.is_expired():
|
||||
raise CredentialExpiredError("Credential has expired.")
|
||||
|
||||
key = credential.pool_key
|
||||
fingerprint = _fingerprint(credential.secret)
|
||||
entry = self._entries.get(key)
|
||||
if entry is not None and entry.secret_fingerprint == fingerprint:
|
||||
entry.last_used_at = time.time()
|
||||
return entry
|
||||
|
||||
async with self._lock_for(key):
|
||||
entry = self._entries.get(key)
|
||||
if entry is not None:
|
||||
if entry.secret_fingerprint == fingerprint:
|
||||
entry.last_used_at = time.time()
|
||||
return entry
|
||||
# The secret rotated behind an unchanged pool key.
|
||||
logger.info("Rebuilding pool key=%s after secret rotation", key)
|
||||
await self.invalidate(key)
|
||||
entry = await self._build_entry(credential)
|
||||
await self._warmup(entry)
|
||||
self._entries[key] = entry
|
||||
return entry
|
||||
|
||||
async def _dispose(self, entry: EngineEntry) -> None:
|
||||
try:
|
||||
await entry.engine.dispose()
|
||||
except Exception as exc: # pragma: no cover - dispose is best effort
|
||||
logger.warning("Failed to dispose pool %s: %s", entry.key, exc)
|
||||
|
||||
async def invalidate(self, key: PoolKey) -> None:
|
||||
"""Bump the generation and drop the pool; the next lease rebuilds it."""
|
||||
async with self._registry_lock:
|
||||
entry = self._entries.pop(key, None)
|
||||
if entry is None:
|
||||
return
|
||||
self._generations[key] = entry.generation + 1
|
||||
logger.info("Invalidated pool key=%s generation=%s", key, entry.generation)
|
||||
await self._dispose(entry)
|
||||
|
||||
def _assert_write_allowed(self, sql: str) -> None:
|
||||
keyword = _first_sql_keyword(sql)
|
||||
if not self.config.query.allow_write and keyword not in _READ_ONLY_KEYWORDS:
|
||||
raise PermissionError(
|
||||
f"Write operations are disabled (first keyword: {keyword!r}). "
|
||||
"Set QUERY_ALLOW_WRITE=true to enable DML/DDL."
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[aiomysql.Connection]:
|
||||
if self._pool is None:
|
||||
raise RuntimeError("Connection pool is not started")
|
||||
async with self._pool.acquire() as conn:
|
||||
yield conn
|
||||
async def connect(self, credential: Credential) -> AsyncGenerator[AsyncConnection]:
|
||||
entry = await self._get_entry(credential)
|
||||
generation = entry.generation
|
||||
async with self._semaphore:
|
||||
try:
|
||||
async with entry.engine.connect() as conn:
|
||||
if self._generations.get(entry.key) != generation:
|
||||
raise StaleConnectionError(
|
||||
f"Connection for {entry.key} belongs to generation "
|
||||
f"{generation}, current is {self._generations.get(entry.key)}."
|
||||
)
|
||||
entry.last_used_at = time.time()
|
||||
yield conn
|
||||
except (SQLAlchemyError, OSError) as exc:
|
||||
await self.invalidate(entry.key)
|
||||
raise PoolUnavailableError(
|
||||
f"Query failed on {self.db.url}: {exc}"
|
||||
) from exc
|
||||
|
||||
@staticmethod
|
||||
async def _stream_rows(
|
||||
conn: AsyncConnection,
|
||||
statement: str,
|
||||
bound: dict[str, Any],
|
||||
row_cap: int,
|
||||
byte_cap: int,
|
||||
) -> tuple[list[str], list[dict[str, Any]], bool]:
|
||||
"""Read through a server-side cursor, stopping as soon as a budget is hit."""
|
||||
result = await conn.stream(text(statement), bound)
|
||||
columns = list(result.keys())
|
||||
rows: list[dict[str, Any]] = []
|
||||
truncated = False
|
||||
total = 0
|
||||
try:
|
||||
async for row in result:
|
||||
if len(rows) >= row_cap:
|
||||
truncated = True
|
||||
break
|
||||
item = dict(zip(columns, row))
|
||||
total += _row_bytes(item)
|
||||
if total > byte_cap:
|
||||
truncated = True
|
||||
break
|
||||
rows.append(item)
|
||||
finally:
|
||||
await result.close()
|
||||
return columns, rows, truncated
|
||||
|
||||
|
||||
async def execute_query(
|
||||
self,
|
||||
session_id: str,
|
||||
sql: str,
|
||||
params: Mapping[str, Any] | Sequence[Any] | None = None,
|
||||
auth_context: Credential | None = None,
|
||||
*,
|
||||
max_rows: int | None = None,
|
||||
max_bytes: int | None = None,
|
||||
) -> QueryResult:
|
||||
"""Validate, route and run one statement; the lease is always released."""
|
||||
credential = self._effective_credential(session_id, auth_context)
|
||||
limits = self.config.query
|
||||
row_cap = min(max_rows or limits.max_rows, limits.max_rows)
|
||||
byte_cap = min(max_bytes or limits.max_bytes, limits.max_bytes)
|
||||
|
||||
statement, bound = bind_parameters(sql, params)
|
||||
guard = (
|
||||
self.guard
|
||||
if row_cap == self.guard.max_rows
|
||||
else replace(self.guard, max_rows=row_cap)
|
||||
)
|
||||
statement = guard.validate(statement)
|
||||
self.policy.check(statement, credential.database or self.db.schema)
|
||||
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
async with self.connect(credential) as conn:
|
||||
async with asyncio.timeout(limits.timeout_seconds):
|
||||
if not returns_rows(statement):
|
||||
result = await conn.execute(text(statement), bound)
|
||||
await conn.commit()
|
||||
return QueryResult(
|
||||
columns=[],
|
||||
rows=[],
|
||||
row_count=result.rowcount,
|
||||
session_id=session_id,
|
||||
elapsed_ms=(time.perf_counter() - started) * 1000,
|
||||
)
|
||||
columns, rows, truncated = await self._stream_rows(
|
||||
conn, statement, bound, row_cap, byte_cap
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Query failed session=%s principal=%s: %s",
|
||||
session_id,
|
||||
credential.principal,
|
||||
exc,
|
||||
)
|
||||
raise
|
||||
|
||||
return QueryResult(
|
||||
columns=columns,
|
||||
rows=rows,
|
||||
row_count=len(rows),
|
||||
truncated=truncated,
|
||||
session_id=session_id,
|
||||
elapsed_ms=(time.perf_counter() - started) * 1000,
|
||||
)
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
sql: str,
|
||||
args: tuple[Any, ...] | None = None,
|
||||
params: Mapping[str, Any] | Sequence[Any] | None = None,
|
||||
*,
|
||||
database: str | None = None,
|
||||
max_rows: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Execute *sql* and return a structured result dict."""
|
||||
self._assert_write_allowed(sql)
|
||||
limit = max_rows if max_rows is not None else self.config.query.max_rows
|
||||
async with self.acquire() as conn:
|
||||
if database:
|
||||
await conn.select_db(database)
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await cur.execute(sql, args)
|
||||
rows = await cur.fetchmany(limit + 1)
|
||||
truncated = len(rows) > limit
|
||||
if truncated:
|
||||
rows = rows[:limit]
|
||||
return {
|
||||
"columns": [d[0] for d in cur.description] if cur.description else [],
|
||||
"rows": [dict(r) for r in rows],
|
||||
"row_count": len(rows),
|
||||
"truncated": truncated,
|
||||
"max_rows": limit,
|
||||
}
|
||||
credential: Credential | None = None,
|
||||
) -> QueryResult:
|
||||
return await self.execute_query(
|
||||
DEFAULT_SESSION, sql, params, credential or self.global_credential()
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------ health checks
|
||||
|
||||
async def health_check(self, entry: EngineEntry) -> bool:
|
||||
try:
|
||||
async with entry.engine.connect() as conn:
|
||||
await asyncio.wait_for(
|
||||
conn.execute(text(HEALTH_CHECK_SQL)),
|
||||
timeout=self.db.connect_timeout,
|
||||
)
|
||||
except (SQLAlchemyError, OSError, asyncio.TimeoutError) as exc:
|
||||
logger.warning("Health check failed key=%s: %s", entry.key, exc)
|
||||
return False
|
||||
entry.last_healthy_at = time.time()
|
||||
return True
|
||||
|
||||
async def _health_loop(self) -> None:
|
||||
while not self._closed:
|
||||
try:
|
||||
await asyncio.sleep(self.db.health_check_interval)
|
||||
await self._sweep()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # pragma: no cover - loop must survive
|
||||
logger.exception("Health sweep failed: %s", exc)
|
||||
|
||||
|
||||
async def _sweep(self) -> None:
|
||||
now = time.time()
|
||||
for session_id in [
|
||||
sid for sid, cred in self._sessions.items() if cred.is_expired(now)
|
||||
]:
|
||||
self.unbind_session(session_id)
|
||||
|
||||
async with self._registry_lock:
|
||||
entries = list(self._entries.values())
|
||||
|
||||
for entry in entries:
|
||||
if entry.expires_at is not None and now >= entry.expires_at:
|
||||
reason = "credential expired"
|
||||
elif not await self.health_check(entry):
|
||||
reason = "unhealthy"
|
||||
else:
|
||||
continue
|
||||
logger.info("Recycling pool key=%s reason=%s", entry.key, reason)
|
||||
await self.invalidate(entry.key)
|
||||
|
||||
# Keep the global pool hot instead of waiting for the next request.
|
||||
global_credential = self.global_credential()
|
||||
if global_credential.pool_key not in self._entries and not self._closed:
|
||||
try:
|
||||
await self._get_entry(global_credential)
|
||||
except PoolUnavailableError as exc:
|
||||
logger.error("Global pool rebuild failed: %s", exc)
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import URL, inspect
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
|
||||
from mini_db_mcp.config import Config
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self, settings: Config):
|
||||
self.settings = settings
|
||||
self.database = settings.database
|
||||
url = URL.create(
|
||||
"mysql+aiomysql",
|
||||
username=self.database.user,
|
||||
password=self.database.password,
|
||||
host=self.database.host,
|
||||
port=self.database.port,
|
||||
database=self.database.schema or None,
|
||||
)
|
||||
self.engine: AsyncEngine = create_async_engine(
|
||||
url,
|
||||
pool_size=self.database.pool_size,
|
||||
max_overflow=self.database.max_overflow,
|
||||
pool_timeout=self.database.pool_timeout,
|
||||
pool_recycle=self.database.pool_recycle,
|
||||
pool_pre_ping=self.database.pool_pre_ping,
|
||||
connect_args={"connect_timeout": self.database.connect_timeout},
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.engine.dispose()
|
||||
|
||||
|
||||
async def list_tables(
|
||||
self,
|
||||
*,
|
||||
keyword: str | None,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> tuple[list[dict[str, Any]], str | None]:
|
||||
"""Return one name-sorted page of tables and views.
|
||||
|
||||
``keyword`` filters object names case-insensitively. ``cursor`` must be
|
||||
the last object name returned by the previous page. The returned cursor
|
||||
is ``None`` when no additional objects remain.
|
||||
|
||||
Raises:
|
||||
ValueError: If the cursor is not present in the filtered result.
|
||||
"""
|
||||
schema = self.settings.database.schema or None
|
||||
|
||||
def inspect_tables(sync_connection):
|
||||
inspector = inspect(sync_connection)
|
||||
table_names = set(inspector.get_table_names(schema=schema))
|
||||
try:
|
||||
view_names = set(inspector.get_view_names(schema=schema))
|
||||
except NotImplementedError:
|
||||
view_names = set()
|
||||
return table_names, view_names
|
||||
|
||||
async with self.engine.connect() as connection:
|
||||
table_names, view_names = await connection.run_sync(inspect_tables)
|
||||
|
||||
keyword_cf = keyword.casefold() if keyword else None
|
||||
names = sorted(
|
||||
name
|
||||
for name in table_names | view_names
|
||||
if keyword_cf is None or keyword_cf in name.casefold()
|
||||
)
|
||||
|
||||
start = 0
|
||||
if cursor is not None:
|
||||
try:
|
||||
start = names.index(cursor) + 1
|
||||
except ValueError as exc:
|
||||
raise ValueError("invalid pagination cursor") from exc
|
||||
|
||||
page_names = names[start : start + limit]
|
||||
next_cursor = page_names[-1] if start + len(page_names) < len(names) else None
|
||||
|
||||
objects: list[dict[str, Any]] = []
|
||||
for name in page_names:
|
||||
object_type = "view" if name in view_names else "table"
|
||||
objects.append(
|
||||
{
|
||||
"schema_name": schema,
|
||||
"object_name": name,
|
||||
"object_type": object_type,
|
||||
"description": None,
|
||||
}
|
||||
)
|
||||
return objects, next_cursor
|
||||
@@ -0,0 +1,171 @@
|
||||
from typing import Annotated, Any
|
||||
from pydantic import Field
|
||||
|
||||
from mini_db_mcp.domain.domain_model import DomainBaseModel, DomainDefinition, ChildToolDefinition
|
||||
|
||||
|
||||
class DomainCatalog(DomainBaseModel):
|
||||
"""Catalog for top-domains and child tools."""
|
||||
|
||||
domains: Annotated[tuple[DomainDefinition, ...], Field(min_length=1, max_length=12)]
|
||||
|
||||
|
||||
def _input_schema(
|
||||
properties: dict[str, Any],
|
||||
*,
|
||||
required: tuple[str, ...] = (),
|
||||
any_of: tuple[dict[str, Any], ...] = (),
|
||||
one_of: tuple[dict[str, Any], ...] = (),
|
||||
) -> dict[str, Any]:
|
||||
schema: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"additionalProperties": False,
|
||||
}
|
||||
if required:
|
||||
schema["required"] = list(required)
|
||||
if any_of:
|
||||
schema["anyOf"] = list(any_of)
|
||||
if one_of:
|
||||
schema["oneOf"] = list(one_of)
|
||||
return schema
|
||||
|
||||
|
||||
|
||||
_DETAIL_OUTPUT = {"type": "object", "additionalProperties": True}
|
||||
|
||||
_LIST_TABLES_OUTPUT = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"schema_name": {"type": ["string", "null"]},
|
||||
"object_name": {"type": "string"},
|
||||
"object_type": {"type": "string", "enum": ["table", "view"]},
|
||||
"description": {"type": ["string", "null"]},
|
||||
},
|
||||
"required": [
|
||||
"schema_name",
|
||||
"object_name",
|
||||
"object_type",
|
||||
"description",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"next_cursor": {
|
||||
"type": ["string", "null"],
|
||||
"description": "Cursor for the next page, or null when no more objects remain.",
|
||||
},
|
||||
"truncated": {"type": "boolean"},
|
||||
},
|
||||
"required": ["items", "next_cursor", "truncated"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
def _domain(
|
||||
*,
|
||||
name: str,
|
||||
title: str,
|
||||
description: str,
|
||||
children: tuple[ChildToolDefinition, ...],
|
||||
) -> DomainDefinition:
|
||||
return DomainDefinition(
|
||||
domain_name=name,
|
||||
domain_title=title,
|
||||
domain_description=description,
|
||||
children=children
|
||||
)
|
||||
|
||||
|
||||
def _child(
|
||||
*,
|
||||
domain: str,
|
||||
name: str,
|
||||
title: str,
|
||||
description: str,
|
||||
input_schema: dict[str, Any],
|
||||
output_schema: dict[str, Any] = _DETAIL_OUTPUT,
|
||||
) -> ChildToolDefinition:
|
||||
return ChildToolDefinition(
|
||||
tool_name=name,
|
||||
tool_title=title,
|
||||
tool_description=description,
|
||||
input_schema=input_schema,
|
||||
output_schema=output_schema,
|
||||
)
|
||||
|
||||
|
||||
def _string(
|
||||
description: str,
|
||||
*,
|
||||
enum: tuple[str, ...] | None = None,
|
||||
pattern: str | None = None,
|
||||
max_length: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
schema: dict[str, Any] = {"type": "string", "description": description}
|
||||
if enum is not None:
|
||||
schema["enum"] = list(enum)
|
||||
if pattern is not None:
|
||||
schema["pattern"] = pattern
|
||||
if max_length is not None:
|
||||
schema["maxLength"] = max_length
|
||||
return schema
|
||||
|
||||
|
||||
def _string_array(
|
||||
description: str,
|
||||
*,
|
||||
enum: tuple[str, ...] | None = None,
|
||||
min_items: int = 1,
|
||||
) -> dict[str, Any]:
|
||||
item_schema: dict[str, Any] = {"type": "string"}
|
||||
if enum is not None:
|
||||
item_schema["enum"] = list(enum)
|
||||
return {
|
||||
"type": "array",
|
||||
"description": description,
|
||||
"items": item_schema,
|
||||
"minItems": min_items,
|
||||
"uniqueItems": True,
|
||||
}
|
||||
|
||||
|
||||
DOMAIN_DEFINITIONS = (
|
||||
_domain(
|
||||
name="metadata",
|
||||
title="Metadata",
|
||||
description="Provides semantic metadata and a catalog of available business entities.",
|
||||
children=(
|
||||
_child(
|
||||
domain="metadata",
|
||||
name="list_tables",
|
||||
title="List Tables",
|
||||
description="Lists authorized database tables and views available for querying.",
|
||||
input_schema=_input_schema(
|
||||
{
|
||||
"keyword": _string("Optional case-insensitive object-name filter."),
|
||||
"cursor": {
|
||||
"type": ["string", "null"],
|
||||
"description": "Cursor returned by the previous page; omit for the first page.",
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of matching objects to return.",
|
||||
"minimum": 1,
|
||||
"default": 100,
|
||||
},
|
||||
},
|
||||
),
|
||||
output_schema=_LIST_TABLES_OUTPUT,
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
DOMAIN_CATALOG = DomainCatalog(domains=DOMAIN_DEFINITIONS)
|
||||
@@ -0,0 +1,72 @@
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
ToolHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
|
||||
|
||||
|
||||
class HandlerRegistryError(ValueError):
|
||||
"""Raised when handler providers contain an invalid binding."""
|
||||
|
||||
|
||||
class HandlerArgumentsError(ValueError):
|
||||
"""Raised when a handler cannot represent validated domain arguments."""
|
||||
|
||||
|
||||
class HandlerProvider(Protocol):
|
||||
"""Expose the child tools owned by one domain."""
|
||||
|
||||
domain_name: str
|
||||
|
||||
@property
|
||||
def handler_names(self) -> frozenset[str]: ...
|
||||
|
||||
def resolve_handler(self, name: str) -> ToolHandler | None: ...
|
||||
|
||||
|
||||
class DomainHandlerRegistry:
|
||||
"""Resolve a domain first, then one of its child tools."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
providers: Iterable[HandlerProvider] = (),
|
||||
) -> None:
|
||||
providers_by_domain: dict[str, HandlerProvider] = {}
|
||||
self._providers = tuple(providers)
|
||||
for provider in self._providers:
|
||||
if provider.domain_name in providers_by_domain:
|
||||
raise HandlerRegistryError(
|
||||
f"Duplicate domain handler: {provider.domain_name}"
|
||||
)
|
||||
for name in provider.handler_names:
|
||||
handler = provider.resolve_handler(name)
|
||||
if handler is None:
|
||||
raise HandlerRegistryError(
|
||||
f"Domain {provider.domain_name} declares missing handler: {name}"
|
||||
)
|
||||
providers_by_domain[provider.domain_name] = provider
|
||||
self._providers_by_domain = providers_by_domain
|
||||
|
||||
def resolve_domain(self, domain_name: str) -> HandlerProvider | None:
|
||||
return self._providers_by_domain.get(domain_name)
|
||||
|
||||
def resolve_handler(
|
||||
self,
|
||||
domain_name: str,
|
||||
child_tool: str,
|
||||
) -> ToolHandler | None:
|
||||
provider = self.resolve_domain(domain_name)
|
||||
return provider.resolve_handler(child_tool) if provider is not None else None
|
||||
|
||||
def has_handler(self, domain_name: str, child_tool: str) -> bool:
|
||||
"""Return whether the current runtime can resolve one handler."""
|
||||
return self.resolve_handler(domain_name, child_tool) is not None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ToolHandler",
|
||||
"HandlerRegistryError",
|
||||
"HandlerArgumentsError",
|
||||
"HandlerProvider",
|
||||
"DomainHandlerRegistry"
|
||||
]
|
||||
@@ -0,0 +1,15 @@
|
||||
from mini_db_mcp.domain.domain_handler import DomainHandlerRegistry
|
||||
from mini_db_mcp.domain.handler.metadata_handler import MetadataHandler
|
||||
from mini_db_mcp.domain.runtime.metadata_runtime import MetadataRuntime, MetadataStore
|
||||
|
||||
|
||||
def create_domain_registry(*,
|
||||
metadata_store: MetadataStore
|
||||
) -> DomainHandlerRegistry:
|
||||
"""Build the domain runtimes, handlers, and registry."""
|
||||
metadata_runtime = MetadataRuntime(metadata_store)
|
||||
return DomainHandlerRegistry(
|
||||
[
|
||||
MetadataHandler(metadata_runtime)
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,167 @@
|
||||
from typing import Any
|
||||
from mcp.types import Tool, ToolAnnotations
|
||||
|
||||
from mini_db_mcp.domain.domain_catalog import DomainCatalog
|
||||
from mini_db_mcp.domain.domain_model import DomainDefinition, DomainErrorCode
|
||||
|
||||
|
||||
DOMAIN_DISCOVERY_DESCRIPTION_SUFFIX = (
|
||||
" Choose one matching domain; do not speculatively open others. Call with "
|
||||
"an empty object to discover its exact child tools and schemas, then call "
|
||||
"it again with the returned child_tool and child_arguments."
|
||||
)
|
||||
|
||||
|
||||
_DOMAIN_MANIFEST_OUTPUT_SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"mode": {"type": "string", "const": "manifest"},
|
||||
"domain": {"type": "string"},
|
||||
"generated_at": {"type": "string", "format": "date-time"},
|
||||
"children": {
|
||||
"type": "array",
|
||||
"items": {"type": "object"},
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"mode",
|
||||
"domain",
|
||||
"generated_at",
|
||||
"children",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
_DOMAIN_EXECUTION_OUTPUT_SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"mode": {"type": "string", "const": "result"},
|
||||
"domain": {"type": "string"},
|
||||
"child_tool": {"type": "string"},
|
||||
"data": {"type": "object"},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"request_id": {"type": "string"},
|
||||
"duration_ms": {"type": "number", "minimum": 0},
|
||||
"source": {"type": "string"},
|
||||
"truncated": {"type": "boolean"},
|
||||
},
|
||||
"required": [
|
||||
"request_id",
|
||||
"duration_ms",
|
||||
"source",
|
||||
"truncated",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"warnings": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"mode",
|
||||
"domain",
|
||||
"child_tool",
|
||||
"data",
|
||||
"metadata",
|
||||
"warnings",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
DOMAIN_ERROR_OUTPUT_SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"mode": {"type": "string", "const": "error"},
|
||||
"domain": {"type": "string", "minLength": 1},
|
||||
"child_tool": {"type": "string", "minLength": 1},
|
||||
"error": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"enum": [code.value for code in DomainErrorCode],
|
||||
},
|
||||
"message": {"type": "string", "minLength": 1},
|
||||
"retryable": {"type": "boolean"},
|
||||
"details": {
|
||||
"type": "object",
|
||||
"additionalProperties": True,
|
||||
},
|
||||
},
|
||||
"required": ["code", "message", "retryable", "details"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"required": ["mode", "domain", "child_tool", "error"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
DOMAIN_TOOL_INPUT_SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"child_tool": {
|
||||
"type": "string",
|
||||
"minLength": 1,
|
||||
"description": (
|
||||
"Exact child tool name returned by this domain manifest."
|
||||
),
|
||||
},
|
||||
"child_arguments": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Arguments validated against the selected child tool schema."
|
||||
),
|
||||
},
|
||||
},
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
DOMAIN_TOOL_OUTPUT_SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"oneOf": [
|
||||
_DOMAIN_MANIFEST_OUTPUT_SCHEMA,
|
||||
_DOMAIN_EXECUTION_OUTPUT_SCHEMA,
|
||||
DOMAIN_ERROR_OUTPUT_SCHEMA,
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class DomainManifest():
|
||||
"""Expose top-level domain."""
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
catalog: DomainCatalog
|
||||
) -> None:
|
||||
self._catalog = catalog
|
||||
self._domain_names = frozenset(
|
||||
domain.domain_name for domain in self._catalog.domains
|
||||
)
|
||||
self._domain_tools = tuple(
|
||||
_build_top_level_tool(domain) for domain in self._catalog.domains
|
||||
)
|
||||
|
||||
|
||||
async def list_tools(self) -> list[Tool]:
|
||||
return list(self._domain_tools)
|
||||
|
||||
|
||||
def _build_top_level_tool(domain: DomainDefinition) -> Tool:
|
||||
return Tool(
|
||||
name=domain.domain_name,
|
||||
title=domain.domain_title,
|
||||
description=domain.domain_description + DOMAIN_DISCOVERY_DESCRIPTION_SUFFIX,
|
||||
input_schema=DOMAIN_TOOL_INPUT_SCHEMA,
|
||||
output_schema=DOMAIN_TOOL_OUTPUT_SCHEMA,
|
||||
annotations=ToolAnnotations(
|
||||
title=domain.domain_title,
|
||||
read_only_hint=True
|
||||
),
|
||||
)
|
||||
@@ -1,13 +1,55 @@
|
||||
from enum import StrEnum
|
||||
import json
|
||||
from typing import Annotated, Any, Literal, Self, cast, TypedDict
|
||||
|
||||
from typing import Annotated, Any, Literal, Self, cast
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
JsonValue
|
||||
JsonValue,
|
||||
StringConstraints,
|
||||
Field,
|
||||
model_validator,
|
||||
field_validator,
|
||||
field_serializer
|
||||
)
|
||||
|
||||
JsonObject = dict[str, JsonValue]
|
||||
from mini_db_mcp.utils.json_util import JsonObject, to_json_value
|
||||
|
||||
|
||||
MAX_DOMAIN_MANIFEST_BYTES = 16 * 1024
|
||||
|
||||
|
||||
Identity = Annotated[
|
||||
str,
|
||||
StringConstraints(
|
||||
min_length=1,
|
||||
max_length=128,
|
||||
pattern=r"^[a-z][a-z0-9_]*$",
|
||||
strip_whitespace=True,
|
||||
),
|
||||
]
|
||||
|
||||
NonEmptyText = Annotated[
|
||||
str,
|
||||
StringConstraints(
|
||||
min_length=1,
|
||||
max_length=128,
|
||||
strip_whitespace=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class DomainErrorCode(StrEnum):
|
||||
"""Stable errors returned by hierarchical domain tools."""
|
||||
|
||||
DOMAIN_DISCOVERY_DENIED = "DOMAIN_DISCOVERY_DENIED"
|
||||
CHILD_TOOL_NOT_FOUND = "CHILD_TOOL_NOT_FOUND"
|
||||
CHILD_MANIFEST_STALE = "CHILD_MANIFEST_STALE"
|
||||
CHILD_ARGUMENTS_INVALID = "CHILD_ARGUMENTS_INVALID"
|
||||
CHILD_RESULT_INVALID = "CHILD_RESULT_INVALID"
|
||||
CHILD_CAPABILITY_UNAVAILABLE = "CHILD_CAPABILITY_UNAVAILABLE"
|
||||
CHILD_CONFIRMATION_REQUIRED = "CHILD_CONFIRMATION_REQUIRED"
|
||||
CHILD_EXECUTION_TIMEOUT = "CHILD_EXECUTION_TIMEOUT"
|
||||
CHILD_EXECUTION_FAILED = "CHILD_EXECUTION_FAILED"
|
||||
|
||||
|
||||
class DomainBaseModel(BaseModel):
|
||||
@@ -20,8 +62,9 @@ class DomainBaseModel(BaseModel):
|
||||
validate_default=True,
|
||||
)
|
||||
|
||||
def to_wire(self) -> JsonObject:
|
||||
def to_json_object(self) -> JsonObject:
|
||||
"""Return a JSON-compatible representation with protocol aliases."""
|
||||
|
||||
return cast(
|
||||
JsonObject,
|
||||
self.model_dump(
|
||||
@@ -33,10 +76,95 @@ class DomainBaseModel(BaseModel):
|
||||
|
||||
def to_canonical_json(self) -> str:
|
||||
"""Serialize deterministically for manifests, tests, and hashing."""
|
||||
|
||||
return json.dumps(
|
||||
self.to_wire(),
|
||||
self.to_json_object(),
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
sort_keys=True,
|
||||
)
|
||||
|
||||
|
||||
class DomainDefinition(DomainBaseModel):
|
||||
"""Top-level domain definition."""
|
||||
|
||||
domain_name: Identity
|
||||
domain_title: NonEmptyText
|
||||
domain_description: NonEmptyText
|
||||
children: Annotated[
|
||||
tuple[ChildToolDefinition, ...],
|
||||
Field(min_length=1, max_length=12),
|
||||
]
|
||||
max_manifest_bytes: Annotated[
|
||||
int,
|
||||
Field(ge=1, le=MAX_DOMAIN_MANIFEST_BYTES),
|
||||
] = MAX_DOMAIN_MANIFEST_BYTES
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_domain(self) -> Self:
|
||||
_require_unique(
|
||||
tuple(child.tool_name for child in self.children),
|
||||
"child tool names",
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class ChildToolDefinition(DomainBaseModel):
|
||||
"""Exact child tool of a domain."""
|
||||
|
||||
tool_name: Identity
|
||||
tool_title: NonEmptyText
|
||||
tool_description: NonEmptyText
|
||||
input_schema: JsonObject
|
||||
output_schema: JsonObject
|
||||
|
||||
|
||||
class DomainToolRequest(DomainBaseModel):
|
||||
"""Common request shape for discovery and exact child execution."""
|
||||
|
||||
child_tool: Identity | None = None
|
||||
child_arguments: JsonObject | None = None
|
||||
|
||||
@field_validator("child_arguments", mode="before")
|
||||
@classmethod
|
||||
def _normalize_arguments_input(cls, value: Any) -> JsonValue:
|
||||
_reject_tuples(value)
|
||||
return to_json_value(value)
|
||||
|
||||
@field_serializer("child_arguments")
|
||||
def _serialize_arguments(self, value: JsonObject | None) -> JsonObject | None:
|
||||
return cast(JsonObject | None, to_json_value(value))
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_request_mode(self) -> Self:
|
||||
if self.child_tool is None and self.child_arguments is not None:
|
||||
raise ValueError("Arguments require child_tool")
|
||||
return self
|
||||
|
||||
@property
|
||||
def is_discovery(self) -> bool:
|
||||
return self.child_tool is None
|
||||
|
||||
@property
|
||||
def execution_arguments(self) -> JsonObject:
|
||||
return cast(
|
||||
JsonObject,
|
||||
to_json_value(self.child_arguments or {}),
|
||||
)
|
||||
|
||||
|
||||
def _require_unique(values: tuple[str, ...], field_name: str) -> None:
|
||||
if len(values) != len(set(values)):
|
||||
raise ValueError(f"{field_name} must not contain duplicates")
|
||||
|
||||
|
||||
def _reject_tuples(value: Any) -> None:
|
||||
if isinstance(value, tuple):
|
||||
raise ValueError("JSON arrays must be provided as lists")
|
||||
if isinstance(value, dict):
|
||||
for item in value.values():
|
||||
_reject_tuples(item)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
_reject_tuples(item)
|
||||
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Mapping
|
||||
from types import MappingProxyType
|
||||
|
||||
from mini_db_mcp.domain.domain_handler import ToolHandler
|
||||
from mini_db_mcp.domain.runtime.metadata_runtime import MetadataRuntime
|
||||
|
||||
|
||||
class MetadataHandler:
|
||||
|
||||
domain_name = "metadata"
|
||||
|
||||
|
||||
def __init__(self, runtime: MetadataRuntime) -> None:
|
||||
_handlers = {
|
||||
"list_tables": runtime.list_tables,
|
||||
}
|
||||
self._handlers: Mapping[str, ToolHandler] = MappingProxyType(_handlers)
|
||||
|
||||
|
||||
@property
|
||||
def handler_names(self) -> frozenset[str]:
|
||||
return frozenset(self._handlers)
|
||||
|
||||
def resolve_handler(self, name: str) -> ToolHandler | None:
|
||||
return self._handlers.get(name)
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
class MetadataStore(Protocol):
|
||||
"""Database capabilities required by the metadata runtime."""
|
||||
|
||||
async def list_tables(
|
||||
self,
|
||||
*,
|
||||
keyword: str | None,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> tuple[list[dict[str, Any]], str | None]: ...
|
||||
|
||||
|
||||
class MetadataRuntime:
|
||||
"""Runtime for metadata domain."""
|
||||
|
||||
def __init__(self, store: MetadataStore) -> None:
|
||||
self._store = store
|
||||
|
||||
async def list_tables(self, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
keyword = arguments.get("keyword")
|
||||
if keyword is not None and not isinstance(keyword, str):
|
||||
raise ValueError("keyword must be a string or null")
|
||||
|
||||
cursor = arguments.get("cursor")
|
||||
if cursor is not None and not isinstance(cursor, str):
|
||||
raise ValueError("cursor must be a string or null")
|
||||
|
||||
limit = arguments.get("limit", 100)
|
||||
if not isinstance(limit, int) or isinstance(limit, bool) or limit < 1:
|
||||
raise ValueError("limit must be a positive integer")
|
||||
|
||||
items, next_cursor = await self._store.list_tables(
|
||||
keyword=keyword,
|
||||
limit=limit,
|
||||
cursor=cursor,
|
||||
)
|
||||
return {
|
||||
"items": items,
|
||||
"next_cursor": next_cursor,
|
||||
"truncated": next_cursor is not None,
|
||||
}
|
||||
+7
-22
@@ -1,30 +1,13 @@
|
||||
import logging
|
||||
import sys
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
import asyncio
|
||||
|
||||
from mcp.server import Server
|
||||
from mini_db_mcp.config import Config
|
||||
from mini_db_mcp.server import MiniServer
|
||||
from mini_db_mcp import __version__
|
||||
|
||||
from mcp.types import (
|
||||
INTERNAL_ERROR,
|
||||
INVALID_PARAMS,
|
||||
CallToolResult,
|
||||
GetPromptResult,
|
||||
ListPromptsResult,
|
||||
ListResourcesResult,
|
||||
ListToolsResult,
|
||||
ReadResourceResult,
|
||||
TextContent,
|
||||
TextResourceContents,
|
||||
)
|
||||
|
||||
from .config import Config
|
||||
from .prompts import MissingPromptArgumentError, PromptsManager, UnknownPromptError
|
||||
from .server import MiniServer
|
||||
from . import __version__
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
def _setup_logging(level: str) -> None:
|
||||
logging.basicConfig(
|
||||
@@ -41,6 +24,8 @@ def main() -> None:
|
||||
transport = config.server.transport
|
||||
if transport == "stdio":
|
||||
asyncio.run(server.start_stdio())
|
||||
elif transport == "http":
|
||||
asyncio.run(server.start_http())
|
||||
else:
|
||||
logger.error("Unsupported transport: %s", transport)
|
||||
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
|
||||
import sqlglot
|
||||
from sqlglot import exp
|
||||
from sqlglot.errors import ParseError, TokenError
|
||||
|
||||
|
||||
class SecurityError(Exception):
|
||||
"""Base class for all query rejections."""
|
||||
|
||||
|
||||
class SqlValidationError(SecurityError):
|
||||
pass
|
||||
|
||||
|
||||
class PermissionDeniedError(SecurityError):
|
||||
pass
|
||||
|
||||
|
||||
MAX_SQL_LENGTH = 200_000
|
||||
# Doris first, MySQL as a fallback for constructs the Doris dialect misses.
|
||||
PARSE_DIALECTS = ("doris", "mysql")
|
||||
|
||||
# exp.Query covers SELECT, set operations and subqueries.
|
||||
_READ_ONLY_NODES = (exp.Query, exp.Show, exp.Describe, exp.Use)
|
||||
_WRITE_NODES = (
|
||||
exp.Insert,
|
||||
exp.Update,
|
||||
exp.Delete,
|
||||
exp.Drop,
|
||||
exp.Create,
|
||||
exp.Alter,
|
||||
exp.TruncateTable,
|
||||
)
|
||||
# Rejected regardless of allow_write: they read files, block workers or leak timing.
|
||||
_FORBIDDEN_FUNCTIONS = frozenset({"load_file", "sleep", "benchmark", "outfile"})
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def parse_statement(sql: str) -> exp.Expression:
|
||||
"""Parse a single statement, rejecting anything the dialects cannot model.
|
||||
|
||||
The returned tree is shared through the cache and must never be mutated.
|
||||
"""
|
||||
if not sql or not sql.strip():
|
||||
raise SqlValidationError("Empty SQL statement.")
|
||||
if len(sql) > MAX_SQL_LENGTH:
|
||||
raise SqlValidationError(f"Statement exceeds {MAX_SQL_LENGTH} characters.")
|
||||
if "\x00" in sql:
|
||||
raise SqlValidationError("Statement contains a NUL byte.")
|
||||
|
||||
last_error: Exception | None = None
|
||||
for dialect in PARSE_DIALECTS:
|
||||
try:
|
||||
trees = sqlglot.parse(sql, dialect=dialect)
|
||||
except (ParseError, TokenError, RecursionError) as exc:
|
||||
last_error = exc
|
||||
continue
|
||||
|
||||
statements = [tree for tree in trees if tree is not None]
|
||||
if not statements:
|
||||
raise SqlValidationError("Empty SQL statement.")
|
||||
if len(statements) > 1:
|
||||
raise SqlValidationError("Multiple statements per request are not allowed.")
|
||||
|
||||
root = statements[0]
|
||||
# sqlglot degrades DCL/ADMIN and other unmodelled syntax to Command.
|
||||
if isinstance(root, exp.Command) or root.find(exp.Command) is not None:
|
||||
raise SqlValidationError("Administrative or unsupported statement.")
|
||||
return root
|
||||
|
||||
raise SqlValidationError(f"Statement could not be parsed: {last_error}")
|
||||
|
||||
|
||||
def table_names(root: exp.Expression) -> set[str]:
|
||||
"""Physical tables referenced by the tree, with CTE aliases removed."""
|
||||
ctes = {cte.alias_or_name.lower() for cte in root.find_all(exp.CTE)}
|
||||
names: set[str] = set()
|
||||
for table in root.find_all(exp.Table):
|
||||
name = table.name.lower()
|
||||
if not name:
|
||||
continue
|
||||
parts = [part.lower() for part in (table.catalog, table.db) if part]
|
||||
if not parts and name in ctes:
|
||||
continue
|
||||
names.add(".".join([*parts, name]))
|
||||
return names
|
||||
|
||||
|
||||
def extract_tables(sql: str) -> set[str]:
|
||||
return table_names(parse_statement(sql))
|
||||
|
||||
|
||||
def returns_rows(sql: str) -> bool:
|
||||
return isinstance(parse_statement(sql), _READ_ONLY_NODES)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SqlGuard:
|
||||
"""Statement level validation"""
|
||||
|
||||
allow_write: bool = False
|
||||
max_rows: int = 1000
|
||||
|
||||
def validate(self, sql: str) -> str:
|
||||
statement = sql.strip()
|
||||
if statement.endswith(";"):
|
||||
statement = statement[:-1].strip()
|
||||
root = parse_statement(statement)
|
||||
self._reject_dangerous(root)
|
||||
|
||||
if isinstance(root, _WRITE_NODES):
|
||||
if not self.allow_write:
|
||||
raise PermissionDeniedError(
|
||||
f"{type(root).__name__.upper()} is rejected in read-only mode."
|
||||
)
|
||||
return statement
|
||||
if isinstance(root, _READ_ONLY_NODES):
|
||||
return self._apply_row_limit(statement, root)
|
||||
raise SqlValidationError(f"Unsupported statement type: {type(root).__name__}.")
|
||||
|
||||
def _reject_dangerous(self, root: exp.Expression) -> None:
|
||||
if root.find(exp.Into) is not None:
|
||||
raise SqlValidationError("SELECT ... INTO is not allowed.")
|
||||
for func in root.find_all(exp.Anonymous):
|
||||
if func.name.lower() in _FORBIDDEN_FUNCTIONS:
|
||||
raise SqlValidationError(f"Function '{func.name}' is not allowed.")
|
||||
|
||||
def _apply_row_limit(self, statement: str, root: exp.Expression) -> str:
|
||||
"""Append a limit to the original text; the AST is only used to decide."""
|
||||
if not isinstance(root, exp.Query) or root.args.get("limit") is not None:
|
||||
return statement
|
||||
return f"{statement} LIMIT {self.max_rows}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PermissionPolicy:
|
||||
"""Per-principal access rules, resolved before the statement runs."""
|
||||
|
||||
allowed_databases: frozenset[str] = frozenset()
|
||||
denied_tables: frozenset[str] = frozenset()
|
||||
allow_write: bool = False
|
||||
|
||||
def check(self, sql: str, default_database: str | None) -> None:
|
||||
root = parse_statement(sql.strip().rstrip(";").strip())
|
||||
if isinstance(root, _WRITE_NODES) and not self.allow_write:
|
||||
raise PermissionDeniedError("Principal is not allowed to write.")
|
||||
|
||||
for table in table_names(root):
|
||||
database, _, name = table.rpartition(".")
|
||||
database = database or (default_database or "").lower()
|
||||
if self.allowed_databases and database not in self.allowed_databases:
|
||||
raise PermissionDeniedError(f"Database '{database}' is not accessible.")
|
||||
qualified = f"{database}.{name}" if database else name
|
||||
if name in self.denied_tables or qualified in self.denied_tables:
|
||||
raise PermissionDeniedError(f"Table '{qualified}' is not accessible.")
|
||||
+89
-16
@@ -1,31 +1,36 @@
|
||||
from __future__ import annotations
|
||||
from . import __version__
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Protocol
|
||||
|
||||
from mcp.server import Server
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
GetPromptResult,
|
||||
ListPromptsResult,
|
||||
ListResourcesResult,
|
||||
ListToolsResult,
|
||||
ReadResourceResult,
|
||||
PaginatedRequestParams,
|
||||
TextContent,
|
||||
TextResourceContents,
|
||||
Resource,
|
||||
Tool,
|
||||
Prompt
|
||||
)
|
||||
|
||||
from .config import Config
|
||||
from .prompts import PromptsManager
|
||||
from . import __version__
|
||||
from mini_db_mcp.config import Config
|
||||
from mini_db_mcp.db import Database
|
||||
from mini_db_mcp.domain.domain_catalog import DOMAIN_CATALOG
|
||||
from mini_db_mcp.domain.domain_handler_factory import create_domain_registry
|
||||
from mini_db_mcp.utils.pagination import paginate_list
|
||||
from mini_db_mcp.prompts import PromptsManager
|
||||
from mini_db_mcp.tools import ToolsManager
|
||||
|
||||
|
||||
NAME = "mini-db-mcp"
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ResourcesProtocol(Protocol):
|
||||
async def list_resources(self) -> list[Resource]: ...
|
||||
|
||||
@@ -35,7 +40,7 @@ class ResourcesProtocol(Protocol):
|
||||
class ToolsProtocol(Protocol):
|
||||
async def list_tools(self) -> list[Tool]: ...
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> str: ...
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: ...
|
||||
|
||||
|
||||
class PromptsProtocol(Protocol):
|
||||
@@ -51,25 +56,61 @@ class MiniServer():
|
||||
def __init__(self, config: Config) -> None:
|
||||
self.config = config
|
||||
self.prompts_manager = PromptsManager()
|
||||
self.database = Database(config)
|
||||
self.handler_registry = create_domain_registry(
|
||||
metadata_store=self.database,
|
||||
)
|
||||
self.tools_manager = ToolsManager(
|
||||
catalog=DOMAIN_CATALOG,
|
||||
registry=self.handler_registry,
|
||||
)
|
||||
self.server = _create_mcp_server(
|
||||
prompts_manager=self.prompts_manager
|
||||
prompts_manager=self.prompts_manager,
|
||||
tools_manager=self.tools_manager,
|
||||
)
|
||||
|
||||
async def start_stdio(self) -> None:
|
||||
from mcp.server.stdio import stdio_server
|
||||
|
||||
async with stdio_server() as streams:
|
||||
read_stream, write_stream = streams
|
||||
await self.server.run(
|
||||
read_stream,
|
||||
write_stream,
|
||||
self.server.create_initialization_options(),
|
||||
try:
|
||||
async with stdio_server() as streams:
|
||||
read_stream, write_stream = streams
|
||||
await self.server.run(
|
||||
read_stream,
|
||||
write_stream,
|
||||
self.server.create_initialization_options(),
|
||||
)
|
||||
finally:
|
||||
await self.database.close()
|
||||
|
||||
async def start_http(self) -> None:
|
||||
"""Serve MCP over stateless Streamable HTTP at ``/mcp``."""
|
||||
import uvicorn
|
||||
|
||||
app = self.server.streamable_http_app(
|
||||
streamable_http_path="/mcp",
|
||||
json_response=True,
|
||||
stateless_http=True,
|
||||
host=self.config.server.host,
|
||||
)
|
||||
server = uvicorn.Server(
|
||||
uvicorn.Config(
|
||||
app,
|
||||
host=self.config.server.host,
|
||||
port=self.config.server.port,
|
||||
log_level=self.config.server.log_level.lower(),
|
||||
)
|
||||
)
|
||||
try:
|
||||
await server.serve()
|
||||
finally:
|
||||
await self.database.close()
|
||||
|
||||
|
||||
def _create_mcp_server(
|
||||
*,
|
||||
prompts_manager: PromptsProtocol | None = None,
|
||||
tools_manager: ToolsProtocol | None = None,
|
||||
) -> Server:
|
||||
|
||||
async def on_list_prompts(ctx, params):
|
||||
@@ -78,11 +119,43 @@ def _create_mcp_server(
|
||||
async def on_get_prompt(ctx, params):
|
||||
return await prompts_manager.get_prompt(params.name, dict(params.arguments or {}))
|
||||
|
||||
async def list_tools(
|
||||
ctx,
|
||||
params: PaginatedRequestParams | None,
|
||||
) -> ListToolsResult:
|
||||
tools = await tools_manager.list_tools()
|
||||
page_items, next_cursor = paginate_list(
|
||||
tools,
|
||||
cursor=params.cursor if params else None,
|
||||
page_size=10,
|
||||
identifier=lambda tool: tool.name,
|
||||
)
|
||||
logger.info("Returning %d of %d tools", len(page_items), len(tools))
|
||||
return ListToolsResult(
|
||||
tools=page_items,
|
||||
next_cursor=next_cursor,
|
||||
)
|
||||
|
||||
async def call_tool(ctx, params):
|
||||
# ToolsManager returns either a manifest (discovery) or a child result.
|
||||
result = await tools_manager.call_tool(
|
||||
params.name,
|
||||
dict(params.arguments or {}),
|
||||
)
|
||||
# Publish the same payload as structured data and backward-compatible text.
|
||||
return CallToolResult(
|
||||
content=[TextContent(type="text", text=json.dumps(result, ensure_ascii=False))],
|
||||
structuredContent=result,
|
||||
isError=result.get("mode") == "error",
|
||||
)
|
||||
|
||||
server = Server(
|
||||
name=NAME,
|
||||
version=__version__,
|
||||
on_list_prompts=on_list_prompts,
|
||||
on_get_prompt=on_get_prompt,
|
||||
on_list_tools=list_tools,
|
||||
on_call_tool=call_tool,
|
||||
)
|
||||
return server
|
||||
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from mcp.types import Tool
|
||||
from mini_db_mcp.domain.domain_manifest import DomainManifest
|
||||
from mini_db_mcp.domain.domain_catalog import DomainCatalog
|
||||
from mini_db_mcp.domain.domain_handler import DomainHandlerRegistry
|
||||
from mini_db_mcp.domain.domain_model import DomainErrorCode, DomainToolRequest
|
||||
|
||||
|
||||
class ToolsManager:
|
||||
"""Expose domain tools through a two-call discovery and execution flow.
|
||||
|
||||
The first call uses empty arguments to return the selected domain's child
|
||||
manifest. The second call supplies ``child_tool`` and ``child_arguments``
|
||||
to execute exactly one child through the domain registry.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
catalog: DomainCatalog,
|
||||
registry: DomainHandlerRegistry,
|
||||
) -> None:
|
||||
self._catalog = catalog
|
||||
self._domains = {domain.domain_name: domain for domain in catalog.domains}
|
||||
self._domain_manifest_service = DomainManifest(catalog=catalog)
|
||||
self._registry = registry
|
||||
|
||||
async def list_tools(self) -> list[Tool]:
|
||||
return await self._domain_manifest_service.list_tools()
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
domain = self._domains.get(name)
|
||||
if domain is None:
|
||||
raise ValueError(f"Unknown domain tool: {name}")
|
||||
|
||||
request = DomainToolRequest(**arguments)
|
||||
# First call: expose only this domain's child schemas to the model.
|
||||
if request.is_discovery:
|
||||
return {
|
||||
"mode": "manifest",
|
||||
"domain": name,
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
"children": [child.to_json_object() for child in domain.children],
|
||||
}
|
||||
|
||||
child_tool = request.child_tool
|
||||
assert child_tool is not None
|
||||
# Second call: route the selected child to its owning domain handler.
|
||||
handler = self._registry.resolve_handler(name, child_tool)
|
||||
if handler is None:
|
||||
return self._error_result(
|
||||
name,
|
||||
child_tool,
|
||||
DomainErrorCode.CHILD_TOOL_NOT_FOUND,
|
||||
f"Unknown child tool {child_tool!r} for domain {name!r}.",
|
||||
)
|
||||
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
data = await handler(dict(request.execution_arguments))
|
||||
except ValueError as exc:
|
||||
return self._error_result(
|
||||
name,
|
||||
child_tool,
|
||||
DomainErrorCode.CHILD_ARGUMENTS_INVALID,
|
||||
str(exc),
|
||||
)
|
||||
except Exception as exc:
|
||||
return self._error_result(
|
||||
name,
|
||||
child_tool,
|
||||
DomainErrorCode.CHILD_EXECUTION_FAILED,
|
||||
str(exc),
|
||||
retryable=True,
|
||||
)
|
||||
|
||||
return {
|
||||
"mode": "result",
|
||||
"domain": name,
|
||||
"child_tool": child_tool,
|
||||
"data": data,
|
||||
"metadata": {
|
||||
"request_id": uuid.uuid4().hex,
|
||||
"duration_ms": (time.perf_counter() - started_at) * 1000,
|
||||
"source": "database",
|
||||
"truncated": bool(data.get("truncated", False)),
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _error_result(
|
||||
domain: str,
|
||||
child_tool: str,
|
||||
code: DomainErrorCode,
|
||||
message: str,
|
||||
*,
|
||||
retryable: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"mode": "error",
|
||||
"domain": domain,
|
||||
"child_tool": child_tool,
|
||||
"error": {
|
||||
"code": code.value,
|
||||
"message": message,
|
||||
"retryable": retryable,
|
||||
"details": {},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
from collections.abc import Mapping
|
||||
from math import isfinite
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
type JsonObject = dict[str, JsonValue]
|
||||
|
||||
|
||||
def to_json_value(value: object) -> JsonValue:
|
||||
|
||||
if value is None or isinstance(value, str | bool | int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
if not isfinite(value):
|
||||
raise ValueError("JSON numbers must be finite")
|
||||
return value
|
||||
if isinstance(value, Mapping):
|
||||
result: JsonObject = {}
|
||||
for key, item in value.items():
|
||||
if not isinstance(key, str):
|
||||
raise TypeError(
|
||||
f"JSON object keys must be strings, got {type(key).__name__}"
|
||||
)
|
||||
result[key] = to_json_value(item)
|
||||
return result
|
||||
if isinstance(value, list | tuple):
|
||||
return [to_json_value(item) for item in value]
|
||||
raise TypeError(
|
||||
f"{type(value).__name__} cannot be converted to a JSON value"
|
||||
)
|
||||
@@ -0,0 +1,38 @@
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import TypeVar
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
MAX_PAGE_SIZE = 10
|
||||
|
||||
|
||||
def paginate_list(
|
||||
items: Sequence[T],
|
||||
*,
|
||||
cursor: str | None,
|
||||
page_size: int,
|
||||
identifier: Callable[[T], str],
|
||||
) -> tuple[list[T], str | None]:
|
||||
"""Return a deterministic page using the prior item's identifier as cursor."""
|
||||
|
||||
if not 1 <= page_size <= MAX_PAGE_SIZE:
|
||||
raise ValueError(f"page_size must be in the range 1-{MAX_PAGE_SIZE}")
|
||||
|
||||
ordered = sorted(items, key=identifier)
|
||||
identifiers = [identifier(item) for item in ordered]
|
||||
if len(identifiers) != len(set(identifiers)):
|
||||
raise ValueError("list contains duplicate identifiers")
|
||||
|
||||
start = 0
|
||||
if cursor is not None:
|
||||
try:
|
||||
start = identifiers.index(cursor) + 1
|
||||
except ValueError as exc:
|
||||
raise ValueError("invalid pagination cursor") from exc
|
||||
if start >= len(ordered):
|
||||
raise ValueError("invalid pagination cursor")
|
||||
|
||||
page_items = ordered[start : start + page_size]
|
||||
end = start + len(page_items)
|
||||
next_cursor = identifiers[end - 1] if end < len(ordered) else None
|
||||
return page_items, next_cursor
|
||||
Reference in New Issue
Block a user