metadata tools
This commit is contained in:
Chen, Li
@@ -15,23 +15,19 @@ class DatabaseConfig(BaseSettings):
|
||||
password: str = ""
|
||||
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
|
||||
max_bytes: int = 262_144
|
||||
max_cell_chars: int = 4_096
|
||||
timeout_seconds: int = 30
|
||||
allow_write: bool = False
|
||||
|
||||
|
||||
@@ -1,569 +0,0 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
from typing import Any, Sequence
|
||||
|
||||
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
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
HEALTH_CHECK_SQL = "SELECT 1"
|
||||
GLOBAL_PRINCIPAL = "__global__"
|
||||
DEFAULT_SESSION = "__anonymous__"
|
||||
|
||||
|
||||
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:
|
||||
"""Owns one AsyncEngine per (credential kind, principal, database).
|
||||
|
||||
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.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:
|
||||
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:
|
||||
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)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
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,
|
||||
params: Mapping[str, Any] | Sequence[Any] | None = None,
|
||||
*,
|
||||
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)
|
||||
+155
-13
@@ -1,11 +1,19 @@
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import URL, inspect
|
||||
from datetime import datetime, date, time
|
||||
from decimal import Decimal
|
||||
from uuid import UUID
|
||||
import base64
|
||||
from sqlalchemy import URL, inspect, Table, MetaData, select
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from mini_db_mcp.config import Config
|
||||
from mini_db_mcp.utils.pagination_util import paginate_list
|
||||
|
||||
|
||||
class QueryExecutionError(RuntimeError):
|
||||
pass
|
||||
|
||||
class Database:
|
||||
def __init__(self, settings: Config):
|
||||
self.settings = settings
|
||||
@@ -63,24 +71,21 @@ class Database:
|
||||
table_names, view_names = await connection.run_sync(inspect_tables)
|
||||
|
||||
keyword_cf = keyword.casefold() if keyword else None
|
||||
names = sorted(
|
||||
sorted_table_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
|
||||
paged_table_names, next_cursor = paginate_list(
|
||||
sorted_table_names,
|
||||
cursor=cursor,
|
||||
page_size=limit,
|
||||
identifier=lambda name: name,
|
||||
)
|
||||
|
||||
objects: list[dict[str, Any]] = []
|
||||
for name in page_names:
|
||||
for name in paged_table_names:
|
||||
object_type = "view" if name in view_names else "table"
|
||||
objects.append(
|
||||
{
|
||||
@@ -91,3 +96,140 @@ class Database:
|
||||
}
|
||||
)
|
||||
return objects, next_cursor
|
||||
|
||||
|
||||
async def describe_table(
|
||||
self,
|
||||
*,
|
||||
table_name: str,
|
||||
schema: str | None = None,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Return one cursor-paginated page of a table's columns."""
|
||||
effective_schema = schema or self.settings.database.schema or None
|
||||
|
||||
def inspect_table(sync_connection):
|
||||
inspector = inspect(sync_connection)
|
||||
columns = inspector.get_columns(table_name, schema=effective_schema)
|
||||
primary_keys = set(
|
||||
(
|
||||
inspector.get_pk_constraint(
|
||||
table_name,
|
||||
schema=effective_schema,
|
||||
)
|
||||
or {}
|
||||
).get("constrained_columns") or []
|
||||
)
|
||||
try:
|
||||
table_description = (
|
||||
inspector.get_table_comment(
|
||||
table_name,
|
||||
schema=effective_schema,
|
||||
)
|
||||
or {}
|
||||
).get("text")
|
||||
except NotImplementedError:
|
||||
table_description = None
|
||||
return columns, primary_keys, table_description
|
||||
|
||||
try:
|
||||
async with self.engine.connect() as connection:
|
||||
columns, primary_keys, table_description = await connection.run_sync(
|
||||
inspect_table
|
||||
)
|
||||
except SQLAlchemyError as exc:
|
||||
raise QueryExecutionError(f"Failed to read table metadata: {exc}") from exc
|
||||
|
||||
paged_columns, next_cursor = paginate_list(
|
||||
columns,
|
||||
cursor=cursor,
|
||||
page_size=limit,
|
||||
identifier=lambda column: column["name"],
|
||||
)
|
||||
visible_columns: list[dict[str, Any]] = []
|
||||
for column in paged_columns:
|
||||
visible_columns.append(
|
||||
{
|
||||
"name": column["name"],
|
||||
"data_type": str(column.get("type")),
|
||||
"nullable": column.get("nullable"),
|
||||
"primary_key": column["name"] in primary_keys,
|
||||
"description": column.get("comment"),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"schema_name": effective_schema,
|
||||
"table_name": table_name,
|
||||
"table_description": table_description,
|
||||
"columns": visible_columns,
|
||||
"next_cursor": next_cursor,
|
||||
"truncated": next_cursor is not None,
|
||||
}
|
||||
|
||||
|
||||
async def sample_values(
|
||||
self,
|
||||
*,
|
||||
table_name: str,
|
||||
schema: str | None = None,
|
||||
column: str,
|
||||
limit: int
|
||||
) -> tuple[list[Any], bool]:
|
||||
effective_schema = schema or self.settings.database.schema or None
|
||||
metadata = MetaData()
|
||||
try:
|
||||
async with self.engine.connect() as connection:
|
||||
table_obj = await connection.run_sync(
|
||||
lambda sync_connection: Table(
|
||||
table_name,
|
||||
metadata,
|
||||
schema=effective_schema,
|
||||
autoload_with=sync_connection,
|
||||
)
|
||||
)
|
||||
except SQLAlchemyError as exc:
|
||||
raise QueryExecutionError("metadata_error", f"反射表结构失败：{exc}") from exc
|
||||
try:
|
||||
column_obj = table_obj.c[column]
|
||||
except KeyError as exc:
|
||||
raise QueryExecutionError(
|
||||
"metadata_error", f"列不存在：{table_name}.{column}"
|
||||
) from exc
|
||||
statement = (
|
||||
select(column_obj)
|
||||
.where(column_obj.is_not(None))
|
||||
.distinct()
|
||||
.limit(limit + 1)
|
||||
)
|
||||
try:
|
||||
async with self.engine.connect() as connection:
|
||||
result = await connection.execute(statement)
|
||||
values = [row[0] for row in result]
|
||||
except SQLAlchemyError as exc:
|
||||
raise QueryExecutionError("query_failed", f"采样查询失败：{exc}") from exc
|
||||
truncated = len(values) > limit
|
||||
return [
|
||||
_serialize_scalar(
|
||||
value,
|
||||
max_chars=self.settings.query.max_cell_chars,
|
||||
)
|
||||
for value in values[:limit]
|
||||
], truncated
|
||||
|
||||
|
||||
def _serialize_scalar(value: Any, *, max_chars: int) -> Any:
|
||||
if value is None or isinstance(value, (bool, int, float)):
|
||||
return value
|
||||
if isinstance(value, Decimal):
|
||||
return str(value)
|
||||
if isinstance(value, (datetime, date, time, UUID)):
|
||||
return value.isoformat() if hasattr(value, "isoformat") else str(value)
|
||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||
text = base64.b64encode(bytes(value)).decode("ascii")
|
||||
else:
|
||||
text = str(value)
|
||||
if len(text) > max_chars:
|
||||
return text[:max_chars] + "…[truncated]"
|
||||
return text
|
||||
@@ -66,6 +66,57 @@ _LIST_TABLES_OUTPUT = {
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
_DESCRIBE_TABLE_OUTPUT = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"schema_name": {"type": ["string", "null"]},
|
||||
"table_name": {"type": "string"},
|
||||
"table_description": {"type": ["string", "null"]},
|
||||
"columns": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"data_type": {"type": "string"},
|
||||
"nullable": {"type": ["boolean", "null"]},
|
||||
"primary_key": {"type": "boolean"},
|
||||
"description": {"type": ["string", "null"]},
|
||||
},
|
||||
"required": [
|
||||
"name",
|
||||
"data_type",
|
||||
"nullable",
|
||||
"primary_key",
|
||||
"description",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"next_cursor": {"type": ["string", "null"]},
|
||||
"truncated": {"type": "boolean"},
|
||||
},
|
||||
"required": [
|
||||
"schema_name",
|
||||
"table_name",
|
||||
"table_description",
|
||||
"columns",
|
||||
"next_cursor",
|
||||
"truncated",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
_SAMPLE_VALUES_OUTPUT = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"values": {"type": "array"},
|
||||
"truncated": {"type": "boolean"},
|
||||
},
|
||||
"required": ["values", "truncated"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
def _domain(
|
||||
*,
|
||||
@@ -163,6 +214,52 @@ DOMAIN_DEFINITIONS = (
|
||||
),
|
||||
output_schema=_LIST_TABLES_OUTPUT,
|
||||
),
|
||||
_child(
|
||||
domain="metadata",
|
||||
name="describe_table",
|
||||
title="Describe Table",
|
||||
description="Returns a cursor-paginated description of table columns.",
|
||||
input_schema=_input_schema(
|
||||
{
|
||||
"table_name": _string("Table or view name."),
|
||||
"schema": _string("Optional schema name."),
|
||||
"cursor": {
|
||||
"type": ["string", "null"],
|
||||
"description": "Last column name from the previous page.",
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of columns to return.",
|
||||
"minimum": 1,
|
||||
"default": 100,
|
||||
},
|
||||
},
|
||||
required=("table_name",),
|
||||
),
|
||||
output_schema=_DESCRIBE_TABLE_OUTPUT,
|
||||
),
|
||||
_child(
|
||||
domain="metadata",
|
||||
name="sample_values",
|
||||
title="Sample Values",
|
||||
description="Returns distinct non-null sample values for one column.",
|
||||
input_schema=_input_schema(
|
||||
{
|
||||
"table_name": _string("Table or view name."),
|
||||
"schema": _string("Optional schema name."),
|
||||
"column": _string("Column name."),
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of distinct values to return.",
|
||||
"minimum": 1,
|
||||
"maximum": 1000,
|
||||
"default": 20,
|
||||
},
|
||||
},
|
||||
required=("table_name", "column"),
|
||||
),
|
||||
output_schema=_SAMPLE_VALUES_OUTPUT,
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
@@ -13,6 +13,8 @@ class MetadataHandler:
|
||||
def __init__(self, runtime: MetadataRuntime) -> None:
|
||||
_handlers = {
|
||||
"list_tables": runtime.list_tables,
|
||||
"describe_table": runtime.describe_table,
|
||||
"sample_values": runtime.sample_values,
|
||||
}
|
||||
self._handlers: Mapping[str, ToolHandler] = MappingProxyType(_handlers)
|
||||
|
||||
|
||||
@@ -9,8 +9,25 @@ class MetadataStore(Protocol):
|
||||
*,
|
||||
keyword: str | None,
|
||||
limit: int,
|
||||
cursor: str | None) -> tuple[list[dict[str, Any]], str | None]: ...
|
||||
|
||||
async def describe_table(
|
||||
self,
|
||||
*,
|
||||
schema: str | None = None,
|
||||
table_name: str,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> tuple[list[dict[str, Any]], str | None]: ...
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def sample_values(
|
||||
self,
|
||||
*,
|
||||
table_name: str,
|
||||
schema: str | None,
|
||||
column: str,
|
||||
limit: int,
|
||||
) -> tuple[list[Any], bool]: ...
|
||||
|
||||
|
||||
class MetadataRuntime:
|
||||
@@ -41,4 +58,53 @@ class MetadataRuntime:
|
||||
"items": items,
|
||||
"next_cursor": next_cursor,
|
||||
"truncated": next_cursor is not None,
|
||||
}
|
||||
}
|
||||
|
||||
async def describe_table(self, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
table_name = arguments.get("table_name")
|
||||
if not isinstance(table_name, str) or not table_name:
|
||||
raise ValueError("table_name must be a non-empty string")
|
||||
|
||||
schema = arguments.get("schema")
|
||||
if schema is not None and not isinstance(schema, str):
|
||||
raise ValueError("schema 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", 200)
|
||||
if not isinstance(limit, int) or isinstance(limit, bool) or not 1 <= limit <= 200:
|
||||
raise ValueError("limit must be a positive integer between 1 and 200")
|
||||
|
||||
return await self._store.describe_table(
|
||||
table_name=table_name,
|
||||
schema=schema,
|
||||
limit=limit,
|
||||
cursor=cursor,
|
||||
)
|
||||
|
||||
async def sample_values(self, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
table_name = arguments.get("table_name")
|
||||
if not isinstance(table_name, str) or not table_name:
|
||||
raise ValueError("table_name must be a non-empty string")
|
||||
|
||||
column = arguments.get("column")
|
||||
if not isinstance(column, str) or not column:
|
||||
raise ValueError("column must be a non-empty string")
|
||||
|
||||
schema = arguments.get("schema")
|
||||
if schema is not None and not isinstance(schema, str):
|
||||
raise ValueError("schema must be a string or null")
|
||||
|
||||
limit = arguments.get("limit", 20)
|
||||
if (not isinstance(limit, int) or isinstance(limit, bool) or not 1 <= limit <= 50):
|
||||
raise ValueError("limit must be an integer between 1 and 50")
|
||||
|
||||
values, truncated = await self._store.sample_values(
|
||||
table_name=table_name,
|
||||
schema=schema,
|
||||
column=column,
|
||||
limit=limit,
|
||||
)
|
||||
return {"values": values, "truncated": truncated}
|
||||
@@ -20,7 +20,7 @@ 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.utils.pagination_util import paginate_list
|
||||
from mini_db_mcp.prompts import PromptsManager
|
||||
from mini_db_mcp.tools import ToolsManager
|
||||
|
||||
|
||||
@@ -37,7 +37,17 @@ class ToolsManager:
|
||||
if domain is None:
|
||||
raise ValueError(f"Unknown domain tool: {name}")
|
||||
|
||||
request = DomainToolRequest(**arguments)
|
||||
raw_child_tool = arguments.get("child_tool") if isinstance(arguments, dict) else None
|
||||
error_child_tool = raw_child_tool if isinstance(raw_child_tool, str) and raw_child_tool else "request"
|
||||
try:
|
||||
request = DomainToolRequest(**arguments)
|
||||
except (TypeError, ValueError) as exc:
|
||||
return self._error_result(
|
||||
name,
|
||||
error_child_tool,
|
||||
DomainErrorCode.CHILD_ARGUMENTS_INVALID,
|
||||
str(exc),
|
||||
)
|
||||
# First call: expose only this domain's child schemas to the model.
|
||||
if request.is_discovery:
|
||||
return {
|
||||
@@ -49,6 +59,22 @@ class ToolsManager:
|
||||
|
||||
child_tool = request.child_tool
|
||||
assert child_tool is not None
|
||||
child = next(
|
||||
(candidate for candidate in domain.children if candidate.tool_name == child_tool),
|
||||
None,
|
||||
)
|
||||
if child is not None:
|
||||
validation_error = _validate_arguments(
|
||||
request.execution_arguments,
|
||||
child.input_schema,
|
||||
)
|
||||
if validation_error is not None:
|
||||
return self._error_result(
|
||||
name,
|
||||
child_tool,
|
||||
DomainErrorCode.CHILD_ARGUMENTS_INVALID,
|
||||
validation_error,
|
||||
)
|
||||
# Second call: route the selected child to its owning domain handler.
|
||||
handler = self._registry.resolve_handler(name, child_tool)
|
||||
if handler is None:
|
||||
@@ -112,3 +138,49 @@ class ToolsManager:
|
||||
"details": {},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _validate_arguments(arguments: dict[str, Any], schema: dict[str, Any]) -> str | None:
|
||||
properties = schema.get("properties", {})
|
||||
missing = [name for name in schema.get("required", []) if name not in arguments]
|
||||
if missing:
|
||||
return f"Missing required argument(s): {', '.join(missing)}"
|
||||
|
||||
if schema.get("additionalProperties") is False:
|
||||
unknown = sorted(set(arguments) - set(properties))
|
||||
if unknown:
|
||||
return f"Unknown argument(s): {', '.join(unknown)}"
|
||||
|
||||
for name, value in arguments.items():
|
||||
property_schema = properties.get(name)
|
||||
if property_schema is None:
|
||||
continue
|
||||
expected = property_schema.get("type")
|
||||
expected_types = expected if isinstance(expected, list) else [expected]
|
||||
if not _matches_json_type(value, expected_types):
|
||||
return f"Argument {name!r} must have type {expected}."
|
||||
if isinstance(value, int) and not isinstance(value, bool):
|
||||
minimum = property_schema.get("minimum")
|
||||
maximum = property_schema.get("maximum")
|
||||
if minimum is not None and value < minimum:
|
||||
return f"Argument {name!r} must be at least {minimum}."
|
||||
if maximum is not None and value > maximum:
|
||||
return f"Argument {name!r} must be at most {maximum}."
|
||||
return None
|
||||
|
||||
|
||||
def _matches_json_type(value: Any, expected_types: list[Any]) -> bool:
|
||||
for expected in expected_types:
|
||||
if expected == "null" and value is None:
|
||||
return True
|
||||
if expected == "string" and isinstance(value, str):
|
||||
return True
|
||||
if expected == "integer" and isinstance(value, int) and not isinstance(value, bool):
|
||||
return True
|
||||
if expected == "boolean" and isinstance(value, bool):
|
||||
return True
|
||||
if expected == "object" and isinstance(value, dict):
|
||||
return True
|
||||
if expected == "array" and isinstance(value, list):
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -3,14 +3,14 @@ from typing import TypeVar
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
MAX_PAGE_SIZE = 10
|
||||
MAX_PAGE_SIZE = 100
|
||||
|
||||
|
||||
def paginate_list(
|
||||
items: Sequence[T],
|
||||
*,
|
||||
cursor: str | None,
|
||||
page_size: int,
|
||||
page_size: int = MAX_PAGE_SIZE,
|
||||
identifier: Callable[[T], str],
|
||||
) -> tuple[list[T], str | None]:
|
||||
"""Return a deterministic page using the prior item's identifier as cursor."""
|
||||
Reference in New Issue
Block a user