DEV Community

SAI RAM
SAI RAM

Posted on Originally published at anvilry.vercel.app

Why I Built cost-guard-mcp: Pre-Flight Cost Guardrails for AI Agents Talking to Data Warehouses

Description of the GIF

The query that shouldn't have run

Give an AI agent a warehouse MCP server and, sooner or later, it will write a query that scans a full multi-terabyte table and quietly runs up a bill in the hundreds of dollars because it forgot a partition filter, or it will SELECT * from something it thought was small and get back four million rows straight into its own context window. Neither of these requires a malicious agent or a bad prompt. It just requires an agent doing what agents do: writing plausible SQL against a schema it partially understands, then running it. A human DBA would eyeball the query, guess the table size, maybe run EXPLAIN first out of habit. An agent given a generic "run this SQL" tool has no equivalent instinct, and as far as I can tell, no warehouse MCP server on the market tells the agent what a query will cost or how much data it will return before that query actually executes.

That gap is the entire reason cost-guard-mcp exists. It is a Model Context Protocol server that sits in front of BigQuery and Snowflake and gives an agent three tools: describe_engine_capabilities(engine) to learn what an engine can and can't tell you, estimate_query_cost(engine, sql, warehouse) to get a cost and byte estimate before running anything, and run_query_bounded(engine, sql, max_bytes_billed, max_rows, max_estimated_cost_usd) to actually execute a query with those caps enforced. It is a genuinely small project — 15 Python files, 715 lines under src/cost_guard_mcp/ — and it is genuinely new. The first commit landed 2026-09-12 at 22:23 IST, the feature-complete v1 landed a few hours later at 02:27, v0.1.0 shipped that same day, and v0.1.1 — the hardening pass I'll get to below — shipped today, 2026-09-13. The whole history of this project fits inside about eighteen hours. I'm not going to pretend otherwise; there's no multi-month backstory here, just a focused tool built fast and then immediately hardened.

Alternative description of the GIF

An estimate that tells you how much to trust itself

The design decision I care about most in this codebase isn't the warehouse integration — it's that estimate_query_cost never returns a bare number. Every CostEstimate it produces carries an accuracy_tier field, and that field is not a comment or a README promise, it's a required field on the Pydantic model in src/cost_guard_mcp/types.py with no default. You cannot construct a CostEstimate without deciding what tier it belongs to. AGENTS.md in the repo states the intent behind this directly: it's "enforced structurally by the Pydantic model, not by convention." That distinction matters more than it sounds. A convention is something a future contributor forgets. A required field with no default is something the type checker refuses to let you skip.

There are three tiers: PRECISE, UPPER_BOUND, and HEURISTIC. BigQuery is where PRECISE actually happens, and it happens because BigQuery's dry-run API is not a client-side guess — it's a real submission to BigQuery's own query planner. dry_run() in src/cost_guard_mcp/engines/bigquery.py builds a bigquery.QueryJobConfig(dry_run=True, use_query_cache=False), sends it through client.query(), and reads back total_bytes_processed from the real API response. BigQuery validates and fully plans the query server-side without executing or billing it, so the byte figure that comes back is the same figure the real execution would have produced. That's what earns the PRECISE label.

Snowflake gets a structurally different treatment, and the code is explicit about why. explain_estimate() runs EXPLAIN USING JSON {sql}, parses the returned plan, and reads plan['GlobalStats']['bytesAssigned'] as its byte figure — but that number describes what the query planner expects to scan, not what actually gets scanned, and Snowflake's own documentation (quoted directly in a code comment) says runtime plan optimizations "can reduce the number of partitions and bytes scanned." Because of that, explain_estimate() doesn't have a downgrade table the way BigQuery does — it hardcodes AccuracyTier.UPPER_BOUND on line 84, unconditionally, because there is no better tier available to fall from. BigQuery's dry-run and Snowflake's EXPLAIN are answering genuinely different questions: one is "what will this actually cost," the other is "what is the most this should cost."

The place this stops being an abstract distinction and becomes something you can watch happen is inside BigQuery itself. BigQuery's dry-run response includes its own confidence signal — totalBytesProcessedAccuracy — and dry_run() reads it via query_job._properties['statistics']['query']['totalBytesProcessedAccuracy'], defaulting to 'UNKNOWN' if it's missing, then maps it through a fixed table: PRECISE stays PRECISE; LOWER_BOUND, UPPER_BOUND, UNKNOWN, and anything else all fall through to AccuracyTier.UPPER_BOUND via the dict's own default. A unit test, test_dry_run_treats_unrecognized_accuracy_value_as_upper_bound, feeds it a value BigQuery has never returned before — 'SOME_FUTURE_VALUE' — and asserts the tool still downgrades safely rather than defaulting to PRECISE. The practical consequence: run the exact same SQL text against an ordinary settled table and you get PRECISE with an empty caveats list. Run that identical SQL against a table with a pending streaming buffer, or a wildcard, or a federated source, and the tool automatically flips to UPPER_BOUND and attaches a caveat quoting BigQuery's own raw accuracy string verbatim — something like "BigQuery reported this estimate's own accuracy as 'LOWER_BOUND', not PRECISE — treating it conservatively as UPPER_BOUND." Nothing about the query changed. Only BigQuery's own confidence in its byte count changed, and the tool surfaces that shift instead of quietly reporting one undifferentiated number both times.

