From b63c9c2c1f2746aa53c1b894ae6a636c60794be0 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 10 Aug 2026 15:59:44 +0100 Subject: [PATCH 1/2] Add duckdb-analyst: sandbox-free SQL analysis over Overture/Natural Earth/STAC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three tools (list_sources, query, chart) over a single, hardened DuckDB connection: curated views for Overture Maps divisions/places and Natural Earth countries/populated places, plus STAC search via the community `stac` extension pointed at Planetary Computer. No code-execution sandbox needed — callers only ever write read-only SQL. Security model (see connection.py's module docstring for the full writeup): statement-shape validation (single SELECT/WITH, denylisted introspection calls) is defense-in-depth only; the real controls are DuckDB's disabled_filesystems + lock_configuration (verified empirically against a real DuckDB limitation where this only works after one remote-read "warmup" — undocumented upstream, reproduced and worked around here) and, primarily, the deployment shipping no local secrets for this connection to ever reach. --- pyproject.toml | 2 + toolsets/duckdb-analyst/pyproject.toml | 14 + .../src/duckdb_analyst/__init__.py | 1 + .../src/duckdb_analyst/connection.py | 513 ++++++++++++++++++ .../src/duckdb_analyst/security.py | 73 +++ .../src/duckdb_analyst/tools.py | 188 +++++++ .../tests/test_duckdb_analyst.py | 147 +++++ toolsets/duckdb-analyst/toolset.yaml | 13 + uv.lock | 42 ++ 9 files changed, 993 insertions(+) create mode 100644 toolsets/duckdb-analyst/pyproject.toml create mode 100644 toolsets/duckdb-analyst/src/duckdb_analyst/__init__.py create mode 100644 toolsets/duckdb-analyst/src/duckdb_analyst/connection.py create mode 100644 toolsets/duckdb-analyst/src/duckdb_analyst/security.py create mode 100644 toolsets/duckdb-analyst/src/duckdb_analyst/tools.py create mode 100644 toolsets/duckdb-analyst/tests/test_duckdb_analyst.py create mode 100644 toolsets/duckdb-analyst/toolset.yaml 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..79b4fe1 --- /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/STAC datasets (Overture Maps, Natural Earth, STAC) 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..cff37f7 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py @@ -0,0 +1,513 @@ +"""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 just before it, so the ``stac`` community extension can be + loaded once here but never again from caller SQL. +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``, and the community +``stac`` extension) — nothing under ``duckdb-analyst`` ever calls +``INSTALL``/``LOAD`` again after connection setup, and layer 3 makes sure +caller SQL couldn't anyway. ``httpfs``/``spatial`` are core DuckDB +extensions, so their versions are pinned transitively by this toolset's +pinned ``duckdb`` dependency. The community ``stac`` extension +(https://github.com/ahuarte47/duckdb-stac) has no independent version pin +available through ``INSTALL ... FROM community`` — it always resolves to +the latest build published for the installed DuckDB version. That's an +accepted, small supply-chain trust boundary (single maintainer, MIT +licensed, reviewed before adding here), not something this file can pin +away. +""" + +from typing import NotRequired, TypedDict + +import duckdb + +#: Overture release to read. Deliberately pinned, not globbed +#: (``release/*/...``): the public bucket keeps more than one release +#: available at a time, so a wildcard here would silently union +#: possibly-incompatible releases into one view. Bump this string +#: periodically — see https://docs.overturemaps.org/getting-data/duckdb/ for +#: the current release list. +OVERTURE_RELEASE = "2026-07-22.0" +_OVERTURE_BUCKET = "s3://overturemaps-us-west-2" + +# 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" +) + +#: Microsoft Planetary Computer's public, no-auth STAC API — same catalog +#: `stac-explorer` already uses (see its `STAC_ROOT`). `STAC_Search`'s `url` +#: argument is this endpoint's Item Search route. +PC_STAC_SEARCH_URL = "https://planetarycomputer.microsoft.com/api/stac/v1/search" + +#: 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}") + + # Pinned extensions only — see the module docstring, layer "Extensions". + con.execute("INSTALL httpfs") + con.execute("LOAD httpfs") + con.execute("INSTALL spatial") + con.execute("LOAD spatial") + con.execute("INSTALL stac FROM community") + con.execute("LOAD stac") + + # Overture's bucket is genuinely anonymous/unsigned — no CREATE SECRET, + # no credentials, just the region (verified directly against the live + # bucket; this is also the pattern Overture's own docs recommend). + con.execute("SET s3_region = 'us-west-2'") + + _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 allow_community_extensions = false") + 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"""...""")` directly, 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 (a + # pinned release date, hardcoded dataset URLs) — never caller input, + # which is what actually makes string-built SQL dangerous. + divisions_path = ( + f"{_OVERTURE_BUCKET}/release/{OVERTURE_RELEASE}/" + "theme=divisions/type=division_area/*" + ) + divisions_sql = f""" + CREATE VIEW overture_divisions AS + SELECT + id, + subtype, + class, + country, + region, + admin_level, + names."primary" AS primary_name, + is_land, + is_territorial, + bbox.xmin AS bbox_xmin, + bbox.ymin AS bbox_ymin, + bbox.xmax AS bbox_xmax, + bbox.ymax AS bbox_ymax, + geometry + FROM read_parquet('{divisions_path}') + """ # noqa: S608 + con.execute(divisions_sql) + + places_path = ( + f"{_OVERTURE_BUCKET}/release/{OVERTURE_RELEASE}/theme=places/type=place/*" + ) + places_sql = f""" + CREATE VIEW overture_places AS + SELECT + id, + names."primary" AS primary_name, + categories."primary" AS category, + confidence, + operating_status, + basic_category, + brand.names."primary" AS brand_name, + addresses[1].country AS country, + addresses[1].region AS region, + addresses[1].locality AS locality, + bbox.xmin AS bbox_xmin, + bbox.ymin AS bbox_ymin, + bbox.xmax AS bbox_xmax, + bbox.ymax AS bbox_ymax, + geometry + FROM read_parquet('{places_path}') + """ # noqa: S608 + con.execute(places_sql) + + 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]] = { + "overture_divisions": { + "subtype": { + "name": "subtype", + "type": "", + "description": "country / region / county / locality / etc.", + "good_for": ["x", "color"], + }, + "country": { + "name": "country", + "type": "", + "description": "ISO 3166-1 alpha-2 country code.", + "good_for": ["x", "color"], + }, + "region": { + "name": "region", + "type": "", + "description": "ISO 3166-2 principal-subdivision code.", + "good_for": ["x", "color"], + }, + "admin_level": { + "name": "admin_level", + "type": "", + "description": "Depth in the administrative hierarchy (0=country).", + "good_for": ["x", "color"], + }, + "primary_name": { + "name": "primary_name", + "type": "", + "description": "The division's primary display name.", + "good_for": ["x", "color"], + }, + "geometry": { + "name": "geometry", + "type": "", + "description": ( + "Polygon/MultiPolygon boundary. Not chart data as-is — wrap " + "with ST_AsText(geometry) or ST_AsGeoJSON(geometry) if you " + "need it in a query/chart result." + ), + }, + }, + "overture_places": { + "category": { + "name": "category", + "type": "", + "description": "Primary Overture place-category taxonomy value.", + "good_for": ["x", "color"], + }, + "confidence": { + "name": "confidence", + "type": "", + "description": "Model confidence the place exists, 0-1.", + "good_for": ["y", "color"], + }, + "country": { + "name": "country", + "type": "", + "description": "Country code of the place's first listed address.", + "good_for": ["x", "color"], + }, + "operating_status": { + "name": "operating_status", + "type": "", + "description": "open / temporarily_closed / permanently_closed.", + "good_for": ["color"], + }, + "geometry": { + "name": "geometry", + "type": "", + "description": ( + "Point location. Wrap with ST_X(geometry)/ST_Y(geometry) " + "for plain lon/lat numbers usable as x/y chart channels." + ), + }, + }, + "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 = { + "overture_divisions": ( + "Overture Maps administrative divisions (countries, regions, " + f"counties, ...), release {OVERTURE_RELEASE}, polygons. Good for " + "breakdowns by country/region/admin_level." + ), + "overture_places": ( + "Overture Maps points of interest (POIs) — businesses, landmarks, " + f"etc., release {OVERTURE_RELEASE}, points. Good for category/" + "confidence/operating-status breakdowns; global scale (billions of " + "rows across all files), so always filter or LIMIT." + ), + "natural_earth_countries": ( + "Natural Earth 1:110m country polygons with basic demographic/" + "economic attributes — a lightweight companion to overture_divisions " + "when you just need country-level name/population/GDP, not a full " + "administrative hierarchy." + ), + "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 = { + "overture_divisions": ( + "SELECT primary_name, country, admin_level FROM overture_divisions " + "WHERE subtype = 'country' LIMIT 20" + ), + "overture_places": ( + "SELECT category, count(*) AS n FROM overture_places " + "WHERE country = 'US' GROUP BY category ORDER BY n DESC LIMIT 20" + ), + "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]: + sources = [ + SourceInfo( + name=view, + kind="view", + description=_VIEW_DESCRIPTIONS[view], + example_sql=_EXAMPLE_SQL[view], + columns=_describe_view(con, view), + ) + for view in ( + "overture_divisions", + "overture_places", + "natural_earth_countries", + "natural_earth_places", + ) + ] + # A hand-written example, not SQL this module ever executes — it's + # display-only text for list_sources, so ruff's S608 heuristic (which + # can't tell "this is documentation" from "this runs") is a false + # positive here. + stac_example_sql = ( + 'SELECT collection, id, datetime, "eo:cloud_cover" ' # noqa: S608 + f"FROM STAC_Search('{PC_STAC_SEARCH_URL}', " + "collections := ['sentinel-2-l2a'], " + "bbox := [-122.5, 47.5, -122.0, 47.8], " + "datetime := '2024-06-01/2024-06-30', max_items := 20)" + ) + sources.append( + SourceInfo( + name="STAC_Search", + kind="table_function", + description=( + "SQL table function (from the `stac` DuckDB extension) that " + "searches any STAC API's Item Search endpoint. Pointed here " + "at Microsoft Planetary Computer's public, no-auth catalog " + f"({PC_STAC_SEARCH_URL}); pass a different `url` for another " + "STAC API. Common result columns: collection, id, datetime, " + "bbox, geometry, plus collection-specific properties (e.g. " + '"eo:cloud_cover") — these vary by collection, so run ' + "`DESCRIBE` on a search first if unsure." + ), + example_sql=stac_example_sql, + ) + ) + return sources + + +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..a72cd17 --- /dev/null +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py @@ -0,0 +1,188 @@ +"""LangChain tools for duckdb-analyst: sandbox-free SQL analysis over public +parquet/STAC 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 — STAC +search is just another SQL table function reachable through `query`/`chart`, +documented 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 and SQL table functions available to + `query`/`chart`: Overture Maps divisions/places, Natural Earth countries/ + populated places, and STAC search — with column descriptions and which + columns work well as x/y/color chart channels. Call this before writing + SQL against an unfamiliar source. + """ + 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..9de39be --- /dev/null +++ b/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py @@ -0,0 +1,147 @@ +"""Tests for duckdb-analyst. + +These hit the real network — DuckDB's own httpfs/spatial/stac 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 small, +filtered queries against the curated views (parquet row-group pruning means +these don't pull the whole multi-hundred-MB dataset over the wire). + +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 { + "overture_divisions", + "overture_places", + "natural_earth_countries", + "natural_earth_places", + "STAC_Search", + } <= names + for source in result["sources"]: + assert source["description"] + assert source["example_sql"] + + +def test_list_sources_flags_chart_friendly_columns(): + result = list_sources.invoke({}) + divisions = next(s for s in result["sources"] if s["name"] == "overture_divisions") + country_column = next(c for c in divisions["columns"] if c["name"] == "country") + assert "x" in country_column["good_for"] or "color" in country_column["good_for"] + + +async def test_query_against_curated_view(): + result = await query.ainvoke( + { + "sql": ( + "SELECT primary_name, country FROM overture_divisions " + "WHERE subtype = 'country' AND primary_name IS NOT NULL" + ), + "limit": 5, + } + ) + assert not is_error(result) + assert 0 < len(result["rows"]) <= 5 + assert result["row_count"] == len(result["rows"]) + assert "primary_name" in result["rows"][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(): + # A LIMIT-bounded, single-file-ish query (like the curated-view query + # test above) rather than an unfiltered GROUP BY: aggregating across the + # whole multi-GB overture_divisions view would need to scan every row + # group instead of stopping early, and blow the 30s query timeout. + spec = { + "mark": "bar", + "encoding": {"x": {"field": "primary_name"}, "y": {"field": "country"}}, + } + result = await chart.ainvoke( + { + "sql": ( + "SELECT primary_name, country FROM overture_divisions " + "WHERE subtype = 'country' AND primary_name IS NOT NULL" + ), + "spec": spec, + "limit": 10, + } + ) + assert not is_error(result) + assert result["spec"]["mark"] == "bar" + assert result["spec"]["encoding"] == spec["encoding"] + assert len(result["spec"]["data"]["values"]) <= 10 + + +# --------------------------------------------------------------------------- +# 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" }, From 7b4276a4cc2966bdda1fd1c7285fc63509f64ad6 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 10 Aug 2026 17:21:47 +0100 Subject: [PATCH 2/2] Narrow scope to small datasets: drop Overture and STAC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both were advertised in list_sources but neither was usable in practice, and no query test covered either one. Overture: the public parquet layout gives no partition pruning on the fields an analyst filters on, so any aggregate scans the whole theme. Measured against this toolset's memory/thread caps and the 30s watchdog: a bare SELECT ... LIMIT 3 took ~2.8s, adding WHERE country = 'PT' took ~17.4s, and a GROUP BY exceeded the timeout. A dataset that size wants a local copy, which is just another CREATE VIEW here. STAC: STAC_Search against Planetary Computer fails outright — HTTP 422 on the Item Search route, 405 on the catalog root. The `stac` community extension was also the only unpinnable dependency in the toolset, since INSTALL ... FROM community always resolves to the latest build. Removing it lets allow_community_extensions = false move ahead of the first INSTALL, so community extensions are now refused outright rather than allowed once and closed off afterwards. Tests now exercise every advertised source with a real query, and test_list_sources_advertises_nothing_untested keeps that true. Suite drops from 22.6s to 3.2s; 17 tests, security regression suite unchanged. --- toolsets/duckdb-analyst/pyproject.toml | 2 +- .../src/duckdb_analyst/connection.py | 274 ++++-------------- .../src/duckdb_analyst/tools.py | 17 +- .../tests/test_duckdb_analyst.py | 126 ++++++-- 4 files changed, 160 insertions(+), 259 deletions(-) diff --git a/toolsets/duckdb-analyst/pyproject.toml b/toolsets/duckdb-analyst/pyproject.toml index 79b4fe1..4c0d295 100644 --- a/toolsets/duckdb-analyst/pyproject.toml +++ b/toolsets/duckdb-analyst/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "duckdb-analyst" version = "0.1.0" -description = "Read-only, sandbox-free SQL analysis over public parquet/STAC datasets (Overture Maps, Natural Earth, STAC) via DuckDB." +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", diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py b/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py index cff37f7..6207c32 100644 --- a/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/connection.py @@ -52,9 +52,11 @@ 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 just before it, so the ``stac`` community extension can be - loaded once here but never again from caller SQL. + 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 @@ -63,33 +65,39 @@ 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``, and the community -``stac`` extension) — nothing under ``duckdb-analyst`` ever calls -``INSTALL``/``LOAD`` again after connection setup, and layer 3 makes sure -caller SQL couldn't anyway. ``httpfs``/``spatial`` are core DuckDB -extensions, so their versions are pinned transitively by this toolset's -pinned ``duckdb`` dependency. The community ``stac`` extension -(https://github.com/ahuarte47/duckdb-stac) has no independent version pin -available through ``INSTALL ... FROM community`` — it always resolves to -the latest build published for the installed DuckDB version. That's an -accepted, small supply-chain trust boundary (single maintainer, MIT -licensed, reviewed before adding here), not something this file can pin -away. +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 -#: Overture release to read. Deliberately pinned, not globbed -#: (``release/*/...``): the public bucket keeps more than one release -#: available at a time, so a wildcard here would silently union -#: possibly-incompatible releases into one view. Bump this string -#: periodically — see https://docs.overturemaps.org/getting-data/duckdb/ for -#: the current release list. -OVERTURE_RELEASE = "2026-07-22.0" -_OVERTURE_BUCKET = "s3://overturemaps-us-west-2" - # 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 @@ -109,11 +117,6 @@ "geojson/ne_110m_populated_places.geojson" ) -#: Microsoft Planetary Computer's public, no-auth STAC API — same catalog -#: `stac-explorer` already uses (see its `STAC_ROOT`). `STAC_Search`'s `url` -#: argument is this endpoint's Item Search route. -PC_STAC_SEARCH_URL = "https://planetarycomputer.microsoft.com/api/stac/v1/search" - #: 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. @@ -146,18 +149,20 @@ def _build_connection() -> duckdb.DuckDBPyConnection: con.execute(f"SET memory_limit = '{_MEMORY_LIMIT}'") con.execute(f"SET threads = {_THREADS}") - # Pinned extensions only — see the module docstring, layer "Extensions". + # 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") - con.execute("INSTALL stac FROM community") - con.execute("LOAD stac") - # Overture's bucket is genuinely anonymous/unsigned — no CREATE SECRET, - # no credentials, just the region (verified directly against the live - # bucket; this is also the pattern Overture's own docs recommend). - con.execute("SET s3_region = 'us-west-2'") + # 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) @@ -165,7 +170,6 @@ def _build_connection() -> duckdb.DuckDBPyConnection: # 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 allow_community_extensions = false") con.execute("SET lock_configuration = true") return con @@ -185,63 +189,13 @@ def _warm_up_remote_filesystem(con: duckdb.DuckDBPyConnection) -> None: 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"""...""")` directly, so + # 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 (a - # pinned release date, hardcoded dataset URLs) — never caller input, - # which is what actually makes string-built SQL dangerous. - divisions_path = ( - f"{_OVERTURE_BUCKET}/release/{OVERTURE_RELEASE}/" - "theme=divisions/type=division_area/*" - ) - divisions_sql = f""" - CREATE VIEW overture_divisions AS - SELECT - id, - subtype, - class, - country, - region, - admin_level, - names."primary" AS primary_name, - is_land, - is_territorial, - bbox.xmin AS bbox_xmin, - bbox.ymin AS bbox_ymin, - bbox.xmax AS bbox_xmax, - bbox.ymax AS bbox_ymax, - geometry - FROM read_parquet('{divisions_path}') - """ # noqa: S608 - con.execute(divisions_sql) - - places_path = ( - f"{_OVERTURE_BUCKET}/release/{OVERTURE_RELEASE}/theme=places/type=place/*" - ) - places_sql = f""" - CREATE VIEW overture_places AS - SELECT - id, - names."primary" AS primary_name, - categories."primary" AS category, - confidence, - operating_status, - basic_category, - brand.names."primary" AS brand_name, - addresses[1].country AS country, - addresses[1].region AS region, - addresses[1].locality AS locality, - bbox.xmin AS bbox_xmin, - bbox.ymin AS bbox_ymin, - bbox.xmax AS bbox_xmax, - bbox.ymax AS bbox_ymax, - geometry - FROM read_parquet('{places_path}') - """ # noqa: S608 - con.execute(places_sql) - + # 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 @@ -280,81 +234,6 @@ def _register_views(con: duckdb.DuckDBPyConnection) -> None: # 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]] = { - "overture_divisions": { - "subtype": { - "name": "subtype", - "type": "", - "description": "country / region / county / locality / etc.", - "good_for": ["x", "color"], - }, - "country": { - "name": "country", - "type": "", - "description": "ISO 3166-1 alpha-2 country code.", - "good_for": ["x", "color"], - }, - "region": { - "name": "region", - "type": "", - "description": "ISO 3166-2 principal-subdivision code.", - "good_for": ["x", "color"], - }, - "admin_level": { - "name": "admin_level", - "type": "", - "description": "Depth in the administrative hierarchy (0=country).", - "good_for": ["x", "color"], - }, - "primary_name": { - "name": "primary_name", - "type": "", - "description": "The division's primary display name.", - "good_for": ["x", "color"], - }, - "geometry": { - "name": "geometry", - "type": "", - "description": ( - "Polygon/MultiPolygon boundary. Not chart data as-is — wrap " - "with ST_AsText(geometry) or ST_AsGeoJSON(geometry) if you " - "need it in a query/chart result." - ), - }, - }, - "overture_places": { - "category": { - "name": "category", - "type": "", - "description": "Primary Overture place-category taxonomy value.", - "good_for": ["x", "color"], - }, - "confidence": { - "name": "confidence", - "type": "", - "description": "Model confidence the place exists, 0-1.", - "good_for": ["y", "color"], - }, - "country": { - "name": "country", - "type": "", - "description": "Country code of the place's first listed address.", - "good_for": ["x", "color"], - }, - "operating_status": { - "name": "operating_status", - "type": "", - "description": "open / temporarily_closed / permanently_closed.", - "good_for": ["color"], - }, - "geometry": { - "name": "geometry", - "type": "", - "description": ( - "Point location. Wrap with ST_X(geometry)/ST_Y(geometry) " - "for plain lon/lat numbers usable as x/y chart channels." - ), - }, - }, "natural_earth_countries": { "name": { "name": "name", @@ -403,22 +282,11 @@ def _register_views(con: duckdb.DuckDBPyConnection) -> None: } _VIEW_DESCRIPTIONS = { - "overture_divisions": ( - "Overture Maps administrative divisions (countries, regions, " - f"counties, ...), release {OVERTURE_RELEASE}, polygons. Good for " - "breakdowns by country/region/admin_level." - ), - "overture_places": ( - "Overture Maps points of interest (POIs) — businesses, landmarks, " - f"etc., release {OVERTURE_RELEASE}, points. Good for category/" - "confidence/operating-status breakdowns; global scale (billions of " - "rows across all files), so always filter or LIMIT." - ), "natural_earth_countries": ( "Natural Earth 1:110m country polygons with basic demographic/" - "economic attributes — a lightweight companion to overture_divisions " - "when you just need country-level name/population/GDP, not a full " - "administrative hierarchy." + "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 " @@ -428,14 +296,6 @@ def _register_views(con: duckdb.DuckDBPyConnection) -> None: } _EXAMPLE_SQL = { - "overture_divisions": ( - "SELECT primary_name, country, admin_level FROM overture_divisions " - "WHERE subtype = 'country' LIMIT 20" - ), - "overture_places": ( - "SELECT category, count(*) AS n FROM overture_places " - "WHERE country = 'US' GROUP BY category ORDER BY n DESC LIMIT 20" - ), "natural_earth_countries": ( "SELECT name, continent, population_estimate FROM " "natural_earth_countries ORDER BY population_estimate DESC LIMIT 20" @@ -463,7 +323,7 @@ def _describe_view(con: duckdb.DuckDBPyConnection, view: str) -> list[ColumnInfo def _build_sources(con: duckdb.DuckDBPyConnection) -> list[SourceInfo]: - sources = [ + return [ SourceInfo( name=view, kind="view", @@ -471,42 +331,8 @@ def _build_sources(con: duckdb.DuckDBPyConnection) -> list[SourceInfo]: example_sql=_EXAMPLE_SQL[view], columns=_describe_view(con, view), ) - for view in ( - "overture_divisions", - "overture_places", - "natural_earth_countries", - "natural_earth_places", - ) + for view in ("natural_earth_countries", "natural_earth_places") ] - # A hand-written example, not SQL this module ever executes — it's - # display-only text for list_sources, so ruff's S608 heuristic (which - # can't tell "this is documentation" from "this runs") is a false - # positive here. - stac_example_sql = ( - 'SELECT collection, id, datetime, "eo:cloud_cover" ' # noqa: S608 - f"FROM STAC_Search('{PC_STAC_SEARCH_URL}', " - "collections := ['sentinel-2-l2a'], " - "bbox := [-122.5, 47.5, -122.0, 47.8], " - "datetime := '2024-06-01/2024-06-30', max_items := 20)" - ) - sources.append( - SourceInfo( - name="STAC_Search", - kind="table_function", - description=( - "SQL table function (from the `stac` DuckDB extension) that " - "searches any STAC API's Item Search endpoint. Pointed here " - "at Microsoft Planetary Computer's public, no-auth catalog " - f"({PC_STAC_SEARCH_URL}); pass a different `url` for another " - "STAC API. Common result columns: collection, id, datetime, " - "bbox, geometry, plus collection-specific properties (e.g. " - '"eo:cloud_cover") — these vary by collection, so run ' - "`DESCRIBE` on a search first if unsure." - ), - example_sql=stac_example_sql, - ) - ) - return sources CON = _build_connection() diff --git a/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py b/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py index a72cd17..bb1d7c1 100644 --- a/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py +++ b/toolsets/duckdb-analyst/src/duckdb_analyst/tools.py @@ -1,11 +1,12 @@ """LangChain tools for duckdb-analyst: sandbox-free SQL analysis over public -parquet/STAC datasets via a security-hardened, read-only DuckDB connection. +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 — STAC -search is just another SQL table function reachable through `query`/`chart`, -documented via `list_sources` (see `connection.SOURCES`). +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. @@ -131,11 +132,13 @@ async def watchdog() -> None: @tool def list_sources() -> ListSourcesResult: - """List the pre-registered views and SQL table functions available to - `query`/`chart`: Overture Maps divisions/places, Natural Earth countries/ - populated places, and STAC search — with column descriptions and which + """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( diff --git a/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py b/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py index 9de39be..2b4a609 100644 --- a/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py +++ b/toolsets/duckdb-analyst/tests/test_duckdb_analyst.py @@ -1,11 +1,19 @@ """Tests for duckdb-analyst. -These hit the real network — DuckDB's own httpfs/spatial/stac 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 small, -filtered queries against the curated views (parquet row-group pruning means -these don't pull the whole multi-hundred-MB dataset over the wire). +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. @@ -23,31 +31,41 @@ def test_list_sources_covers_every_dataset_family(): result = list_sources.invoke({}) names = {source["name"] for source in result["sources"]} - assert { - "overture_divisions", - "overture_places", - "natural_earth_countries", - "natural_earth_places", - "STAC_Search", - } <= names + 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({}) - divisions = next(s for s in result["sources"] if s["name"] == "overture_divisions") - country_column = next(c for c in divisions["columns"] if c["name"] == "country") - assert "x" in country_column["good_for"] or "color" in country_column["good_for"] + 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 primary_name, country FROM overture_divisions " - "WHERE subtype = 'country' AND primary_name IS NOT NULL" + "SELECT name, continent, population_estimate " + "FROM natural_earth_countries ORDER BY population_estimate DESC" ), "limit": 5, } @@ -55,7 +73,57 @@ async def test_query_against_curated_view(): assert not is_error(result) assert 0 < len(result["rows"]) <= 5 assert result["row_count"] == len(result["rows"]) - assert "primary_name" in result["rows"][0] + 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(): @@ -80,19 +148,19 @@ async def test_query_enforces_hard_row_cap_server_side(): async def test_chart_fills_in_data_values_and_preserves_spec(): - # A LIMIT-bounded, single-file-ish query (like the curated-view query - # test above) rather than an unfiltered GROUP BY: aggregating across the - # whole multi-GB overture_divisions view would need to scan every row - # group instead of stopping early, and blow the 30s query timeout. spec = { "mark": "bar", - "encoding": {"x": {"field": "primary_name"}, "y": {"field": "country"}}, + "encoding": { + "x": {"field": "continent", "type": "nominal", "sort": "-y"}, + "y": {"field": "population", "type": "quantitative"}, + }, } result = await chart.ainvoke( { "sql": ( - "SELECT primary_name, country FROM overture_divisions " - "WHERE subtype = 'country' AND primary_name IS NOT NULL" + "SELECT continent, SUM(population_estimate) AS population " + "FROM natural_earth_countries GROUP BY continent " + "ORDER BY population DESC" ), "spec": spec, "limit": 10, @@ -101,7 +169,11 @@ async def test_chart_fills_in_data_values_and_preserves_spec(): assert not is_error(result) assert result["spec"]["mark"] == "bar" assert result["spec"]["encoding"] == spec["encoding"] - assert len(result["spec"]["data"]["values"]) <= 10 + 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]) # ---------------------------------------------------------------------------