Skip to content

Add querypilot proxy: Postgres wire-protocol proxy that enforces QueryPilot policies - #29

Merged
nickklos10 merged 1 commit into
mainfrom
feat/pg-wire-proxy
Oct 21, 2026
Merged

nickklos10 merged 1 commit into
mainfrom
feat/pg-wire-proxy

Conversation

@nickklos10

Copy link
Copy Markdown
Owner

Summary

querypilot proxy is a Postgres wire-protocol server that puts QueryPilot's validator and access policy in front of a real database. Any client connects to it unchanged: psql, psycopg, SQLAlchemy/LangChain, an existing MCP database server, Cursor/Claude Code, or a BI tool. Each one gets read-only enforcement, row filters, masking, row caps, audit, and errors it can repair, with no SDK integration.

pip install "querypilot[proxy]"
querypilot proxy --upstream postgresql://readonly@db/prod --listen 127.0.0.1:6543 --policy querypilot.yaml

Every statement takes exactly one route (proxy/statements.py):

  • Session commands. BEGIN/COMMIT/ROLLBACK/savepoints, SET/RESET/SHOW of allowlisted settings, and DEALLOCATE are regenerated from their tokens and never forwarded verbatim. BEGIN always becomes BEGIN ... READ ONLY.
  • Catalog queries. Queries that read only pg_catalog/information_schema and call only allowlisted functions pass through. That is what drivers, psql \d and ORM reflection need. pg_stats, pg_stat_*, pg_statistic*, query-string functions (query_to_xml, ts_stat) and non-pg_catalog functions are denied.
  • Everything else goes through SQLValidator. Only its rewritten SQL (row filters, LIMIT cap) runs upstream, inside a READ ONLY transaction, on a session with default_transaction_read_only=on and a statement timeout.
  • Anything unrecognized is rejected (fails closed). That includes COPY, PREPARE, DO, EXPLAIN, SET search_path/role, and SQL that does not parse.

Results are masked by output position from ValidationResult.output_masks. Redacted or hashed columns are described as text. Rejections are Postgres ErrorResponses with a SQLSTATE, the validator's agent-safe message, and the repair hint. Every statement goes to the audit sink as proxy_<route>.

Implementation

  • Wire protocol: pygwire, a pure-Python sans-I/O codec under MIT, maintained in 2026. riffq (Arrow-typed, pre-release) and buenavista (last released 2024) were the alternatives considered. Upstream calls use libpq at the psycopg.pq level, so values, type OIDs/typmods and text or binary encodings pass through unchanged. Every upstream call uses the extended protocol, so Postgres itself refuses any second statement smuggled into forwarded text.
  • Coverage:
    • Simple query (several statements per message).
    • Extended Parse/Bind/Describe/Execute/Close/Sync/Flush, with named statements and portals, and max_rows with PortalSuspended.
    • Transactions, savepoints, and 25P02 after an error.
    • CancelRequest, ParameterStatus after SET, NegotiateProtocolVersion, SSLRequest answered with N, and an optional password (required off loopback).
  • JSONLAuditSink(max_cached=...) bounds its in-memory copy for long-running servers.
  • Design note: docs/proxy.md covers protocol coverage, known gaps (no TLS termination, cleartext auth, no server-side cursors/EXPLAIN/COPY, LIMIT $1) and why the proxy fails closed. There is also a new README section.

