|
| 1 | +"""chDB-only introspection tools for mcp-clickhouse. |
| 2 | +
|
| 3 | +These tools are exposed ONLY in chDB-only mode (chDB enabled, ClickHouse server |
| 4 | +disabled). In that mode the ClickHouse-server tools are not registered, so the |
| 5 | +bare canonical names are free; these take them (``list_databases``, |
| 6 | +``list_tables``, ``describe_table``, ``get_sample_data``, ``list_functions``) |
| 7 | +and operate on the in-process chDB engine, alongside the existing |
| 8 | +``run_chdb_select_query`` tool. |
| 9 | +
|
| 10 | +Implementation: this module is a thin adapter over ``chdb.agents.ChDBTool`` (the |
| 11 | +canonical, cross-language chDB agent-tool contract). ChDBTool owns the shared |
| 12 | +behavior — read-only enforcement, parameter binding, identifier quoting, result |
| 13 | +caps, typed errors, and an optional query timeout — so mcp-clickhouse does not |
| 14 | +re-implement any of it. ChDBTool never mutates the externally owned session it |
| 15 | +is handed: at construction it verifies that the session's ``readonly`` setting |
| 16 | +matches the declared ``read_only`` flag and fails with a CONFIG_MISMATCH error |
| 17 | +otherwise (the server locks the session to ``readonly=2`` at init). The tool |
| 18 | +signatures here keep the ClickHouse-server-parity ``(database, table)`` shape |
| 19 | +and map onto ChDBTool's ``database=`` qualifier. See chdb.agents.CONTRACT.md. |
| 20 | +
|
| 21 | +Decoupling: self-contained; injected by the single caller |
| 22 | +(``register_chdb_only_tools``) with a chDB-client factory, the shared executor, |
| 23 | +and a query-timeout provider. The ClickHouse code path does not import this. |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import asyncio |
| 29 | +import concurrent.futures |
| 30 | +import json |
| 31 | +import logging |
| 32 | +from typing import Callable, Optional |
| 33 | + |
| 34 | +from fastmcp.exceptions import ToolError |
| 35 | +from fastmcp.tools import Tool |
| 36 | + |
| 37 | +logger = logging.getLogger(__name__) |
| 38 | + |
| 39 | + |
| 40 | +def register_chdb_only_tools( |
| 41 | + mcp, |
| 42 | + *, |
| 43 | + max_result_bytes: int, |
| 44 | + create_client: Callable[[], object], |
| 45 | + query_executor: concurrent.futures.ThreadPoolExecutor, |
| 46 | + query_timeout: Callable[[], int], |
| 47 | + allow_write: bool = False, |
| 48 | +) -> None: |
| 49 | + """Register the chDB-only introspection tools on the FastMCP instance. |
| 50 | +
|
| 51 | + Call exactly once, only in chDB-only mode. Arguments are injected to keep |
| 52 | + this module decoupled from the ClickHouse server code: |
| 53 | +
|
| 54 | + - ``max_result_bytes``: output byte cap, taken from ``ChDBConfig`` by the |
| 55 | + caller (not re-read from the environment here). |
| 56 | + - ``create_client``: returns the shared chDB client (a chdb Session). |
| 57 | + - ``query_executor``: thread pool to run blocking chDB work off the event loop. |
| 58 | + - ``query_timeout``: per-query timeout in seconds (also fed to ChDBTool's |
| 59 | + engine-side ``max_execution_time`` as defense in depth). |
| 60 | + - ``allow_write``: must match the session's mode. When False (default) the |
| 61 | + session must already be under ``SET readonly=2`` (the server applies it |
| 62 | + at init) so writes/DDL are rejected by the engine. |
| 63 | + """ |
| 64 | + from chdb.agents import ChDBError, ChDBTool |
| 65 | + |
| 66 | + # One ChDBTool over the shared session. It probes (never mutates) the |
| 67 | + # session's readonly mode — construction fails with CONFIG_MISMATCH if the |
| 68 | + # flag disagrees with the session — and applies the byte cap and the |
| 69 | + # engine-side timeout. |
| 70 | + tool = ChDBTool( |
| 71 | + session=create_client(), |
| 72 | + read_only=not allow_write, |
| 73 | + max_bytes=max_result_bytes, |
| 74 | + max_execution_time=query_timeout(), |
| 75 | + ) |
| 76 | + |
| 77 | + async def _run(work: Callable[[], object], tool_name: str) -> str: |
| 78 | + """Run a ChDBTool call on the executor with a timeout; return JSON text. |
| 79 | +
|
| 80 | + Output mirrors the ClickHouse-server tools: success is bare JSON, and any |
| 81 | + failure (engine error or timeout) is raised as a ``ToolError`` — the same |
| 82 | + way ``run_query`` surfaces errors — rather than returned as a string.""" |
| 83 | + timeout = query_timeout() |
| 84 | + future = query_executor.submit(work) |
| 85 | + try: |
| 86 | + result = await asyncio.wait_for(asyncio.wrap_future(future), timeout=timeout) |
| 87 | + except asyncio.TimeoutError: |
| 88 | + future.cancel() |
| 89 | + logger.warning("chDB %s timed out after %ss", tool_name, timeout) |
| 90 | + raise ToolError(f"chDB {tool_name} timed out after {timeout} seconds") |
| 91 | + except ChDBError as err: |
| 92 | + raise ToolError(err.message) |
| 93 | + except Exception as err: # noqa: BLE001 |
| 94 | + logger.error("chDB %s failed: %s", tool_name, err) |
| 95 | + raise ToolError(str(err)) |
| 96 | + return json.dumps(result, ensure_ascii=False, default=str) |
| 97 | + |
| 98 | + async def list_databases() -> str: |
| 99 | + """List databases in the in-process chDB engine.""" |
| 100 | + return await _run(tool.list_databases, "list_databases") |
| 101 | + |
| 102 | + async def list_tables(database: str) -> str: |
| 103 | + """List tables in a chDB database. |
| 104 | +
|
| 105 | + Args: |
| 106 | + database: Database name (plain SQL identifier). |
| 107 | + """ |
| 108 | + return await _run(lambda: tool.list_tables(database), "list_tables") |
| 109 | + |
| 110 | + async def describe_table(database: str, table: str) -> str: |
| 111 | + """Return column names and types for a chDB table. |
| 112 | +
|
| 113 | + Args: |
| 114 | + database: Database name (plain identifier). |
| 115 | + table: Table name (plain identifier). |
| 116 | + """ |
| 117 | + return await _run(lambda: tool.describe(table, database=database), "describe_table") |
| 118 | + |
| 119 | + async def get_sample_data(database: str, table: str, limit: int = 10) -> str: |
| 120 | + """Return the first rows of a chDB table. |
| 121 | +
|
| 122 | + Args: |
| 123 | + database: Database name (plain identifier). |
| 124 | + table: Table name (plain identifier). |
| 125 | + limit: Maximum rows to return; clamped to [1, 1000]. |
| 126 | + """ |
| 127 | + n = max(1, min(int(limit), 1000)) |
| 128 | + return await _run( |
| 129 | + lambda: tool.get_sample_data(table, database=database, limit=n).rows, |
| 130 | + "get_sample_data", |
| 131 | + ) |
| 132 | + |
| 133 | + async def list_functions(pattern: Optional[str] = None) -> str: |
| 134 | + """List SQL functions available in the chDB engine. |
| 135 | +
|
| 136 | + Args: |
| 137 | + pattern: Optional case-insensitive substring filter on the function |
| 138 | + name. |
| 139 | + """ |
| 140 | + # `pattern` is a substring; ChDBTool's `like` is a raw LIKE pattern, so |
| 141 | + # wrap with %...% to preserve the substring-match semantics. |
| 142 | + like = f"%{pattern}%" if pattern else None |
| 143 | + return await _run(lambda: tool.list_functions(like=like), "list_functions") |
| 144 | + |
| 145 | + tools = ( |
| 146 | + (list_databases, "list_databases", "List databases in the in-process chDB engine."), |
| 147 | + (list_tables, "list_tables", "List tables in a chDB database."), |
| 148 | + (describe_table, "describe_table", "Return column names and types for a chDB table."), |
| 149 | + (get_sample_data, "get_sample_data", "Return the first rows of a chDB table."), |
| 150 | + (list_functions, "list_functions", "List SQL functions available in the chDB engine."), |
| 151 | + ) |
| 152 | + for fn, name, description in tools: |
| 153 | + mcp.add_tool(Tool.from_function(fn, name=name, description=description)) |
| 154 | + logger.info("chDB-only introspection tools registered (%d tools, over chdb.agents.ChDBTool)", len(tools)) |
0 commit comments