diff --git a/pyproject.toml b/pyproject.toml index 2b108b1..cfdcefe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "credential-demo", "hello", "stac-explorer", + "duckdb-analyst", ] [dependency-groups] @@ -41,6 +42,7 @@ members = ["toolsets/*"] credential-demo = { workspace = true } hello = { workspace = true } stac-explorer = { workspace = true } +duckdb-analyst = { workspace = true } [tool.ruff] line-length = 88 diff --git a/toolsets/duckdb-analyst/pyproject.toml b/toolsets/duckdb-analyst/pyproject.toml new file mode 100644 index 0000000..4c0d295 --- /dev/null +++ b/toolsets/duckdb-analyst/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "duckdb-analyst" +version = "0.1.0" +description = "Read-only, sandbox-free SQL analysis over public parquet datasets (Natural Earth, plus any public parquet/CSV URL) via DuckDB." +requires-python = ">=3.12,<3.14" +dependencies = [ + "duckdb>=1.5.5,<2.0.0", + "langchain-core<2.0.0,>=1.4.6", + "mcp-toolsets-runtime", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/__init__.py b/toolsets/duckdb-analyst/src/duckdb_analyst/__init__.py new file mode 100644 index 0000000..b846975 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/__init__.py @@ -0,0 +1 @@ +"""duckdb-analyst toolset.""" diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py b/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py new file mode 100644 index 0000000..6207c32 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py @@ -0,0 +1,339 @@ +"""The one DuckDB connection this toolset ever opens, and its security model. + +Read this before touching ``tools.py``. The connection is built once, at +import time, in :func:`_build_connection`: extensions are installed, the +curated views are created, and then the connection is locked down before any +caller-supplied SQL ever runs against it. ``query``/``chart`` in ``tools.py`` +only ever get a ``cursor()`` off the already-locked-down :data:`CON`. + +SECURITY MODEL — four layers, in decreasing order of how much they actually +protect this connection from a hostile ``SELECT``: + +1. **The deployment has no local secrets reachable, full stop.** This is the + PRIMARY control, and it lives outside this file entirely: nothing in + Python can compensate for a `.env`, a mounted Secret, or an ambient AWS + credential sitting in this container's filesystem or environment. This + toolset's ``pyproject.toml`` declares no ``CREDENTIAL_HEADERS``, needs no + secrets in ``toolset.yaml``, and the shared ``Dockerfile`` ships a bare + ``python:3.12-slim-bookworm`` runtime image with nothing baked in beyond + the venv. If that ever changes for this toolset specifically, this + connection's threat model changes with it — grep for "PRIMARY control" + before adding any credential to this toolset's deployment. + +2. **DuckDB is configured so its own filesystem access can't reach local + disk while ``https://``/``s3://`` reads keep working.** DuckDB's + documented, all-or-nothing knob (``SET enable_external_access = false``) + was tried first and rejected: it disables ``read_csv``/``read_parquet``/ + ``read_json`` reading from *any* external source, remote included — see + https://duckdb.org/docs/stable/operations_manual/securing_duckdb/overview. + ``SET disabled_filesystems = 'LocalFileSystem'`` is DuckDB's documented, + finer-grained alternative (it's literally their own worked example for + blocking ``read_csv('/etc/passwd', ...)``), and empirically it *does* + deliver local-blocked/remote-allowed — but only once the connection's + remote filesystem path has been exercised at least once beforehand. Prior + to that first remote read, DuckDB's own httpfs/spatial code paths still + touch local disk internally (e.g. to resolve/cache a filesystem handle), + so a `disabled_filesystems` set on a "cold" connection blocks the first + remote read too — a real DuckDB limitation, tracked upstream as + https://github.com/duckdb/duckdb/issues/15734 for the plain-httpfs case, + and reproduced here for ``spatial``'s ``ST_Read`` as well. Verified + directly against this toolset's pinned DuckDB version (see the + ``dependencies`` pin in ``pyproject.toml``): one warmup read of *any* + remote URL, of *any* kind (plain https, s3://, or ``ST_Read``), before + locking down, is sufficient — every remote kind then keeps working for + arbitrary new URLs afterwards, while local paths stay blocked. That is + exactly what ``_warm_up_remote_filesystem`` below does, and it runs + before ``_register_views`` for good measure (creating each view already + touches its remote source once, but the explicit warmup makes the + ordering requirement self-documenting instead of incidental). This is a + real, verified technical control, not a fig leaf — but it rests on an + *undocumented* DuckDB behavior (one-time-touch-then-cached), which is one + more reason layer 1 is the control this toolset actually depends on, not + this one. +3. ``SET lock_configuration = true`` is set last, after every setting above, + so no runtime SQL — even something that slips past ``security.py``'s + statement-shape filter — can loosen any of it. ``SET + allow_community_extensions = false`` runs first, before any ``INSTALL``: + this toolset needs only core extensions (``httpfs``, ``spatial``), so + community extensions are refused outright rather than allowed once and + then closed off. +4. ``security.validate_select_only`` (defense-in-depth, checked before any + caller SQL reaches this connection): single-statement, ``SELECT``/``WITH`` + only, plus a small denylist for introspection functions + (``duckdb_secrets()`` etc.) that are otherwise valid inside a plain + ``SELECT``. This does **not** stop ``read_text('/etc/passwd')`` — that + starts with ``SELECT`` and is one statement, so it sails through this + check. It's stopped by layer 2 (and, ultimately, guaranteed by layer 1). + +Extensions are pinned by name (``httpfs``, ``spatial``) — nothing under +``duckdb-analyst`` ever calls ``INSTALL``/``LOAD`` again after connection +setup, and layer 3 makes sure caller SQL couldn't anyway. Both are core +DuckDB extensions, so their versions are pinned transitively by this +toolset's pinned ``duckdb`` dependency. No community extension is used, so +there is no unpinnable dependency here. + +SCOPE — what this toolset is for: + +Small, fast datasets read over the network, plus ad hoc ``https://``/ +``s3://`` parquet and CSV URLs. The curated views are deliberately +lightweight (Natural Earth 1:110m, a few hundred KB each), so a full +``GROUP BY`` over one of them returns in well under a second. + +Large remote datasets are out of scope, and that is a measured limit rather +than an untested guess. Overture Maps was tried here first and removed: its +themes are hundreds of GB, the public parquet layout gives no partition +pruning for the fields an analyst actually filters on (``country``, +``locality``), and so any aggregate query scans the whole theme. Measured +against ``_MEMORY_LIMIT``/``_THREADS`` below and the 30s watchdog in +``security.py``: a bare ``SELECT ... LIMIT 3`` took ~2.8s, adding ``WHERE +country = 'PT'`` took ~17.4s, and a ``GROUP BY`` exceeded the timeout +outright. A dataset that size wants a local copy (this is the same +conclusion ``gazet`` reached, which is why it ships local Overture +extracts), not a remote scan — and a local copy is just another +``CREATE VIEW`` here, needing no change to ``tools.py`` or the security +model. +""" + +from typing import NotRequired, TypedDict + +import duckdb + +# Official Natural Earth vector data, unzipped, served directly over plain +# https by the project's own GitHub organization +# (https://github.com/nvkelso/natural-earth-vector — nvkelso is a Natural +# Earth maintainer; this is the canonical vector-data repo, not a random +# mirror). Deliberately NOT the zipped shapefiles naturalearthdata.com links +# to: this DuckDB build's `spatial` extension has no working GDAL vsicurl +# support (verified: `/vsicurl/...` and `/vsizip/vsicurl/...` both fail to +# open, on or off the lockdown below), but `ST_Read` handles a plain +# `https://` URL itself via DuckDB's own httpfs filesystem — the same code +# path `read_parquet`/`read_csv` use, and the one the warmup below exercises. +_NATURAL_EARTH_COUNTRIES_URL = ( + "https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/" + "geojson/ne_110m_admin_0_countries.geojson" +) +_NATURAL_EARTH_PLACES_URL = ( + "https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/" + "geojson/ne_110m_populated_places.geojson" +) + +#: Conservative resource caps for a shared, always-on connection. Sized +#: against this toolset's `toolset.yaml` pod memory limit, with headroom for +#: Python/uvicorn overhead alongside DuckDB itself. +_MEMORY_LIMIT = "700MB" +_THREADS = 2 + + +class ColumnInfo(TypedDict): + """One column of a curated source, as advertised by ``list_sources``.""" + + name: str + type: str + description: NotRequired[str] + #: Vega-Lite encoding channels this column tends to work well as. + good_for: NotRequired[list[str]] + + +class SourceInfo(TypedDict): + """One pre-registered source, as advertised by ``list_sources``.""" + + name: str + kind: str # "view" (curated, fixed schema) | "table_function" (SQL fn) + description: str + example_sql: str + columns: NotRequired[list[ColumnInfo]] + + +def _build_connection() -> duckdb.DuckDBPyConnection: + con = duckdb.connect(":memory:") + con.execute(f"SET memory_limit = '{_MEMORY_LIMIT}'") + con.execute(f"SET threads = {_THREADS}") + + # Core extensions only, and community extensions refused before the first + # INSTALL — see the module docstring, layer "Extensions". + con.execute("SET allow_community_extensions = false") + con.execute("INSTALL httpfs") + con.execute("LOAD httpfs") + con.execute("INSTALL spatial") + con.execute("LOAD spatial") + + # Default region for ad hoc `s3://` reads. Public/anonymous buckets need + # no credentials, just a region. NOTE: `lock_configuration` below means a + # caller cannot change this, so an `s3://` URL in a different region will + # fail — use that bucket's `https://` endpoint instead, or add a region + # here. Plain `https://` reads are unaffected. + con.execute("SET s3_region = 'us-east-1'") + + _warm_up_remote_filesystem(con) + _register_views(con) + + # Lock down. See the module docstring, layer "DuckDB is configured...", + # for why this order (and the warmup above) is load-bearing. + con.execute("SET disabled_filesystems = 'LocalFileSystem'") + con.execute("SET lock_configuration = true") + return con + + +def _warm_up_remote_filesystem(con: duckdb.DuckDBPyConnection) -> None: + """Touch DuckDB's remote filesystem path once before locking local disk. + + See the module docstring for why this is necessary and what it's based + on. Any tiny, reliably-reachable remote read does the job; this one is + deliberately unrelated to the views below so the requirement doesn't + look like an accident of registration order. + """ + con.execute( + "SELECT 1 FROM read_parquet('https://duckdb.org/data/holdings.parquet') LIMIT 1" + ) + + +def _register_views(con: duckdb.DuckDBPyConnection) -> None: + # Each view's SQL is built as a plain string first and executed on its + # own line, rather than inlining `con.execute(f"""..."""), so + # the linter-suppression comment below has an unambiguous line to + # attach to instead of landing inside the SQL string itself. The + # underlying ruff finding ("possible SQL injection") is a false + # positive: the only interpolated values here are module constants + # (hardcoded dataset URLs) — never caller input, which is what actually + # makes string-built SQL dangerous. + countries_sql = f""" + CREATE VIEW natural_earth_countries AS + SELECT + "NAME" AS name, + "ADMIN" AS admin_name, + "ISO_A2" AS iso_a2, + "ISO_A3" AS iso_a3, + "CONTINENT" AS continent, + "SUBREGION" AS subregion, + "POP_EST" AS population_estimate, + "GDP_MD" AS gdp_million_usd, + "INCOME_GRP" AS income_group, + geom AS geometry + FROM ST_Read('{_NATURAL_EARTH_COUNTRIES_URL}') + """ # noqa: S608 + con.execute(countries_sql) + + places_ne_sql = f""" + CREATE VIEW natural_earth_places AS + SELECT + "NAME" AS name, + "ADM0NAME" AS country_name, + "ISO_A2" AS iso_a2, + "LATITUDE" AS latitude, + "LONGITUDE" AS longitude, + "POP_MAX" AS population_max, + "TIMEZONE" AS timezone, + geom AS geometry + FROM ST_Read('{_NATURAL_EARTH_PLACES_URL}') + """ # noqa: S608 + con.execute(places_ne_sql) + + +# Hand-authored notes on top of each view's real, live-introspected schema +# (see _build_sources): the things a schema dump can't tell an LLM, like +# which columns make good chart channels. A column with no entry here still +# appears in list_sources with just its name/type. +_COLUMN_NOTES: dict[str, dict[str, ColumnInfo]] = { + "natural_earth_countries": { + "name": { + "name": "name", + "type": "", + "description": "Country name.", + "good_for": ["x", "color"], + }, + "continent": { + "name": "continent", + "type": "", + "good_for": ["x", "color"], + }, + "population_estimate": { + "name": "population_estimate", + "type": "", + "good_for": ["y", "color"], + }, + "gdp_million_usd": { + "name": "gdp_million_usd", + "type": "", + "description": "GDP estimate, millions of USD.", + "good_for": ["y", "color"], + }, + "income_group": { + "name": "income_group", + "type": "", + "good_for": ["x", "color"], + }, + }, + "natural_earth_places": { + "name": { + "name": "name", + "type": "", + "description": "City/place name.", + "good_for": ["color"], + }, + "latitude": {"name": "latitude", "type": "", "good_for": ["y"]}, + "longitude": {"name": "longitude", "type": "", "good_for": ["x"]}, + "population_max": { + "name": "population_max", + "type": "", + "description": "High estimate of the place's population.", + "good_for": ["y", "color"], + }, + }, +} + +_VIEW_DESCRIPTIONS = { + "natural_earth_countries": ( + "Natural Earth 1:110m country polygons with basic demographic/" + "economic attributes (name, continent, population, GDP, income " + "group). Small and fast — a full GROUP BY over it returns in well " + "under a second. Good for country- and continent-level breakdowns." + ), + "natural_earth_places": ( + "Natural Earth 1:110m populated places (major world cities) as " + "points with lat/lon and population — a small, fast dataset good " + "for quick geographic scatter charts." + ), +} + +_EXAMPLE_SQL = { + "natural_earth_countries": ( + "SELECT name, continent, population_estimate FROM " + "natural_earth_countries ORDER BY population_estimate DESC LIMIT 20" + ), + "natural_earth_places": ( + "SELECT name, longitude, latitude, population_max FROM " + "natural_earth_places ORDER BY population_max DESC LIMIT 50" + ), +} + + +def _describe_view(con: duckdb.DuckDBPyConnection, view: str) -> list[ColumnInfo]: + """Live column name/type for a view, merged with any curated notes. + + Reads the real, already-registered view's schema rather than + hand-listing columns, so this catalog can't drift from what the views + actually return. + """ + notes = _COLUMN_NOTES.get(view, {}) + columns: list[ColumnInfo] = [] + for name, duck_type, *_ in con.execute(f"DESCRIBE {view}").fetchall(): + note = notes.get(name, ColumnInfo(name=name, type="")) + columns.append(ColumnInfo(**{**note, "name": name, "type": str(duck_type)})) + return columns + + +def _build_sources(con: duckdb.DuckDBPyConnection) -> list[SourceInfo]: + return [ + SourceInfo( + name=view, + kind="view", + description=_VIEW_DESCRIPTIONS[view], + example_sql=_EXAMPLE_SQL[view], + columns=_describe_view(con, view), + ) + for view in ("natural_earth_countries", "natural_earth_places") + ] + + +CON = _build_connection() +SOURCES = _build_sources(CON) diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/security.py b/toolsets/duckdb-analyst/src/duckdb_analyst/security.py new file mode 100644 index 0000000..9a85bd1 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/security.py @@ -0,0 +1,73 @@ +"""Statement-shape validation for caller-supplied SQL. + +This is defense-in-depth, not the primary control — read +``connection.py``'s module docstring first for the actual security model. +In particular: none of this stops a file-reading table function +(``read_text``, ``read_csv``, ``glob``, ...) called from inside an otherwise +valid ``SELECT``. That class of risk is handled by ``connection.py`` locking +down DuckDB's filesystem access and by the deployment shipping no local +secrets — not by anything in this module. +""" + +import re + +#: Server-enforced row cap. ``limit`` is clamped into (0, MAX_ROW_LIMIT] before +#: it ever reaches SQL, so a caller cannot bypass it by e.g. omitting a LIMIT +#: clause of their own — see ``connection.wrap_with_limit``. +MAX_ROW_LIMIT = 10_000 +DEFAULT_ROW_LIMIT = 1000 + +#: Wall-clock budget for one query, enforced by a Python-side watchdog that +#: calls ``cursor.interrupt()`` — DuckDB has no native query timeout. +QUERY_TIMEOUT_SECONDS = 30.0 + +_LEADING_STATEMENT = re.compile(r"^\s*(SELECT|WITH)\b", re.IGNORECASE) + +#: Table/scalar functions that are reachable from inside an otherwise valid, +#: single-statement SELECT and disclose runtime configuration or credential +#: state — the "starts with SELECT, one statement" shape check below does not +#: see these, so they need an explicit denylist. Not exhaustive; DuckDB's +#: attack surface here is "whatever a future DuckDB version adds a +#: SELECT-callable introspection function for", so this list is reviewed, not +#: assumed complete. +_DENYLISTED_CALLS = re.compile( + r"\b(duckdb_secrets|duckdb_settings|pragma_[a-z_]*)\s*\(", re.IGNORECASE +) + + +def clamp_limit(limit: int) -> int: + """Clamp a caller-supplied row limit into ``(0, MAX_ROW_LIMIT]``.""" + return max(1, min(limit, MAX_ROW_LIMIT)) + + +def validate_select_only(sql: str) -> str | None: + """Reject anything that isn't a single ``SELECT``/``WITH`` statement. + + Returns an error detail string if ``sql`` is rejected, or ``None`` if it + passes this (shallow, syntax-level) check. Rejects: + + - more than one statement (a ``;`` anywhere but a single optional + trailing one) — defense against ``SELECT 1; ATTACH ...``-style + smuggling of a second statement; + - anything not starting with ``SELECT``/``WITH`` — defense against + ``ATTACH``/``COPY``/``INSTALL``/``PRAGMA``/``CALL``/``SET`` and + friends; + - calls to a small denylist of introspection functions (see + ``_DENYLISTED_CALLS``) that are otherwise perfectly valid inside a + single ``SELECT``. + """ + stripped = sql.strip() + if not stripped: + return "empty query" + + body = stripped[:-1] if stripped.endswith(";") else stripped + if ";" in body: + return "only a single statement is allowed (found an embedded ';')" + + if not _LEADING_STATEMENT.match(body): + return "only SELECT/WITH statements are allowed" + + if match := _DENYLISTED_CALLS.search(body): + return f"{match.group(1)}() is not allowed" + + return None diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py b/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py new file mode 100644 index 0000000..bb1d7c1 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py @@ -0,0 +1,191 @@ +"""LangChain tools for duckdb-analyst: sandbox-free SQL analysis over public +parquet datasets via a security-hardened, read-only DuckDB connection. + +Three tools, deliberately minimal: explore what's available (`list_sources`), +run a SELECT (`query`), or run one and drop the rows straight into a +caller-supplied Vega-Lite spec (`chart`). No bespoke per-dataset tool — a new +dataset is a `CREATE VIEW` in `connection.py`, not a fourth tool, and it +becomes reachable through `query`/`chart` and discoverable via `list_sources` +(see `connection.SOURCES`). + +The actual security model (why arbitrary SQL is safe to run here) lives in +`connection.py`'s module docstring — read that before changing anything here. +This module's own job is narrower: reject obviously-wrong SQL early +(`security.validate_select_only`), enforce the row cap server-side, bound +execution time, and turn DuckDB's own errors into `ToolError`s instead of +raising. +""" + +import asyncio +import json +from typing import Any, NotRequired + +import duckdb +from langchain_core.tools import tool + +from mcp_runtime.tool_result import ToolError, ToolResult + +from duckdb_analyst.connection import CON, SOURCES, SourceInfo +from duckdb_analyst.security import ( + QUERY_TIMEOUT_SECONDS, + clamp_limit, + validate_select_only, +) + + +class ListSourcesResult(ToolResult): + """The pre-registered views/table functions available to query/chart.""" + + sources: NotRequired[list[SourceInfo]] + + +class QueryResult(ToolResult): + """Rows from a validated, capped SELECT against the DuckDB connection.""" + + rows: NotRequired[list[dict[str, Any]]] + row_count: NotRequired[int] + + +class ChartResult(ToolResult): + """A caller-supplied Vega-Lite mark/encoding spec with data filled in.""" + + spec: NotRequired[dict[str, Any]] + + +def _json_safe(value: Any) -> Any: + """Round-trip a DuckDB row through JSON to flatten it into plain types. + + DuckDB hands back Decimals, dates/times, raw WKB bytes for geometry + columns, and nested struct/list values as Python dicts/lists — none of + that is guaranteed JSON-native. Rather than hand-writing a type-by-type + converter (and inevitably missing one), `default=str` is the catch-all: + anything `json.dumps` doesn't already know becomes its `str()`. That + makes e.g. a raw geometry blob a somewhat ugly `"b'\\x01\\x03...'"` + string rather than a serialization crash — `list_sources` tells callers + to wrap geometry columns in `ST_AsText`/`ST_AsGeoJSON` instead of relying + on this fallback for anything they actually want to read. + """ + return json.loads(json.dumps(value, default=str)) + + +def _wrap_with_limit(sql: str, limit: int) -> str: + """Force a hard row cap at the SQL level, not just a Python-side slice. + + Wrapping in a subquery caps rows *inside* DuckDB regardless of whether + the caller's own SQL has a LIMIT — a caller can't get more rows back by + simply omitting one. It also means the top-level statement DuckDB + actually executes is always a SELECT ... FROM (...), which is one more + reason (beyond `validate_select_only`) that stray non-SELECT SQL can't + reach execution here. + """ + # Not a real injection vector, despite how this looks to a linter: `sql` + # has already passed `validate_select_only` (single SELECT/WITH + # statement) by the time this runs, and `limit` is an int already + # clamped by `clamp_limit` — never a caller-controlled string spliced in + # verbatim. + return f"SELECT * FROM (\n{sql}\n) AS _duckdb_analyst_query\nLIMIT {limit}" # noqa: S608 + + +async def _run_query( + sql: str, limit: int +) -> tuple[list[str], list[dict[str, Any]]] | ToolError: + """Validate, execute and fetch a capped SELECT, or return a ToolError. + + Runs on a fresh `CON.cursor()` (cheap; shares the already-locked-down + in-memory database, safe to use concurrently with other calls) inside + `asyncio.to_thread`, guarded by a watchdog that interrupts the cursor + past `QUERY_TIMEOUT_SECONDS` — DuckDB has no native query timeout. + """ + if detail := validate_select_only(sql): + return ToolError(error="invalid_query", detail=detail) + + capped_limit = clamp_limit(limit) + executable = _wrap_with_limit(sql, capped_limit) + cursor = CON.cursor() + finished = asyncio.Event() + + async def watchdog() -> None: + try: + await asyncio.wait_for(finished.wait(), timeout=QUERY_TIMEOUT_SECONDS) + except TimeoutError: + cursor.interrupt() + + watchdog_task = asyncio.create_task(watchdog()) + try: + result = await asyncio.to_thread(cursor.execute, executable) + rows = await asyncio.to_thread(result.fetchall) + columns = [description[0] for description in result.description] + except duckdb.Error as error: + kind = ( + "timeout" + if isinstance(error, duckdb.InterruptException) + else "query_failed" + ) + return ToolError(error=kind, detail=str(error)) + finally: + finished.set() + await watchdog_task + + records = [_json_safe(dict(zip(columns, row, strict=True))) for row in rows] + return columns, records + + +@tool +def list_sources() -> ListSourcesResult: + """List the pre-registered views available to `query`/`chart`: Natural + Earth countries and populated places — with column descriptions and which + columns work well as x/y/color chart channels. Call this before writing + SQL against an unfamiliar source. + + `query`/`chart` are not limited to these: they also read any public + `https://` or `s3://` parquet or CSV URL via `read_parquet`/`read_csv`. + """ + names = ", ".join(source["name"] for source in SOURCES) + return ListSourcesResult( + message=f"{len(SOURCES)} source(s) available: {names}.", sources=SOURCES + ) + + +@tool +async def query(sql: str, limit: int = 1000) -> QueryResult | ToolError: + """Run a read-only SQL SELECT against the curated views from + `list_sources`, or an ad hoc `https://`/`s3://` parquet/CSV URL via + `read_parquet`/`read_csv`, and return the rows as JSON records. + + Only a single `SELECT`/`WITH` statement is allowed. `limit` caps the + rows returned (default 1000, hard max 10000), enforced regardless of + what the query itself requests. + """ + outcome = await _run_query(sql, limit) + if isinstance(outcome, dict): # ToolError + return outcome + _columns, rows = outcome + return QueryResult( + message=f"{len(rows)} row(s) returned.", rows=rows, row_count=len(rows) + ) + + +@tool +async def chart( + sql: str, spec: dict[str, Any], limit: int = 1000 +) -> ChartResult | ToolError: + """Run a SQL query the same way `query` does, then return a completed + Vega-Lite spec: `spec` (your `mark`/`encoding`, no `data` key) with the + query's rows filled in as `spec.data.values`. Rendering is left to + whatever client receives the result — this tool only assembles the spec. + + `limit` caps the rows inlined into the chart (default 1000, hard max + 10000) — keep it small; a Vega-Lite spec with tens of thousands of + inlined rows is unwieldy for most renderers. + """ + outcome = await _run_query(sql, limit) + if isinstance(outcome, dict): # ToolError + return outcome + _columns, rows = outcome + full_spec = {**spec, "data": {"values": rows}} + return ChartResult( + message=f"Chart spec built with {len(rows)} row(s) of data.", spec=full_spec + ) + + +TOOLS = [list_sources, query, chart] diff --git a/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py b/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py new file mode 100644 index 0000000..2b4a609 --- /dev/null +++ b/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py @@ -0,0 +1,219 @@ +"""Tests for duckdb-analyst. + +These hit the real network — DuckDB's own httpfs/spatial extensions have no +mockable transport the way `httpx.MockTransport` covers `stac-explorer`. +Fixtures are picked to be small and stable: a single ad hoc parquet file +DuckDB's own docs use as a demo (~500 bytes), and the curated Natural Earth +1:110m views (a few hundred KB each, so a full aggregate over one is still +sub-second). + +Every curated source this toolset advertises is exercised here with a real +query. That is deliberate: an earlier revision advertised Overture Maps and +a STAC table function in `list_sources` without covering either with a +query test, and both turned out to be unusable in practice (Overture +aggregates blew the 30s timeout; STAC search returned HTTP 422). If you add +a source to `connection.py`, add a query test for it here, or it does not +ship. + +The security regression suite is the important part: every case there must +come back as a `ToolError`, never rows. +""" + +from mcp_runtime.tool_result import is_error + +from duckdb_analyst.tools import chart, list_sources, query + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +def test_list_sources_covers_every_dataset_family(): + result = list_sources.invoke({}) + names = {source["name"] for source in result["sources"]} + assert {"natural_earth_countries", "natural_earth_places"} <= names + for source in result["sources"]: + assert source["description"] + assert source["example_sql"] + + +def test_list_sources_advertises_nothing_untested(): + """Guard the docstring's rule: no advertised source without a query test. + + `list_sources` is what an LLM reads to decide what it can query, so an + entry here that no test exercises is how the Overture/STAC regression + got shipped in the first place. + """ + tested = {"natural_earth_countries", "natural_earth_places"} + advertised = {source["name"] for source in list_sources.invoke({})["sources"]} + assert advertised == tested, ( + f"sources advertised but not query-tested: {advertised - tested}" + ) + + +def test_list_sources_flags_chart_friendly_columns(): + result = list_sources.invoke({}) + countries = next( + s for s in result["sources"] if s["name"] == "natural_earth_countries" + ) + continent = next(c for c in countries["columns"] if c["name"] == "continent") + assert "x" in continent["good_for"] or "color" in continent["good_for"] + + +async def test_query_against_curated_view(): + result = await query.ainvoke( + { + "sql": ( + "SELECT name, continent, population_estimate " + "FROM natural_earth_countries ORDER BY population_estimate DESC" + ), + "limit": 5, + } + ) + assert not is_error(result) + assert 0 < len(result["rows"]) <= 5 + assert result["row_count"] == len(result["rows"]) + assert "name" in result["rows"][0] + + +async def test_query_supports_full_aggregate_over_curated_view(): + """A whole-view GROUP BY must finish well inside the query timeout. + + This is the property that makes the curated sources worth curating — + the Overture views this replaced could not do it. + """ + result = await query.ainvoke( + { + "sql": ( + "SELECT continent, COUNT(*) AS n_countries, " + "SUM(population_estimate) AS population " + "FROM natural_earth_countries GROUP BY continent " + "ORDER BY population DESC" + ) + } + ) + assert not is_error(result) + assert len(result["rows"]) > 1 + assert result["rows"][0]["population"] > 0 + + +async def test_query_against_second_curated_view(): + result = await query.ainvoke( + { + "sql": ( + "SELECT name, country_name, population_max, longitude, latitude " + "FROM natural_earth_places ORDER BY population_max DESC" + ), + "limit": 3, + } + ) + assert not is_error(result) + assert len(result["rows"]) == 3 + assert result["rows"][0]["population_max"] >= result["rows"][-1]["population_max"] + + +async def test_query_can_use_spatial_functions_on_geometry(): + """`spatial` is loaded, so ST_* works on the views' geometry column.""" + result = await query.ainvoke( + { + "sql": ( + "SELECT name, ROUND(ST_Area(geometry), 2) AS area_deg2 " + "FROM natural_earth_countries WHERE name = 'Brazil'" + ) + } + ) + assert not is_error(result) + assert result["rows"][0]["area_deg2"] > 0 + + +async def test_query_against_ad_hoc_public_parquet_url(): + result = await query.ainvoke( + { + "sql": "SELECT * FROM read_parquet('https://duckdb.org/data/holdings.parquet')" + } + ) + assert not is_error(result) + assert len(result["rows"]) > 0 + + +async def test_query_enforces_hard_row_cap_server_side(): + result = await query.ainvoke( + { + "sql": "SELECT * FROM range(100000) AS t(n)", + "limit": 1_000_000, # above MAX_ROW_LIMIT + } + ) + assert not is_error(result) + assert result["row_count"] <= 10_000 + + +async def test_chart_fills_in_data_values_and_preserves_spec(): + spec = { + "mark": "bar", + "encoding": { + "x": {"field": "continent", "type": "nominal", "sort": "-y"}, + "y": {"field": "population", "type": "quantitative"}, + }, + } + result = await chart.ainvoke( + { + "sql": ( + "SELECT continent, SUM(population_estimate) AS population " + "FROM natural_earth_countries GROUP BY continent " + "ORDER BY population DESC" + ), + "spec": spec, + "limit": 10, + } + ) + assert not is_error(result) + assert result["spec"]["mark"] == "bar" + assert result["spec"]["encoding"] == spec["encoding"] + values = result["spec"]["data"]["values"] + assert 0 < len(values) <= 10 + # The caller's encoding must line up with the columns the SQL returned, + # or the spec renders empty in the client. + assert {"continent", "population"} <= set(values[0]) + + +# --------------------------------------------------------------------------- +# Security regression suite: every one of these must be rejected, never +# silently succeed. +# --------------------------------------------------------------------------- + + +async def test_rejects_local_file_read_via_read_text_etc_passwd(): + result = await query.ainvoke({"sql": "SELECT * FROM read_text('/etc/passwd')"}) + assert is_error(result) + + +async def test_rejects_local_file_read_via_read_text_proc_environ(): + result = await query.ainvoke( + {"sql": "SELECT * FROM read_text('/proc/self/environ')"} + ) + assert is_error(result) + + +async def test_rejects_multiple_statements(): + result = await query.ainvoke({"sql": "SELECT 1; ATTACH ':memory:' AS x"}) + assert is_error(result) + + +async def test_rejects_install(): + result = await query.ainvoke({"sql": "INSTALL icu"}) + assert is_error(result) + + +async def test_rejects_attempt_to_loosen_locked_configuration(): + result = await query.ainvoke({"sql": "SET enable_external_access = true"}) + assert is_error(result) + + +async def test_rejects_pragma(): + result = await query.ainvoke({"sql": "PRAGMA database_list"}) + assert is_error(result) + + +async def test_rejects_duckdb_secrets(): + result = await query.ainvoke({"sql": "SELECT * FROM duckdb_secrets()"}) + assert is_error(result) diff --git a/toolsets/duckdb-analyst/toolset.yaml b/toolsets/duckdb-analyst/toolset.yaml new file mode 100644 index 0000000..82fc0d4 --- /dev/null +++ b/toolsets/duckdb-analyst/toolset.yaml @@ -0,0 +1,13 @@ +# Helm values overrides for this toolset (see charts/mcp-toolset/values.yaml). +# +# DuckDB needs more headroom than the chart's 512Mi default even for +# capped, LIMIT-bounded queries over remote parquet — the connection's own +# `memory_limit` (see connection.py) is set to 700MB, sized against this pod +# limit with room left for Python/uvicorn overhead. +resources: + requests: + cpu: 100m + memory: 256Mi + limits: + cpu: "1" + memory: 1Gi diff --git a/uv.lock b/uv.lock index c54f472..b8907a1 100644 --- a/uv.lock +++ b/uv.lock @@ -9,6 +9,7 @@ resolution-markers = [ [manifest] members = [ "credential-demo", + "duckdb-analyst", "hello", "mcp-toolsets", "stac-explorer", @@ -442,6 +443,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/d0/205d54408c08b13550c733c4b85429e7ead111c7f0014309637425520a9a/deprecated-1.3.1-py2.py3-none-any.whl", hash = "sha256:597bfef186b6f60181535a29fbe44865ce137a5079f295b479886c82729d5f3f", size = 11298, upload-time = "2025-10-30T08:19:00.758Z" }, ] +[[package]] +name = "duckdb" +version = "1.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/19/e57151753576373c6696a12022648546cca6038e8833fda2908ee2342d9b/duckdb-1.5.5.tar.gz", hash = "sha256:72f33ee57ca7595b23957671a2cc7f7fe2be0ecc2d68f63abedcfcaa3a5c1238", size = 18066741, upload-time = "2026-07-22T10:55:17.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d6/40/2e05d324400fdaa5656c9f48d6298da421cb034d85e509fa0e6e325cf04b/duckdb-1.5.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d4dd65f8941a604b947e0b9b4b4f7165988e29a23ec0b69b4038520956d9933e", size = 32753858, upload-time = "2026-07-22T10:54:05.514Z" }, + { url = "https://files.pythonhosted.org/packages/79/15/5ceb58ffb5bb8a62b3fd7abb39c41467cdf94850ece02e6d88664dfc75ce/duckdb-1.5.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33db46679b071f108d57139493dee2d37e1f5efcf5c5c039c2969eed11a6c8a7", size = 17368293, upload-time = "2026-07-22T10:54:09.139Z" }, + { url = "https://files.pythonhosted.org/packages/bf/5c/bf02da0b354fe83cca4f95a4fbf762181af466f7d551ab2a093f7698882a/duckdb-1.5.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f0b88535a5d86fdd63dba6ea02ab68c003dfb9e4892b11256ef24c4da208baae", size = 15509131, upload-time = "2026-07-22T10:54:12.228Z" }, + { url = "https://files.pythonhosted.org/packages/ea/a9/5f1f09da421d8e930e0b063d11c1b3f90363f40ede74438cd188afdd13a2/duckdb-1.5.5-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f316eae2323d9a851883fdf2dee91c1f9efe251ab33e14a2272f82a913422ed6", size = 19391959, upload-time = "2026-07-22T10:54:15.551Z" }, + { url = "https://files.pythonhosted.org/packages/4f/98/6549769f158126fa64fd6c1ac2eb59a18282146c939867a3eb31b7c1db07/duckdb-1.5.5-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7a6d2d11859d82a936ebdcb30ce3d8a1cbb3e990bff05c12abb9b54c44fa7bd1", size = 21510909, upload-time = "2026-07-22T10:54:19.681Z" }, + { url = "https://files.pythonhosted.org/packages/af/b7/5753b41d3124838f868f9f523362812d9fc45409e9e4dd70dcbb0a25826e/duckdb-1.5.5-cp312-cp312-win_amd64.whl", hash = "sha256:ddfbdb096c11d51ee22492397d342c90a82e62c5d09961477895934d0a25372f", size = 13168544, upload-time = "2026-07-22T10:54:22.789Z" }, + { url = "https://files.pythonhosted.org/packages/5c/28/44b679c7d46245f8398feae7edac959d1b83d4eb143e25b3fce0630b78bd/duckdb-1.5.5-cp312-cp312-win_arm64.whl", hash = "sha256:2725d2b9ace3a4e75d72fc5a239f6a44b502c580edadb8fb2676db772c5f9282", size = 13988684, upload-time = "2026-07-22T10:54:26.003Z" }, + { url = "https://files.pythonhosted.org/packages/47/37/4a38116e7700720fd152c666292214fd3abdf916496991296d8d1f66efbf/duckdb-1.5.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:cd98829b67788609017e65c761bd42a5dd0f9129441bed8bda4d6881ccf819f0", size = 32754294, upload-time = "2026-07-22T10:54:29.822Z" }, + { url = "https://files.pythonhosted.org/packages/66/42/7d392f1ba1eee0eaf4ab4c8c7a604bfe3536cd63f979cf5c98798664f807/duckdb-1.5.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:feead93c56679b79592d437c62975d39cb67adedffa7592c763baf8160ac7366", size = 17368211, upload-time = "2026-07-22T10:54:33.359Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/0a6f4fa60562faa615e55e15bd1953a2f2b17a8edd8105e5cda215e43457/duckdb-1.5.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:49c963d9469373d7aba8d750d9ea565ab823e94166efed953f184dd9b169b98c", size = 15509136, upload-time = "2026-07-22T10:54:36.369Z" }, + { url = "https://files.pythonhosted.org/packages/e4/cb/023c89f51978545b9fab318581bba0c457a58e7530d2d933e54ae7d8647c/duckdb-1.5.5-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a736217825461732b5442d05a220f3da2e23a0dae114efbf08c9bf171b53098a", size = 19392147, upload-time = "2026-07-22T10:54:39.551Z" }, + { url = "https://files.pythonhosted.org/packages/3e/c5/41bef391fb8b23dbc133c9f2ba016e7a7a8124513d2cc1b430f1897d87e4/duckdb-1.5.5-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:078e6a60dd8eedde5832f45422ca5c4a6b8c837aeabd8a56ca0b7d933f588053", size = 21511060, upload-time = "2026-07-22T10:54:42.788Z" }, + { url = "https://files.pythonhosted.org/packages/07/9f/c44dfc1f924ac29b3252dc1b91393c01d009dbfe9f8ed33f10b986151bd1/duckdb-1.5.5-cp313-cp313-win_amd64.whl", hash = "sha256:6826504277dba513c0c5d71d828456c94d729c9d2482f94b2e289f90a9167e28", size = 13168028, upload-time = "2026-07-22T10:54:46.127Z" }, + { url = "https://files.pythonhosted.org/packages/ca/88/591384b2cd59abddd6f5dc175e60374f9abae6064429f0c4402854c10f44/duckdb-1.5.5-cp313-cp313-win_arm64.whl", hash = "sha256:baa9c5702002fabb559ded2a39008f9f421fcbc7237d388b8213eff1e08858de", size = 13989955, upload-time = "2026-07-22T10:54:49.262Z" }, +] + +[[package]] +name = "duckdb-analyst" +version = "0.1.0" +source = { editable = "toolsets/duckdb-analyst" } +dependencies = [ + { name = "duckdb" }, + { name = "langchain-core" }, + { name = "mcp-toolsets-runtime" }, +] + +[package.metadata] +requires-dist = [ + { name = "duckdb", specifier = ">=1.5.5,<2.0.0" }, + { name = "langchain-core", specifier = ">=1.4.6,<2.0.0" }, + { name = "mcp-toolsets-runtime" }, +] + [[package]] name = "fastapi" version = "0.141.1" @@ -1007,6 +1047,7 @@ version = "0.1.0" source = { virtual = "." } dependencies = [ { name = "credential-demo" }, + { name = "duckdb-analyst" }, { name = "hello" }, { name = "mcp-toolsets-runtime", extra = ["agent"] }, { name = "stac-explorer" }, @@ -1029,6 +1070,7 @@ index = [ [package.metadata] requires-dist = [ { name = "credential-demo", editable = "toolsets/credential-demo" }, + { name = "duckdb-analyst", editable = "toolsets/duckdb-analyst" }, { name = "hello", editable = "toolsets/hello" }, { name = "mcp-toolsets-runtime", extras = ["agent"], specifier = ">=0.3.0,<0.4.0" }, { name = "stac-explorer", editable = "toolsets/stac-explorer" },