diff --git a/src/mini_db_mcp/config.py b/src/mini_db_mcp/config.py index d9af592..61f14ce 100644 --- a/src/mini_db_mcp/config.py +++ b/src/mini_db_mcp/config.py @@ -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 diff --git a/src/mini_db_mcp/connection.py b/src/mini_db_mcp/connection.py deleted file mode 100644 index 16ed9a2..0000000 --- a/src/mini_db_mcp/connection.py +++ /dev/null @@ -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) diff --git a/src/mini_db_mcp/db.py b/src/mini_db_mcp/db.py index 20d7f4c..900b68e 100644 --- a/src/mini_db_mcp/db.py +++ b/src/mini_db_mcp/db.py @@ -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 \ No newline at end of file diff --git a/src/mini_db_mcp/domain/domain_catalog.py b/src/mini_db_mcp/domain/domain_catalog.py index 9b744f5..9f126d5 100644 --- a/src/mini_db_mcp/domain/domain_catalog.py +++ b/src/mini_db_mcp/domain/domain_catalog.py @@ -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, + ), ), ), ) diff --git a/src/mini_db_mcp/domain/handler/metadata_handler.py b/src/mini_db_mcp/domain/handler/metadata_handler.py index c3002b8..6192d14 100644 --- a/src/mini_db_mcp/domain/handler/metadata_handler.py +++ b/src/mini_db_mcp/domain/handler/metadata_handler.py @@ -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) diff --git a/src/mini_db_mcp/domain/runtime/metadata_runtime.py b/src/mini_db_mcp/domain/runtime/metadata_runtime.py index 9305de0..0037515 100644 --- a/src/mini_db_mcp/domain/runtime/metadata_runtime.py +++ b/src/mini_db_mcp/domain/runtime/metadata_runtime.py @@ -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, - } \ No newline at end of file + } + + 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} \ No newline at end of file diff --git a/src/mini_db_mcp/server.py b/src/mini_db_mcp/server.py index e069a03..8a8b84d 100644 --- a/src/mini_db_mcp/server.py +++ b/src/mini_db_mcp/server.py @@ -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 diff --git a/src/mini_db_mcp/tools.py b/src/mini_db_mcp/tools.py index be3a7f5..3bbe727 100644 --- a/src/mini_db_mcp/tools.py +++ b/src/mini_db_mcp/tools.py @@ -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 diff --git a/src/mini_db_mcp/utils/pagination.py b/src/mini_db_mcp/utils/pagination_util.py similarity index 95% rename from src/mini_db_mcp/utils/pagination.py rename to src/mini_db_mcp/utils/pagination_util.py index ceba7bb..c2deb72 100644 --- a/src/mini_db_mcp/utils/pagination.py +++ b/src/mini_db_mcp/utils/pagination_util.py @@ -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."""