diff --git a/README.md b/README.md index 2358b4a3..08387bfc 100644 --- a/README.md +++ b/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 diff --git a/mcp_clickhouse/chdb_safety.py b/mcp_clickhouse/chdb_safety.py new file mode 100644 index 00000000..8e89dfb5 --- /dev/null +++ b/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) diff --git a/mcp_clickhouse/chdb_tools.py b/mcp_clickhouse/chdb_tools.py new file mode 100644 index 00000000..026127f3 --- /dev/null +++ b/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)) diff --git a/mcp_clickhouse/mcp_env.py b/mcp_clickhouse/mcp_env.py index 0492ca6a..57fe2512 100644 --- a/mcp_clickhouse/mcp_env.py +++ b/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. diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index c8a12243..9ddd3b0b 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -23,6 +23,7 @@ from starlette.responses import PlainTextResponse from mcp_clickhouse.chdb_prompt import CHDB_PROMPT +from mcp_clickhouse.chdb_safety import truncate_text from mcp_clickhouse.skills_advisor import CLICKHOUSE_SERVER_INSTRUCTIONS from mcp_clickhouse.mcp_env import TransportType, get_chdb_config, get_config, get_mcp_config @@ -136,6 +137,11 @@ def _resolve_auth(mcp_config) -> Dict[str, Any]: ) _chdb_client = None _chdb_error_message: Optional[str] = None +# Lowercase names of every table function the live chDB build exposes, snapshot +# from system.table_functions at session init. The allowlist scanner consults +# this so it tracks whatever the running engine offers instead of a denylist +# that goes stale. Empty until a chDB session is opened. +_chdb_known_table_functions: frozenset = frozenset() @mcp.custom_route("/health", methods=["GET"]) @@ -678,9 +684,56 @@ def create_chdb_client(): return _chdb_client +# file-like table functions whose first argument is a path/URI the allowlist +# can constrain (kept in sync with chdb.agents.safety.scan_file_paths). +_FILE_LIKE_FUNCTIONS = frozenset({"file", "s3", "url", "hdfs", "azureblobstorage"}) + + +def _chdb_sql_source_violation(query: str) -> Optional[str]: + """Return an error message if ``query`` reaches a source outside the allowlist. + + Enforced only when CHDB_FILE_ALLOWLIST is set. It is a **path allowlist** + (like DuckDB's ``allowed_directories``): file-like sources + (file/s3/url/hdfs/azureBlobStorage) are permitted when their path is under an + allowed prefix and refused otherwise; DSN-based external sources + (postgresql/mysql/remote/…) cannot be constrained by a path and are refused + while the allowlist is active. Returns None when the allowlist is unset + (default — behavior unchanged) or the query is clean. + + Uses chdb.agents.safety (the same path-allowlist logic ChDBTool enforces) so + the raw-SQL path and the ChDBTool path agree. + """ + allow = get_chdb_config().file_allowlist + if not allow: + return None + from chdb.agents.safety import ( + FALLBACK_KNOWN_TABLE_FUNCTIONS, + find_source_calls, + path_allowed, + scan_file_paths, + ) + + outside = [path for _fn, path in scan_file_paths(query) if not path_allowed(path, allow)] + if outside: + return f"refusing file paths outside CHDB_FILE_ALLOWLIST {list(allow)}: {outside}" + + known = _chdb_known_table_functions or FALLBACK_KNOWN_TABLE_FUNCTIONS + non_file = sorted({name for name, _ in find_source_calls(query, known)} - _FILE_LIKE_FUNCTIONS) + if non_file: + return ( + f"refusing non-file external sources {non_file} while CHDB_FILE_ALLOWLIST " + f"is set (a path allowlist cannot constrain a DSN-based source)" + ) + return None + + def execute_chdb_query(query: str): """Execute a query using chDB client.""" client = create_chdb_client() + violation = _chdb_sql_source_violation(query) + if violation is not None: + logger.warning("chDB query rejected by file allowlist: %s", violation) + return {"error": violation} try: res = client.query(query, "JSON") if res.has_error(): @@ -708,7 +761,10 @@ def _process_chdb_result(result) -> str: "status": "error", "message": f"chDB query failed: {result['error']}", }) - return _serialize_tool_result(result) + # Final guardrail: bound the serialized payload even though the engine is + # already capped via max_result_bytes, so a large result can't flood the + # caller's context. + return truncate_text(_serialize_tool_result(result), get_chdb_config().max_result_bytes) def run_chdb_select_query(query: str) -> str: @@ -762,6 +818,52 @@ def chdb_initial_prompt() -> str: return CHDB_PROMPT +def _apply_chdb_session_baseline(client) -> None: + """Apply the chDB-session security baseline once at init. + + - Snapshot system.table_functions so the allowlist scanner tracks the live + engine (falls back to a hardcoded set if the catalog query fails). + - Cap engine work: bound result bytes and abort runaway queries by wall time. + - Unless CHDB_ALLOW_WRITE_ACCESS=true, put the session under SET readonly=2, + which rejects persistent-table writes while keeping table functions usable. + + The result byte cap and timeout are applied BEFORE readonly=2, in case a + future chDB release tightens which settings stay writable under readonly. + """ + global _chdb_known_table_functions + from chdb.agents.safety import FALLBACK_KNOWN_TABLE_FUNCTIONS + + try: + rows = client.query("SELECT lower(name) FROM system.table_functions", "TabSeparated") + names = {line.strip() for line in str(rows).splitlines() if line.strip()} + _chdb_known_table_functions = ( + frozenset(names) if names else FALLBACK_KNOWN_TABLE_FUNCTIONS + ) + except Exception as e: # noqa: BLE001 — defensive against catalog issues + logger.warning( + "system.table_functions unavailable (%s); using fallback denylist", e + ) + _chdb_known_table_functions = FALLBACK_KNOWN_TABLE_FUNCTIONS + + chdb_config = get_chdb_config() + # max_block_size is shrunk because result_overflow_mode='break' only checks + # between blocks; at the default block size a single block can blow past + # max_result_bytes before break fires. With a small block the engine + # overshoots by at most one block and the Python truncate() trims the rest. + setup_parts = [ + "max_block_size = 8192", + f"max_result_bytes = {chdb_config.max_result_bytes}", + "result_overflow_mode = 'break'", + ] + query_timeout = get_mcp_config().query_timeout + if query_timeout and query_timeout > 0: + setup_parts.append(f"max_execution_time = {query_timeout}") + client.query("SET " + ", ".join(setup_parts), "TabSeparated") + + if not chdb_config.allow_write_access: + client.query("SET readonly=2", "TabSeparated") + + def _init_chdb_client(): """Initialize the global chDB client instance.""" global _chdb_error_message @@ -777,6 +879,7 @@ def _init_chdb_client(): import chdb.session as chs client = chs.Session(path=data_path) + _apply_chdb_session_baseline(client) _chdb_error_message = None logger.info(f"Successfully connected to chDB with data_path={data_path}") return client @@ -851,3 +954,33 @@ def _register_chdb_tools(): _register_chdb_tools() + + +# When ONLY chDB is enabled (the ClickHouse server is disabled), the bare tool +# names used above are free, so expose a richer set of read-only chDB +# introspection tools under those canonical names. This is the single +# integration point with the otherwise self-contained chDB tool module: it +# injects the chDB client factory, the shared query executor, and the query +# timeout, so that module never imports the ClickHouse server logic. In any +# mode where the ClickHouse server is enabled, these tools are not registered +# and existing chDB behavior (run_chdb_select_query) is unchanged. +if ( + get_chdb_config().enabled + and os.getenv("CLICKHOUSE_ENABLED", "true").lower() != "true" + and _chdb_client is not None +): + from mcp_clickhouse.chdb_tools import register_chdb_only_tools + + register_chdb_only_tools( + mcp, + max_result_bytes=get_chdb_config().max_result_bytes, + create_client=create_chdb_client, + query_executor=QUERY_EXECUTOR, + query_timeout=lambda: get_mcp_config().query_timeout, + # Must match the session's mode: ChDBTool never mutates an externally + # provided session — it probes the session's readonly setting at + # construction and fails with a CONFIG_MISMATCH error if it disagrees + # with the declared flag. _init_chdb_client has already locked the + # session to readonly=2 unless writes are allowed, so pass the same flag. + allow_write=get_chdb_config().allow_write_access, + ) diff --git a/pyproject.toml b/pyproject.toml index 96398342..1f7d2604 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ Home = "https://github.com/ClickHouse/mcp-clickhouse" [project.optional-dependencies] chdb = [ - "chdb>=3.3.0" + "chdb>=4.2.1" # needs chdb.agents contract 0.2.0: external-session verification + chdb.agents.safety scanner ] dev = [ "ruff", diff --git a/tests/test_chdb_safety.py b/tests/test_chdb_safety.py new file mode 100644 index 00000000..138d2fa6 --- /dev/null +++ b/tests/test_chdb_safety.py @@ -0,0 +1,26 @@ +"""Unit tests for the self-contained chDB output-bounding helper. + +Pure-function tests with no chDB or server dependency. SQL source scanning for +the file allowlist is owned (and tested) upstream in ``chdb.agents.safety``; +its integration into the query gate is covered by test_chdb_security_baseline. +""" + +from mcp_clickhouse.chdb_safety import truncate_text + + +def test_truncate_text_passthrough_when_within_budget(): + assert truncate_text("short", 100) == "short" + + +def test_truncate_text_cuts_and_appends_notice(): + out = truncate_text("x" * 100, 10) + head = out.split("\n\n")[0] + assert head == "x" * 10 + assert "truncated at 10 bytes" in out + + +def test_truncate_text_respects_byte_budget_for_multibyte(): + # 10 two-byte chars = 20 bytes; budget 5 bytes -> at most 5 bytes kept + out = truncate_text("é" * 10, 5) + head = out.split("\n\n")[0] + assert len(head.encode("utf-8")) <= 5 diff --git a/tests/test_chdb_security_baseline.py b/tests/test_chdb_security_baseline.py new file mode 100644 index 00000000..38fdfa31 --- /dev/null +++ b/tests/test_chdb_security_baseline.py @@ -0,0 +1,103 @@ +"""Integration tests for the chDB-only security baseline. + +These open a real in-memory chDB session through ``_init_chdb_client`` and +assert the baseline is actually enforced by the engine: ``SET readonly=2`` +(unless writes are allowed), the result-byte / execution-time caps, and the +CHDB_FILE_ALLOWLIST sandbox over raw SQL. Skipped if chDB is not installed. +""" + +from unittest.mock import patch + +import pytest + +from mcp_clickhouse import mcp_server + +chdb = pytest.importorskip("chdb") + + +@pytest.fixture +def chdb_session(monkeypatch): + """Yield a real in-memory chDB session with the baseline applied. + + The session is also wired in as ``mcp_server._chdb_client`` so the + ``execute_chdb_query`` path can run against it, and closed afterwards. + """ + monkeypatch.setenv("CHDB_ENABLED", "true") + monkeypatch.setenv("CHDB_DATA_PATH", ":memory:") + + client = mcp_server._init_chdb_client() + assert client is not None, mcp_server._chdb_error_message + with patch.object(mcp_server, "_chdb_client", client): + try: + yield client + finally: + client.close() + + +def test_readonly_setting_is_2_by_default(chdb_session): + out = str(chdb_session.query("SELECT value FROM system.settings WHERE name = 'readonly'", "CSV")) + assert "2" in out + + +def test_readonly_blocks_ddl_by_default(chdb_session): + result = mcp_server.execute_chdb_query("CREATE TABLE t (a Int32) ENGINE = Memory") + assert isinstance(result, dict) and "error" in result + + +def test_max_result_bytes_setting_applied(chdb_session): + out = str( + chdb_session.query("SELECT value FROM system.settings WHERE name = 'max_result_bytes'", "CSV") + ) + assert "1048576" in out + + +def test_writes_allowed_when_flag_set(monkeypatch): + monkeypatch.setenv("CHDB_ENABLED", "true") + monkeypatch.setenv("CHDB_DATA_PATH", ":memory:") + monkeypatch.setenv("CHDB_ALLOW_WRITE_ACCESS", "true") + client = mcp_server._init_chdb_client() + assert client is not None, mcp_server._chdb_error_message + try: + out = str(client.query("SELECT value FROM system.settings WHERE name = 'readonly'", "CSV")) + # readonly not forced to 2 when writes are allowed + assert "2" not in out + finally: + client.close() + + +def test_select_passes_through(chdb_session): + result = mcp_server.execute_chdb_query("SELECT 1 AS n") + assert result == [{"n": 1}] + + +def test_allowlist_rejects_external_table_function(chdb_session, monkeypatch): + monkeypatch.setenv("CHDB_FILE_ALLOWLIST", "/data") + # url() must be refused before any network access is attempted. + result = mcp_server.execute_chdb_query("SELECT * FROM url('http://169.254.169.254/', 'CSV')") + assert isinstance(result, dict) and "error" in result + assert "CHDB_FILE_ALLOWLIST" in result["error"] + + +def test_allowlist_allows_safe_query(chdb_session, monkeypatch): + monkeypatch.setenv("CHDB_FILE_ALLOWLIST", "/data") + result = mcp_server.execute_chdb_query("SELECT * FROM numbers(3)") + assert result == [{"number": 0}, {"number": 1}, {"number": 2}] + + +def test_allowlist_is_a_path_allowlist(monkeypatch): + # It's a path allowlist (DuckDB allowed_directories style), not a full sandbox: + monkeypatch.setenv("CHDB_FILE_ALLOWLIST", "/data") + # a file-like path UNDER an allowed prefix is permitted + assert mcp_server._chdb_sql_source_violation("SELECT * FROM file('/data/x.parquet')") is None + # a path OUTSIDE the allowlist is refused + assert mcp_server._chdb_sql_source_violation("SELECT * FROM file('/etc/passwd')") is not None + # a DSN-based source can't be path-checked -> refused while the allowlist is active + assert mcp_server._chdb_sql_source_violation( + "SELECT * FROM postgresql('h:5432', 'db', 't', 'u', 'p')" + ) is not None + + +def test_no_gating_when_allowlist_unset(chdb_session, monkeypatch): + monkeypatch.delenv("CHDB_FILE_ALLOWLIST", raising=False) + # With no allowlist the scanner must not flag anything (0.4.0 behavior). + assert mcp_server._chdb_sql_source_violation("SELECT * FROM url('http://x', 'CSV')") is None diff --git a/tests/test_chdb_tools.py b/tests/test_chdb_tools.py new file mode 100644 index 00000000..24108e60 --- /dev/null +++ b/tests/test_chdb_tools.py @@ -0,0 +1,148 @@ +"""End-to-end tests for the chDB-only introspection tools. + +These drive the tools through a real FastMCP server (via an in-memory FastMCP +client) backed by a real in-process chDB session, so the full path is exercised: +tool dispatch -> async wrapper -> thread-pool execution -> chDB -> truncation. + +Skipped when chDB is not installed (the optional ``[chdb]`` extra). +""" + +import concurrent.futures +import json + +import pytest + +chdb_session = pytest.importorskip("chdb.session") +from fastmcp import Client, FastMCP # noqa: E402 + +from mcp_clickhouse.chdb_tools import register_chdb_only_tools # noqa: E402 + +pytestmark = pytest.mark.asyncio + + +# Module-scoped: chDB is one-session-per-process, so the session and its seeded +# data are created once and reused across the module's tests. Setup is +# idempotent so re-runs in the same process do not fail on existing objects. +@pytest.fixture(scope="module") +def chdb_client(): + """A real in-process chDB session seeded with a demo database/table.""" + session = chdb_session.Session() + session.query("CREATE DATABASE IF NOT EXISTS demo", "TabSeparated") + session.query("DROP TABLE IF EXISTS demo.t", "TabSeparated") + session.query("CREATE TABLE demo.t (id Int32, name String) ENGINE = Memory", "TabSeparated") + session.query("INSERT INTO demo.t VALUES (1, 'a'), (2, 'b'), (3, 'c')", "TabSeparated") + # Mirror the server init flow: the session is locked to readonly=2 before it + # is handed to ChDBTool, which verifies (and never mutates) the readonly + # mode of an externally provided session. + session.query("SET readonly=2", "TabSeparated") + yield session + session.close() + + +@pytest.fixture(scope="module") +def mcp_with_tools(chdb_client): + """A FastMCP server with the chDB-only tools registered against the session.""" + executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) + mcp = FastMCP(name="chdb-test") + register_chdb_only_tools( + mcp, + max_result_bytes=1024 * 1024, + create_client=lambda: chdb_client, + query_executor=executor, + query_timeout=lambda: 30, + ) + yield mcp + executor.shutdown(wait=True) + + +async def test_allow_write_does_not_clobber_session(): + # Regression: registering with allow_write=True must NOT put the shared + # session under readonly=2 — that clobber is irreversible and would also + # strip write access from the legacy run_chdb_select_query on the same session. + session = chdb_session.Session() + executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + try: + register_chdb_only_tools( + FastMCP(name="chdb-write-test"), + max_result_bytes=1024 * 1024, + create_client=lambda: session, + query_executor=executor, + query_timeout=lambda: 30, + allow_write=True, + ) + # the session must still accept a write + session.query("CREATE TABLE wprobe (a Int32) ENGINE = Memory", "TabSeparated") + session.query("INSERT INTO wprobe VALUES (1)", "TabSeparated") + assert "1" in str(session.query("SELECT count() FROM wprobe", "CSV")) + finally: + executor.shutdown(wait=True) + session.close() + + +async def test_all_five_tools_are_registered(mcp_with_tools): + tools = await mcp_with_tools.get_tools() + assert { + "list_databases", + "list_tables", + "describe_table", + "get_sample_data", + "list_functions", + } <= set(tools) + + +async def test_list_databases(mcp_with_tools): + async with Client(mcp_with_tools) as client: + result = await client.call_tool("list_databases", {}) + assert "demo" in result.data + + +async def test_list_tables(mcp_with_tools): + async with Client(mcp_with_tools) as client: + result = await client.call_tool("list_tables", {"database": "demo"}) + assert "t" in result.data + + +async def test_describe_table(mcp_with_tools): + async with Client(mcp_with_tools) as client: + result = await client.call_tool("describe_table", {"database": "demo", "table": "t"}) + assert "id" in result.data and "Int32" in result.data + assert "name" in result.data and "String" in result.data + + +async def test_get_sample_data_respects_limit(mcp_with_tools): + async with Client(mcp_with_tools) as client: + result = await client.call_tool( + "get_sample_data", {"database": "demo", "table": "t", "limit": 2} + ) + # rows are returned as a JSON array; the limit bounds the returned count + rows = json.loads(result.data) + assert len(rows) == 2 + + +async def test_list_functions_filter(mcp_with_tools): + async with Client(mcp_with_tools) as client: + result = await client.call_tool("list_functions", {"pattern": "cosine"}) + assert "cosineDistance" in result.data + + +async def test_malicious_database_name_is_inert(mcp_with_tools): + async with Client(mcp_with_tools) as client: + # ChDBTool binds `database` as a value, so a "; DROP ..." name never + # reaches SQL as code: it matches no database (empty result) and cannot + # drop anything. + result = await client.call_tool( + "list_tables", {"database": "demo; DROP TABLE demo.t"} + ) + assert json.loads(result.data) == [] + # the real table is untouched + result2 = await client.call_tool("list_tables", {"database": "demo"}) + assert "t" in result2.data + + +async def test_bad_query_raises_tool_error(mcp_with_tools): + # errors are raised (like the ClickHouse-server tools), not returned as strings + async with Client(mcp_with_tools) as client: + with pytest.raises(Exception): + await client.call_tool( + "describe_table", {"database": "demo", "table": "nonexistent"} + ) diff --git a/tests/test_config_interface.py b/tests/test_config_interface.py index 9275d388..985b23b9 100644 --- a/tests/test_config_interface.py +++ b/tests/test_config_interface.py @@ -1,6 +1,6 @@ import pytest -from mcp_clickhouse.mcp_env import ClickHouseConfig +from mcp_clickhouse.mcp_env import ChDBConfig, ClickHouseConfig def test_interface_http_when_secure_false(monkeypatch: pytest.MonkeyPatch): @@ -120,3 +120,42 @@ def test_server_host_name_omitted_when_unset(monkeypatch: pytest.MonkeyPatch): client_config = config.get_client_config() assert "server_host_name" not in client_config + + +# --- ChDBConfig: chDB-only security-baseline settings ------------------------ + + +def test_chdb_allow_write_access_defaults_false(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("CHDB_ALLOW_WRITE_ACCESS", raising=False) + assert ChDBConfig().allow_write_access is False + + +def test_chdb_allow_write_access_true(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHDB_ALLOW_WRITE_ACCESS", "true") + assert ChDBConfig().allow_write_access is True + + +def test_chdb_max_result_bytes_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("CHDB_MAX_RESULT_BYTES", raising=False) + assert ChDBConfig().max_result_bytes == 1024 * 1024 + + +def test_chdb_max_result_bytes_override(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHDB_MAX_RESULT_BYTES", "2048") + assert ChDBConfig().max_result_bytes == 2048 + + +@pytest.mark.parametrize("bad", ["0", "-5", "notanint", ""]) +def test_chdb_max_result_bytes_falls_back_on_invalid(bad, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHDB_MAX_RESULT_BYTES", bad) + assert ChDBConfig().max_result_bytes == 1024 * 1024 + + +def test_chdb_file_allowlist_empty_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("CHDB_FILE_ALLOWLIST", raising=False) + assert ChDBConfig().file_allowlist == () + + +def test_chdb_file_allowlist_parses_colon_separated(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHDB_FILE_ALLOWLIST", "/data: /tmp/x :") + assert ChDBConfig().file_allowlist == ("/data", "/tmp/x") diff --git a/uv.lock b/uv.lock index cb04a93b..c96e279e 100644 --- a/uv.lock +++ b/uv.lock @@ -202,7 +202,7 @@ wheels = [ [[package]] name = "chdb" -version = "4.1.6" +version = "4.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "chdb-core" }, @@ -211,12 +211,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/61/8b/b7ec9eff8795a87f462c4b304e82391ecd6e4b238607f130557294a9ae6f/chdb-4.1.6-py3-none-any.whl", hash = "sha256:9acf64920697a9a88608c2c740d5ad5c650d1f75a47af089bd78ae5b6fe40f30", size = 1444649, upload-time = "2026-03-19T15:27:04.271Z" }, + { url = "https://files.pythonhosted.org/packages/91/88/85d2dbe2151d69a759f57f4e5d11d55c506e758b0e8fd0a2fa8d4e982c55/chdb-4.2.1-py3-none-any.whl", hash = "sha256:b5ae04ed5e1d6b51acfe7e2d6efbb176027ef09c3ed2255c6ca439a077a4aa1e", size = 1616586, upload-time = "2026-07-13T07:53:25.938Z" }, ] [[package]] name = "chdb-core" -version = "26.1.0" +version = "26.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -224,10 +224,10 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/e3/43/31e699cb602b8348531b2e36a57c70b7ad9238f6aaeed01dc38242524b34/chdb_core-26.1.0-cp39-abi3-macosx_10_15_x86_64.whl", hash = "sha256:7dd4fc1d008e60e1fb54440bc9a7c41996413f5f164c324acc790d47043264a5", size = 105204496, upload-time = "2026-03-02T13:20:46.813Z" }, - { url = "https://files.pythonhosted.org/packages/b1/6a/872464cbffb7f345643751c4f670d0844aef16b2e82ae3f5ceb95790399e/chdb_core-26.1.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:73be929393fd33136e550e99a5101e3068fc7812dfbb93eacbba9603b9a8e77a", size = 87764655, upload-time = "2026-03-02T12:45:07.585Z" }, - { url = "https://files.pythonhosted.org/packages/11/a9/649b956846bac60b91052583f84077a97bd9cc7dfb164178a4222f92b997/chdb_core-26.1.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a62b85061ec81ca64b74f418edf4eaee6ff82aa915b0d537c7ccb3a3e2ed0e9d", size = 118355514, upload-time = "2026-03-02T13:23:53.088Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b1/266c91ca0cd9c5259dfb8f18efd507ba53611bab10bffc10083b97ff66bf/chdb_core-26.1.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ff3559d024c10a31c5041ab4387e0b7d6a3b5d5bbf6dda884c5a36d4dc50ce47", size = 158111115, upload-time = "2026-03-02T13:16:14.098Z" }, + { url = "https://files.pythonhosted.org/packages/10/64/584953f24530357270bf16747cc9e0b5136d002011eb3713cded079a65e1/chdb_core-26.5.0-cp39-abi3-macosx_10_15_x86_64.whl", hash = "sha256:a4b896feb536b75b0cae0612410e3cc93e8161d41eab9bd55720303aef1aa047", size = 110710168, upload-time = "2026-06-08T03:33:44.677Z" }, + { url = "https://files.pythonhosted.org/packages/63/f4/04fc1be6d5576aa04a22faa168351431c3689fde1b1ac51c54eca2af895f/chdb_core-26.5.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:4902bdf2e58e49567dbfcb8d9b88607edf6d472f6c21f33c4674574e0444bd28", size = 96683215, upload-time = "2026-06-08T01:52:15.76Z" }, + { url = "https://files.pythonhosted.org/packages/ce/08/d97a17a69d54ab21b636c7de668a94e0413d3e8fe5804ffbd2e38996a1e6/chdb_core-26.5.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0004a5528287d3eec757a223e851a3b68a485675881a09bb4fa7b5116035b3c7", size = 122690139, upload-time = "2026-06-08T03:17:08.205Z" }, + { url = "https://files.pythonhosted.org/packages/a5/2d/eb9b51180f2c0934339cf94ae2c031fab7c3dc93191f95aff0f83a0ca741/chdb_core-26.5.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:85ddbf675b4a070aa5abb79456b0a08643273dcf3977dc39f347c8acf846c804", size = 157311782, upload-time = "2026-06-08T02:32:13.171Z" }, ] [[package]] @@ -891,7 +891,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "cachetools", specifier = ">=5.5.0" }, - { name = "chdb", marker = "extra == 'chdb'", specifier = ">=3.3.0" }, + { name = "chdb", marker = "extra == 'chdb'", specifier = ">=4.2.1" }, { name = "clickhouse-connect", specifier = ">=0.8.16" }, { name = "fastmcp", specifier = ">=2.0.0,<3.0.0" }, { name = "pytest", marker = "extra == 'dev'" }, @@ -2046,8 +2046,8 @@ name = "taskgroup" version = "0.2.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "exceptiongroup", marker = "python_full_version < '3.14'" }, - { name = "typing-extensions", marker = "python_full_version < '3.14'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f0/8d/e218e0160cc1b692e6e0e5ba34e8865dbb171efeb5fc9a704544b3020605/taskgroup-0.2.2.tar.gz", hash = "sha256:078483ac3e78f2e3f973e2edbf6941374fbea81b9c5d0a96f51d297717f4752d", size = 11504, upload-time = "2025-01-03T09:24:13.761Z" } wheels = [