The third tier, HEURISTIC, is reserved for Databricks, and I want to be precise about its status: it exists in the type system — types.py defines the enum value, and the agent-facing docstring in server.py mentions it — but there is no engines/databricks.py, no pricing module, nothing. describe_engine_capabilities('databricks') raises ValueError, and a test asserts exactly that. I deferred Databricks rather than ship it, because a Databricks estimate would necessarily be a heuristic guess with no dry-run and no EXPLAIN-equivalent bound behind it, and shipping a guess dressed up as a real number is precisely the failure mode this whole project exists to prevent. Better to leave the third tier as a load-bearing placeholder in the schema than half-honest in the field.

The server itself bakes the sequencing into its own tool docstrings, not just into documentation: describe_engine_capabilities's docstring tells an agent to call it "before estimate_query_cost or run_query_bounded to understand how much to trust the accuracy_tier on their responses for this engine," and estimate_query_cost's docstring spells out what each tier means in the response itself. That's the actual contract an agent operates against — not marketing copy, code the agent reads.

Refusing is safer than guessing

A short description of what is happening in the GIF

The second thread running through this project is fail-closed, least-privilege design, and it shows up as a pattern rather than a single feature. The clearest instance is in run_query_bounded. If a caller passes max_estimated_cost_usd but the estimate that came back has estimated_cost_usd=None — which happens for a BigQuery Editions/capacity-billed project, since capacity billing has no fixed dollar-per-byte rate to convert from — the tool doesn't run the query with that cap silently unenforced. It refuses. The check in src/cost_guard_mcp/tools/run_query_bounded.py returns a normal BoundedQueryResult with status="refused" and reason=RefusalReason.COST_CAP_EXCEEDED, plus a plain-English hint explaining exactly why: it couldn't produce a dollar estimate, so it's refusing rather than running uncapped. This is a normal, successful MCP tool result, not an exception and not the MCP error channel — that's a deliberate choice recorded in the project's own ADR #4, so a caller can distinguish "your query is too expensive" from "your MCP client is broken" cleanly.

Two other examples of the same instinct: Snowflake connections require an explicit SNOWFLAKE_ROLE environment variable with no default, and config.py raises a ConfigError if it's unset, specifically to prevent ever falling back to ACCOUNTADMIN. And every function in both engine modules that touches snowflake.connector or the BigQuery client libraries is wrapped by a sanitize_exceptions decorator, which regex-redacts five categories of secret material — passwords, private keys and PEM blocks, tokens/API keys, and user:password@host URIs — from any exception before it can reach a caller or a log line, because snowflake-connector-python has a documented history of leaking credentials into raw exception text.

The part of this story I find most worth being honest about is that two of the security fixes weren't found by a later audit — they were found and closed the same day the vulnerable code shipped. Git history has the receipts: a commit titled fix(security): validate Snowflake warehouse name before USE WAREHOUSE, because the warehouse parameter — caller-supplied, same trust level as the SQL itself — was being interpolated straight into USE WAREHOUSE {warehouse}. That statement is executed directly against the live Snowflake session via cur.execute(), so an unvalidated warehouse name wasn't just a cosmetic risk: a caller could have appended a semicolon and a second statement after the intended identifier, turning a parameter meant to just pick a compute resource into a general SQL execution primitive scoped by whatever the connection's role could already reach. The second commit, fix(bigquery): validate project ID before interpolating into API filter, whose own message says it was "flagged by automated security review after the previous commit introduced the query filter," closes a related but structurally different gap: the GCP project ID isn't actually reachable from an MCP tool parameter today, but the code treats the identifier as if it could be — validated on the theory that anything string-interpolated into an API filter expression deserves the same allowlist treatment as untrusted input, not just the fields a caller happens to be able to reach right now. Both are closed with allowlist regex validators — _validate_warehouse() and _validate_project_id() — that fail loudly with a ValueError (itself sanitized) rather than let a malformed identifier reach a query-language context. Both incidents are written up together as ADR #6 in DECISIONS.md, which turns them into a standing rule for any engine this project adds in the future.