Tests

  • tests/test_proxy_statements.py: 57 routing tests with no database, which always run.
  • tests/test_proxy_live.py: 25 tests against a real Postgres. The fixture uses QUERYPILOT_TEST_POSTGRES_URL if set, otherwise starts postgres:16-alpine with Docker, otherwise skips. The tests cover:
    • psycopg and psql connecting through the proxy.
    • Writes rejected, including COMMIT; DROP TABLE and BEGIN READ WRITE, with the upstream row count unchanged.
    • The row filter under inner/left/right joins and subqueries.
    • Masking through aliases, expressions, CTEs and binary results.
    • The LIMIT cap.
    • Actionable SQLSTATE and hint on rejections.
    • Transaction recovery, prepared statements, SQLAlchemy reflection, catalog queries that try to read data, portal suspension over raw protocol, the password option, and audit records.
  • Run locally against Postgres 16.2 (Docker isn't installed on the dev machine). Full suite: 899 passed, plus 3 skipped that were already skipped before this change. The Docker path of the fixture has not run yet; this PR's CI on ubuntu-latest is its first real run.

Not tested yet

node-postgres clients (e.g. the MCP reference server), JDBC (DBeaver/Metabase) and asyncpg.

🤖 Generated with Claude Code

…eryPilot policies

Any Postgres client (psql, psycopg, SQLAlchemy, MCP database servers, BI tools)
connects to the proxy unchanged. Each statement is routed through an allowlist:
regenerated session commands (BEGIN/SET/SHOW), catalog-only queries for driver
and ORM introspection, or SQLValidator, whose rewritten SQL (row filters, LIMIT
cap) runs upstream in a READ ONLY transaction. Results are masked by output
position, rejections come back as Postgres errors with the validator's message
and hint, and every statement is audited. Anything unrecognized is rejected.

- pygwire (pure-Python sans-I/O codec) for the wire protocol; libpq results
  passed through raw so types and text/binary encodings are preserved
- simple and extended protocol, named statements, portals with max_rows,
  transactions and savepoints, cancel requests, ParameterStatus
- `querypilot proxy --upstream ... --listen ... --policy querypilot.yaml`
- live tests against real Postgres (env URL or Docker, else skip), psql included
- docs/proxy.md: protocol coverage, known gaps, why it fails closed
- JSONLAuditSink(max_cached=...) bounds its in-memory copy for long-running servers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nickklos10
nickklos10 merged commit 92770c7 into main Oct 21, 2026
4 checks passed
@nickklos10
nickklos10 deleted the feat/pg-wire-proxy branch October 21, 2026 15:20
@gitguardian

gitguardian Bot commented Oct 21, 2026

Copy link
Copy Markdown

⚠️ GitGuardian has uncovered 1 secret following the scan of your pull request.

Please consider investigating the findings and remediating the incidents. Failure to do so may lead to compromising the associated services or software components.

🔎 Detected hardcoded secret in your pull request
GitGuardian id GitGuardian status Secret Commit Filename
38813836 Triggered Generic Password bfcc5c5 tests/test_proxy_live.py View secret
🛠 Guidelines to remediate hardcoded secrets
  1. Understand the implications of revoking this secret by investigating where it is used in your code.
  2. Replace and store your secret safely. Learn here the best practices.
  3. Revoke and rotate this secret.
  4. If possible, rewrite git history. Rewriting git history is not a trivial act. You might completely break other contributing developers' workflow and you risk accidentally deleting legitimate data.

To avoid such incidents in the future consider


🦉 GitGuardian detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request.

1 similar comment
@gitguardian

gitguardian Bot commented Oct 21, 2026

Copy link
Copy Markdown

⚠️ GitGuardian has uncovered 1 secret following the scan of your pull request.

Please consider investigating the findings and remediating the incidents. Failure to do so may lead to compromising the associated services or software components.

🔎 Detected hardcoded secret in your pull request
GitGuardian id GitGuardian status Secret Commit Filename
38813836 Triggered Generic Password bfcc5c5 tests/test_proxy_live.py View secret
🛠 Guidelines to remediate hardcoded secrets
  1. Understand the implications of revoking this secret by investigating where it is used in your code.
  2. Replace and store your secret safely. Learn here the best practices.
  3. Revoke and rotate this secret.
  4. If possible, rewrite git history. Rewriting git history is not a trivial act. You might completely break other contributing developers' workflow and you risk accidentally deleting legitimate data.

To avoid such incidents in the future consider


🦉 GitGuardian detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants