chDB-only mode: add introspection tools + safety baseline (P1)

README.md

@@ -39,6 +39,35 @@ An MCP server for ClickHouse.
   * Query data directly from various sources (files, URLs, databases) without ETL processes.
   * Requires the optional `chdb` extra: `pip install 'mcp-clickhouse[chdb]'`
 
+#### chDB-only introspection tools
+
+When chDB is the only enabled engine (`CHDB_ENABLED=true` and `CLICKHOUSE_ENABLED=false`), the
+following read-only introspection tools are also registered, alongside `run_chdb_select_query`.
+They are not registered when the ClickHouse server is enabled (the ClickHouse tools above own those
+names in that mode). User-supplied identifiers are validated and quoted, and output is truncated to
+`CHDB_MAX_RESULT_BYTES`.
+
+* `list_databases`: List databases in the chDB engine.
+* `list_tables`: List tables in a chDB database. Input: `database` (string).
+* `describe_table`: Return column names and types for a chDB table. Input: `database`, `table` (strings).
+* `get_sample_data`: Return the first rows of a chDB table. Input: `database`, `table` (strings), `limit` (int, default 10, clamped to 1–1000).
+* `list_functions`: List SQL functions available in the chDB engine. Input: `pattern` (optional string, case-insensitive name filter).
+
+#### chDB-only security baseline
+
+In chDB-only mode the chDB session also runs under a security baseline so the agent-facing
+query surface is bounded by default:
+
+* **Read-only by default**: the session is put under `SET readonly=2`, which rejects
+  INSERT/CREATE/DROP/ALTER on persistent tables while keeping `SET` and table functions usable.
+  Set `CHDB_ALLOW_WRITE_ACCESS=true` to allow writes.
+* **Bounded output**: results are capped at `CHDB_MAX_RESULT_BYTES` both at the engine level
+  (`max_result_bytes`) and as a final truncation of the serialized payload.
+* **Bounded runtime**: queries are aborted past the shared query timeout (`max_execution_time`).
+* **Optional source sandbox**: set `CHDB_FILE_ALLOWLIST` to refuse external/file table functions
+  (`file`/`url`/`s3`/`remote`/`postgresql`/`executable`/`python`/…) in raw `run_chdb_select_query`
+  SQL. The list of disallowed functions is taken from the live engine's `system.table_functions`.
+
 ### Health Check Endpoint
 
 When running with HTTP or SSE transport, a health check endpoint is available at `/health`. This endpoint:
@@ -579,6 +608,15 @@ The following environment variables are used to configure the ClickHouse and chD
   * Default: `":memory:"` (in-memory database)
   * Use `:memory:` for in-memory database
   * Use a file path for persistent storage (e.g., `/path/to/chdb/data`)
+* `CHDB_MAX_RESULT_BYTES`: Maximum size (in bytes) of chDB output; capped at the engine level and truncated again on the serialized payload, with a notice
+  * Default: `1048576` (1 MiB)
+  * Applies in chDB-only mode (introspection tools and `run_chdb_select_query`)
+* `CHDB_ALLOW_WRITE_ACCESS`: Allow writes against the chDB session
+  * Default: `"false"` (session runs under `SET readonly=2`)
+  * Set to `"true"` to permit INSERT/CREATE/DROP/ALTER
+* `CHDB_FILE_ALLOWLIST`: Colon-separated path prefixes that sandbox the chDB engine
+  * Default: unset (no gating — behavior unchanged)
+  * When set, `run_chdb_select_query` refuses external/file table functions (`file`/`url`/`s3`/`remote`/…) in raw SQL
 
 #### Example Configurations
 

mcp_clickhouse/chdb_safety.py

@@ -0,0 +1,30 @@
+"""Self-contained output-bounding helper for the chDB query path.
+
+A pure function with no third-party or project dependencies, so it stays
+importable when the optional ``chdb`` extra is not installed. SQL source
+scanning for the file allowlist (table-function detection, path checks) ships
+in ``chdb.agents.safety`` and is imported lazily on the chDB-only code paths.
+
+Identifier/string quoting for the introspection tools is NOT here — those tools
+go through ``chdb.agents.ChDBTool``, which owns quoting and parameter binding.
+"""
+
+from __future__ import annotations
+
+_TRUNCATION_NOTICE = (
+    "\n\n[... output truncated at {limit} bytes; narrow the query or raise "
+    "CHDB_MAX_RESULT_BYTES ...]"
+)
+
+
+def truncate_text(text: str, limit_bytes: int) -> str:
+    """Trim text to at most ``limit_bytes`` UTF-8 bytes, appending a notice if cut.
+
+    Trimming is done on the encoded bytes (not characters) so the byte budget is
+    respected exactly; a partial trailing multi-byte character is dropped.
+    """
+    encoded = text.encode("utf-8")
+    if len(encoded) <= limit_bytes:
+        return text
+    trimmed = encoded[:limit_bytes].decode("utf-8", errors="ignore")
+    return trimmed + _TRUNCATION_NOTICE.format(limit=limit_bytes)

mcp_clickhouse/chdb_tools.py

@@ -0,0 +1,154 @@
+"""chDB-only introspection tools for mcp-clickhouse.
+
+These tools are exposed ONLY in chDB-only mode (chDB enabled, ClickHouse server
+disabled). In that mode the ClickHouse-server tools are not registered, so the
+bare canonical names are free; these take them (``list_databases``,
+``list_tables``, ``describe_table``, ``get_sample_data``, ``list_functions``)
+and operate on the in-process chDB engine, alongside the existing
+``run_chdb_select_query`` tool.
+
+Implementation: this module is a thin adapter over ``chdb.agents.ChDBTool`` (the
+canonical, cross-language chDB agent-tool contract). ChDBTool owns the shared
+behavior — read-only enforcement, parameter binding, identifier quoting, result
+caps, typed errors, and an optional query timeout — so mcp-clickhouse does not
+re-implement any of it. ChDBTool never mutates the externally owned session it
+is handed: at construction it verifies that the session's ``readonly`` setting
+matches the declared ``read_only`` flag and fails with a CONFIG_MISMATCH error
+otherwise (the server locks the session to ``readonly=2`` at init). The tool
+signatures here keep the ClickHouse-server-parity ``(database, table)`` shape
+and map onto ChDBTool's ``database=`` qualifier. See chdb.agents.CONTRACT.md.
+
+Decoupling: self-contained; injected by the single caller
+(``register_chdb_only_tools``) with a chDB-client factory, the shared executor,
+and a query-timeout provider. The ClickHouse code path does not import this.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import json
+import logging
+from typing import Callable, Optional
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools import Tool
+
+logger = logging.getLogger(__name__)
+
+
+def register_chdb_only_tools(
+    mcp,
+    *,
+    max_result_bytes: int,
+    create_client: Callable[[], object],
+    query_executor: concurrent.futures.ThreadPoolExecutor,
+    query_timeout: Callable[[], int],
+    allow_write: bool = False,
+) -> None:
+    """Register the chDB-only introspection tools on the FastMCP instance.
+
+    Call exactly once, only in chDB-only mode. Arguments are injected to keep
+    this module decoupled from the ClickHouse server code:
+
+    - ``max_result_bytes``: output byte cap, taken from ``ChDBConfig`` by the
+      caller (not re-read from the environment here).
+    - ``create_client``: returns the shared chDB client (a chdb Session).
+    - ``query_executor``: thread pool to run blocking chDB work off the event loop.
+    - ``query_timeout``: per-query timeout in seconds (also fed to ChDBTool's
+      engine-side ``max_execution_time`` as defense in depth).
+    - ``allow_write``: must match the session's mode. When False (default) the
+      session must already be under ``SET readonly=2`` (the server applies it
+      at init) so writes/DDL are rejected by the engine.
+    """
+    from chdb.agents import ChDBError, ChDBTool
+
+    # One ChDBTool over the shared session. It probes (never mutates) the
+    # session's readonly mode — construction fails with CONFIG_MISMATCH if the
+    # flag disagrees with the session — and applies the byte cap and the
+    # engine-side timeout.
+    tool = ChDBTool(
+        session=create_client(),
+        read_only=not allow_write,
+        max_bytes=max_result_bytes,
+        max_execution_time=query_timeout(),
+    )
+
+    async def _run(work: Callable[[], object], tool_name: str) -> str:
+        """Run a ChDBTool call on the executor with a timeout; return JSON text.
+
+        Output mirrors the ClickHouse-server tools: success is bare JSON, and any
+        failure (engine error or timeout) is raised as a ``ToolError`` — the same
+        way ``run_query`` surfaces errors — rather than returned as a string."""
+        timeout = query_timeout()
+        future = query_executor.submit(work)
+        try:
+            result = await asyncio.wait_for(asyncio.wrap_future(future), timeout=timeout)
+        except asyncio.TimeoutError:
+            future.cancel()
+            logger.warning("chDB %s timed out after %ss", tool_name, timeout)
+            raise ToolError(f"chDB {tool_name} timed out after {timeout} seconds")
+        except ChDBError as err:
+            raise ToolError(err.message)
+        except Exception as err:  # noqa: BLE001
+            logger.error("chDB %s failed: %s", tool_name, err)
+            raise ToolError(str(err))
+        return json.dumps(result, ensure_ascii=False, default=str)
+
+    async def list_databases() -> str:
+        """List databases in the in-process chDB engine."""
+        return await _run(tool.list_databases, "list_databases")
+
+    async def list_tables(database: str) -> str:
+        """List tables in a chDB database.
+
+        Args:
+            database: Database name (plain SQL identifier).
+        """
+        return await _run(lambda: tool.list_tables(database), "list_tables")
+
+    async def describe_table(database: str, table: str) -> str:
+        """Return column names and types for a chDB table.
+
+        Args:
+            database: Database name (plain identifier).
+            table: Table name (plain identifier).
+        """
+        return await _run(lambda: tool.describe(table, database=database), "describe_table")
+
+    async def get_sample_data(database: str, table: str, limit: int = 10) -> str:
+        """Return the first rows of a chDB table.
+
+        Args:
+            database: Database name (plain identifier).
+            table: Table name (plain identifier).
+            limit: Maximum rows to return; clamped to [1, 1000].
+        """
+        n = max(1, min(int(limit), 1000))
+        return await _run(
+            lambda: tool.get_sample_data(table, database=database, limit=n).rows,
+            "get_sample_data",
+        )
+
+    async def list_functions(pattern: Optional[str] = None) -> str:
+        """List SQL functions available in the chDB engine.
+
+        Args:
+            pattern: Optional case-insensitive substring filter on the function
+                name.
+        """
+        # `pattern` is a substring; ChDBTool's `like` is a raw LIKE pattern, so
+        # wrap with %...% to preserve the substring-match semantics.
+        like = f"%{pattern}%" if pattern else None
+        return await _run(lambda: tool.list_functions(like=like), "list_functions")
+
+    tools = (
+        (list_databases, "list_databases", "List databases in the in-process chDB engine."),
+        (list_tables, "list_tables", "List tables in a chDB database."),
+        (describe_table, "describe_table", "Return column names and types for a chDB table."),
+        (get_sample_data, "get_sample_data", "Return the first rows of a chDB table."),
+        (list_functions, "list_functions", "List SQL functions available in the chDB engine."),
+    )
+    for fn, name, description in tools:
+        mcp.add_tool(Tool.from_function(fn, name=name, description=description))
+    logger.info("chDB-only introspection tools registered (%d tools, over chdb.agents.ChDBTool)", len(tools))

mcp_clickhouse/mcp_env.py

@@ -218,8 +218,20 @@ class ChDBConfig:
 
     Required environment variables:
         CHDB_DATA_PATH: The path to the chDB data directory (only required if CHDB_ENABLED=true)
+
+    Optional environment variables (with defaults):
+        CHDB_ALLOW_WRITE_ACCESS: Allow writes; when false, the session runs under
+            SET readonly=2 (table functions stay usable) (default: false)
+        CHDB_MAX_RESULT_BYTES: Engine result-size cap, also used as the final
+            Python truncation budget (default: 1048576, i.e. 1 MiB)
+        CHDB_FILE_ALLOWLIST: Colon-separated path prefixes. When set, the chDB
+            query tool is sandboxed: raw SQL may not call external/file table
+            functions (file/url/s3/remote/...) (default: unset)
     """
 
+    # Engine result-size cap (1 MiB) when CHDB_MAX_RESULT_BYTES is unset/invalid.
+    DEFAULT_MAX_RESULT_BYTES = 1024 * 1024
+
     def __init__(self):
         """Initialize the configuration from environment variables."""
         if self.enabled:
@@ -238,6 +250,49 @@ def data_path(self) -> str:
         """Get the chDB data path."""
         return os.getenv("CHDB_DATA_PATH", ":memory:")
 
+    @property
+    def allow_write_access(self) -> bool:
+        """Get whether writes are allowed.
+
+        When false (the default), the chDB session is put under SET readonly=2,
+        which rejects INSERT/CREATE/DROP/ALTER on persistent tables while still
+        allowing SET and table functions.
+
+        Default: False
+        """
+        return os.getenv("CHDB_ALLOW_WRITE_ACCESS", "false").lower() == "true"
+
+    @property
+    def max_result_bytes(self) -> int:
+        """Get the result-size cap in bytes (engine cap + Python truncation).
+
+        Falls back to DEFAULT_MAX_RESULT_BYTES when unset or not a positive int.
+
+        Default: 1048576 (1 MiB)
+        """
+        raw = os.getenv("CHDB_MAX_RESULT_BYTES")
+        if not raw:
+            return self.DEFAULT_MAX_RESULT_BYTES
+        try:
+            value = int(raw)
+        except ValueError:
+            return self.DEFAULT_MAX_RESULT_BYTES
+        return value if value > 0 else self.DEFAULT_MAX_RESULT_BYTES
+
+    @property
+    def file_allowlist(self) -> tuple[str, ...]:
+        """Colon-separated file-path allowlist prefixes (like DuckDB's ``allowed_directories``).
+
+        When non-empty, chDB may read file-like sources (file/s3/url/…) only when
+        their path is under one of these prefixes; paths outside — and DSN-based
+        external sources that can't be path-checked — are refused. Empty by
+        default (no restriction). See ``CHDB_FILE_ALLOWLIST``.
+        """
+        raw = os.getenv("CHDB_FILE_ALLOWLIST")
+        if not raw:
+            return ()
+        return tuple(p for p in (part.strip() for part in raw.split(":")) if p)
+
     def get_client_config(self) -> dict:
         """Get the configuration dictionary for chDB client.