One more mechanism worth naming because it's small and clever: row-bounded execution on both engines wraps the caller's SQL as SELECT * FROM (<sql>) AS cost_guard_row_cap LIMIT max_rows + 1 before fetching. The +1 is the whole trick — fetch exactly max_rows and you can't tell "there were exactly that many rows" apart from "there were more, and we truncated." Fetch more than max_rows + 1 and the row cap stops doing its job of keeping the fetch itself bounded. This pattern trips Ruff's bandit-style SQL-injection ruleset (S608) in exactly two places, once per engine, and both are suppressed inline with # noqa: S608 because the string interpolation is the caller's own already-validated SQL being wrapped, not an unvalidated identifier.

Hardening the guardrail tool, immediately

A short description of what is happening in the third GIF

I shipped v0.1.0 on 2026-09-12. One day later, before the project had a single GitHub star, I ran a security pass against my own new code and shipped v0.1.1 with PyPI Trusted Publishing (OIDC via id-token: write, scoped to the pypi environment, with zero stored API token anywhere in release.yml), CodeQL scanning on a weekly cron, OpenSSF Scorecard with SARIF results uploaded to GitHub's Security tab, secret scanning, push protection, Dependabot security updates, and branch protection requiring CI to pass on main. All of that is free for a public repository and was simply off for the first day. Concretely: secret scanning and push protection are what would catch an accidentally committed API key or private key before it ever settles into git history, not after; Dependabot security updates are what would catch a CVE landing in a transitive dependency next month, long after I've stopped actively watching this repo. OpenSSF Scorecard in particular runs a battery of named checks against the repo's own supply-chain hygiene — Dangerous-Workflow for GitHub Actions patterns that let untrusted input execute privileged code, Branch-Protection for exactly the main-branch CI gate I just turned on, and Pinned-Dependencies for whether third-party actions are pinned to a mutable tag or a fixed SHA. It felt right to treat the guardrail tool itself as something that needed guardrails, and to do it before the project attracted any real attention rather than after.

The CI structure separates required from advisory checks rather than treating everything as one gate: tests.yml runs Ruff, mypy, and the unit suite with --cov-fail-under=80, and that's what actually blocks a merge to main. audit.yml runs pip-audit but is deliberately not a required check — its own header comment explains the reasoning: a new CVE in an already-approved, unrelated dependency shouldn't block an unrelated PR from merging. integration.yml is separate again, triggered only by workflow_dispatch or a weekly schedule, and runs the five tests under tests/integration/ against real BigQuery and Snowflake credentials in a live-integration environment. Those five tests are not part of the 60 unit tests I count as the project's real coverage number — I ran pytest tests/ --ignore=tests/integration directly against this checkout and got 60 passed, 0 failed, and one asyncio deprecation warning coming from an early bootstrap test in tests/unit/test_server.py that still calls asyncio.get_event_loop() and says in its own comment that it "documents current state" pending cleanup — a small piece of unfinished housekeeping from the first day, not noise from an external dependency. The integration tests need cloud credentials I'm not going to hand to a CI runner on every PR, so they stay gated and scheduled instead.

DECISIONS.md also documents a bug that's worth naming precisely because it's the same failure class as the fail-closed cap-check I described above: the original run_query_bounded logic used a flat boolean chain that fell through to "allow" whenever a cap was requested but its matching estimate field was None. An automated review caught it. The fix — refuse rather than silently execute unenforced — is now the standing convention for the whole tool.

What this doesn't do yet, on purpose

A short description of what is happening in the fourth GIF

I'd rather state the gaps plainly than let anyone find them the hard way. Databricks isn't supported — the HEURISTIC tier is a type-level placeholder with zero engine code behind it, deferred past v1 because I'm not willing to ship a guess dressed as a measurement. Snowflake's UPPER_BOUND estimate excludes Cortex AI Function ("AI Credits") cost entirely — EXPLAIN's bytesAssigned reflects warehouse compute and scan, not AI-inference calls embedded in the SQL, and both the static capability lookup and the live caveat on every Snowflake estimate say so. And BigQuery projects on Editions or capacity billing get a byte estimate with no dollar figure at all, because capacity billing charges for reserved slot-hours, not bytes scanned — there's no fixed dollar-per-byte rate to convert with, so the tool declines to fabricate one rather than print a number that would be fiction.

None of these are bugs waiting for a fix. They're the same accuracy-tier discipline applied to the project's own documentation that the project applies to its output: say what you don't know instead of rounding it up to something confident-sounding. cost-guard-mcp is two days old, has one GitHub star, zero open issues, and I built it because I got tired of agents finding out a query was expensive after it already ran. An estimate that admits what it doesn't know is more useful, and safer, than one that guesses and calls itself precise — that's the whole idea, and it's the same idea whether you're reading it off a CostEstimate object or off this changelog.

Top comments (0)