pg-guard-mcp 0.1.0__tar.gz

pg_guard_mcp-0.1.0/.env.example

@@ -0,0 +1,9 @@
+# Copy to .env or set in your MCP client's env config. Never commit real
+# credentials — this file documents the shape, it holds no real secret.
+
+# Full libpq connection string. Preferred over the individual PG* vars below.
+PG_GUARD_DSN=host=127.0.0.1 dbname=mydb user=myapp_readonly password=change_me
+
+# Optional tuning (defaults shown):
+# PG_GUARD_ROW_LIMIT=1000
+# PG_GUARD_STATEMENT_TIMEOUT_MS=10000

pg_guard_mcp-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.pyc
+.venv/
+.pytest_cache/
+*.egg-info/
+.env
+dist/
+build/

pg_guard_mcp-0.1.0/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.5
+Name: pg-guard-mcp
+Version: 0.1.0
+Summary: A PostgreSQL MCP server that enforces read-only access at the protocol and privilege level — not by parsing the query string.
+Author: Berkant Acun
+License: MIT
+Keywords: mcp,model-context-protocol,postgres,postgresql,read-only,security
+Requires-Python: >=3.10
+Requires-Dist: mcp>=1.2.0
+Requires-Dist: psycopg[binary]>=3.2.0
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.24.0; extra == 'dev'
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# pg-guard-mcp
+
+A PostgreSQL MCP server that enforces read-only access at the **protocol and privilege level** — not by parsing the query string and hoping.
+
+## Why this exists
+
+The official `@modelcontextprotocol/server-postgres` shipped a read-only mode that a single `COMMIT;` could bypass: it wrapped the agent's query in `BEGIN TRANSACTION READ ONLY` and sent the whole thing as one string. Postgres accepts semicolon-separated multiple statements in that mode, so `SELECT 1; COMMIT; DROP SCHEMA public CASCADE;` closed the read-only transaction early and ran the drop as an ordinary write. The package was deprecated over it. ([Datadog Security Labs writeup](https://securitylabs.datadoghq.com/articles/mcp-vulnerability-case-study-SQL-injection-in-the-postgresql-mcp-server/))
+
+pg-guard-mcp exists because that bug class — "read-only" enforced only by string inspection — is still common across the MCP ecosystem. It defends in three independent layers, so no single mistake is fatal:
+
+1. **Protocol layer (the real boundary).** Every query runs through Postgres's *extended* query protocol (`Parse`/`Bind`/`Execute`), never the simple query protocol. The extended protocol structurally rejects more than one statement per `Parse` message — Postgres itself refuses it, before any of our code runs. This is why the Datadog exploit cannot work here regardless of what string is submitted.
+2. **Session layer.** Every connection sets `default_transaction_read_only = on` at the session level, so even a query that somehow reached the database as a write is rejected by Postgres.
+3. **Pre-flight layer.** Before a query is even sent, it's checked for multiple statements and transaction-control keywords (`COMMIT`, `ROLLBACK`, `BEGIN`, `SAVEPOINT`, ...) and rejected with a clear error. This exists to fail fast and loud, not as the primary defense.
+
+On top of that, connecting with a database role that has had write privileges `REVOKE`d is the recommended (and startup-checked) setup — belt and suspenders at the privilege layer too.
+
+## Tools
+
+| Tool | Does |
+|---|---|
+| `pg_run_query(sql)` | Run one read-only statement, return rows |
+| `pg_explain_query(sql)` | Return the query plan without running it |
+| `pg_list_tables(schema="public")` | List tables/views in a schema |
+| `pg_describe_table(table_name, schema="public")` | List a table's columns |
+| `pg_check_privileges()` | Report any write grant the connected role actually holds — should always come back empty |
+
+## Setup
+
+```bash
+pip install -e ".[dev]"
+export PG_GUARD_DSN="host=127.0.0.1 dbname=mydb user=myapp_readonly password=..."
+python -m pg_guard_mcp.server
+```
+
+See `.env.example` for all supported environment variables, and `scripts/setup_dev_db.sh` for a working example of setting up a properly-restricted read-only role (the setup this project's own tests run against).
+
+## Testing
+
+```bash
+pip install -e ".[dev]"
+pytest tests/ -v
+```
+
+`tests/test_safety.py` is pure-Python and needs no database. `tests/test_db.py` and `tests/test_server.py` run against a real local PostgreSQL instance — including the exact exploit payload that deprecated the official Postgres MCP server — and skip automatically if `pgguard_test` isn't reachable. Run `scripts/setup_dev_db.sh` once to create it.
+
+## Status
+
+Early build, 58 passing tests (unit + live-Postgres integration). Not yet published to PyPI.
+
+## License
+
+MIT

pg_guard_mcp-0.1.0/README.md

@@ -0,0 +1,52 @@
+# pg-guard-mcp
+
+A PostgreSQL MCP server that enforces read-only access at the **protocol and privilege level** — not by parsing the query string and hoping.
+
+## Why this exists
+
+The official `@modelcontextprotocol/server-postgres` shipped a read-only mode that a single `COMMIT;` could bypass: it wrapped the agent's query in `BEGIN TRANSACTION READ ONLY` and sent the whole thing as one string. Postgres accepts semicolon-separated multiple statements in that mode, so `SELECT 1; COMMIT; DROP SCHEMA public CASCADE;` closed the read-only transaction early and ran the drop as an ordinary write. The package was deprecated over it. ([Datadog Security Labs writeup](https://securitylabs.datadoghq.com/articles/mcp-vulnerability-case-study-SQL-injection-in-the-postgresql-mcp-server/))
+
+pg-guard-mcp exists because that bug class — "read-only" enforced only by string inspection — is still common across the MCP ecosystem. It defends in three independent layers, so no single mistake is fatal:
+
+1. **Protocol layer (the real boundary).** Every query runs through Postgres's *extended* query protocol (`Parse`/`Bind`/`Execute`), never the simple query protocol. The extended protocol structurally rejects more than one statement per `Parse` message — Postgres itself refuses it, before any of our code runs. This is why the Datadog exploit cannot work here regardless of what string is submitted.
+2. **Session layer.** Every connection sets `default_transaction_read_only = on` at the session level, so even a query that somehow reached the database as a write is rejected by Postgres.
+3. **Pre-flight layer.** Before a query is even sent, it's checked for multiple statements and transaction-control keywords (`COMMIT`, `ROLLBACK`, `BEGIN`, `SAVEPOINT`, ...) and rejected with a clear error. This exists to fail fast and loud, not as the primary defense.
+
+On top of that, connecting with a database role that has had write privileges `REVOKE`d is the recommended (and startup-checked) setup — belt and suspenders at the privilege layer too.
+
+## Tools
+
+| Tool | Does |
+|---|---|
+| `pg_run_query(sql)` | Run one read-only statement, return rows |
+| `pg_explain_query(sql)` | Return the query plan without running it |
+| `pg_list_tables(schema="public")` | List tables/views in a schema |
+| `pg_describe_table(table_name, schema="public")` | List a table's columns |
+| `pg_check_privileges()` | Report any write grant the connected role actually holds — should always come back empty |
+
+## Setup
+
+```bash
+pip install -e ".[dev]"
+export PG_GUARD_DSN="host=127.0.0.1 dbname=mydb user=myapp_readonly password=..."
+python -m pg_guard_mcp.server
+```
+
+See `.env.example` for all supported environment variables, and `scripts/setup_dev_db.sh` for a working example of setting up a properly-restricted read-only role (the setup this project's own tests run against).
+
+## Testing
+
+```bash
+pip install -e ".[dev]"
+pytest tests/ -v
+```
+
+`tests/test_safety.py` is pure-Python and needs no database. `tests/test_db.py` and `tests/test_server.py` run against a real local PostgreSQL instance — including the exact exploit payload that deprecated the official Postgres MCP server — and skip automatically if `pgguard_test` isn't reachable. Run `scripts/setup_dev_db.sh` once to create it.
+
+## Status
+
+Early build, 58 passing tests (unit + live-Postgres integration). Not yet published to PyPI.
+
+## License
+
+MIT

pg_guard_mcp-0.1.0/pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "pg-guard-mcp"
+version = "0.1.0"
+description = "A PostgreSQL MCP server that enforces read-only access at the protocol and privilege level — not by parsing the query string."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "Berkant Acun" }]
+keywords = ["mcp", "model-context-protocol", "postgresql", "postgres", "read-only", "security"]
+dependencies = [
+    "mcp>=1.2.0",
+    "psycopg[binary]>=3.2.0",
+]
+
+[project.scripts]
+pg-guard-mcp = "pg_guard_mcp.server:main"
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.24.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/pg_guard_mcp"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"

pg_guard_mcp-0.1.0/scripts/setup_dev_db.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Sets up a local pgguard_test database and a genuinely-restricted
+# pgguard_readonly role for running the integration test suite.
+# Requires a local PostgreSQL superuser connection (adjust -U/-h as needed).
+set -euo pipefail
+
+PSQL="${PSQL:-psql}"
+SUPERUSER="${PGGUARD_SETUP_SUPERUSER:-postgres}"
+HOST="${PGGUARD_SETUP_HOST:-127.0.0.1}"
+
+"$PSQL" -h "$HOST" -U "$SUPERUSER" -v ON_ERROR_STOP=1 <<SQL
+DROP DATABASE IF EXISTS pgguard_test;
+CREATE DATABASE pgguard_test;
+DROP ROLE IF EXISTS pgguard_readonly;
+CREATE ROLE pgguard_readonly WITH LOGIN PASSWORD 'pgguard_readonly_dev_pw';
+ALTER ROLE pgguard_readonly SET default_transaction_read_only = on;
+SQL
+
+"$PSQL" -h "$HOST" -U "$SUPERUSER" -d pgguard_test -v ON_ERROR_STOP=1 <<SQL
+CREATE TABLE users (id serial PRIMARY KEY, email text NOT NULL);
+INSERT INTO users (email) VALUES ('a@example.com'), ('b@example.com');
+REVOKE ALL ON ALL TABLES IN SCHEMA public FROM pgguard_readonly;
+REVOKE ALL ON SCHEMA public FROM pgguard_readonly;
+GRANT CONNECT ON DATABASE pgguard_test TO pgguard_readonly;
+GRANT USAGE ON SCHEMA public TO pgguard_readonly;
+GRANT SELECT ON ALL TABLES IN SCHEMA public TO pgguard_readonly;
+ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO pgguard_readonly;
+SQL
+
+echo "pgguard_test database and pgguard_readonly role are ready."

pg_guard_mcp-0.1.0/src/pg_guard_mcp/__init__.py

@@ -0,0 +1,4 @@
+"""pg-guard-mcp — a PostgreSQL MCP server that enforces read-only access
+at the protocol and privilege level, not by parsing the query string."""
+
+__version__ = "0.1.0"

pg_guard_mcp-0.1.0/src/pg_guard_mcp/db.py

@@ -0,0 +1,97 @@
+"""
+The real security boundary.
+
+Every query is executed through psycopg's *extended* query protocol
+(Parse/Bind/Execute), never the simple query protocol. Postgres's extended
+protocol structurally refuses to run more than one statement inside a
+single Parse message — that is a server-side wire-protocol rule, enforced
+by Postgres itself, not a client-side check that a cleverly-crafted string
+could talk its way around.
+
+psycopg only falls back to the simple protocol when `execute()` is called
+with `params=None`. `ReadOnlyConnection.query()` always passes params
+(defaulting to an empty tuple), which forces the extended-protocol path
+even for queries that bind no parameters at all. This is the mechanism
+that makes the Datadog-documented exploit against the official
+`@modelcontextprotocol/server-postgres` — smuggling `COMMIT;` plus a write
+past a `BEGIN TRANSACTION READ ONLY` wrapper — structurally impossible
+here: Postgres rejects the multi-statement string before any of it runs.
+
+This is layered with two more independent defenses, so no single mistake
+in this file is fatal:
+- session layer: `default_transaction_read_only = on` is set right after
+  connecting, so even a write that somehow reached Postgres is refused;
+- pre-flight layer: `validate_readonly_query()` (safety.py) rejects an
+  obviously dangerous query before it's sent at all, with a clear error.
+
+The privilege layer — connecting as a role with write grants revoked — is
+enforced by Postgres itself and is the caller's responsibility to set up;
+see README.md.
+"""
+
+from __future__ import annotations
+
+from types import TracebackType
+
+import psycopg
+from psycopg.rows import dict_row
+
+from .safety import validate_readonly_query
+
+DEFAULT_STATEMENT_TIMEOUT_MS = 10_000
+DEFAULT_ROW_LIMIT = 1_000
+
+
+class ReadOnlyConnection:
+    """A single-use, read-only-enforced Postgres connection. Use as a
+    context manager: `with ReadOnlyConnection(dsn) as conn: conn.query(...)`."""
+
+    def __init__(
+        self,
+        dsn: str,
+        *,
+        statement_timeout_ms: int = DEFAULT_STATEMENT_TIMEOUT_MS,
+        row_limit: int = DEFAULT_ROW_LIMIT,
+    ) -> None:
+        self._dsn = dsn
+        self._statement_timeout_ms = statement_timeout_ms
+        self._row_limit = row_limit
+        self._conn: psycopg.Connection | None = None
+
+    def __enter__(self) -> "ReadOnlyConnection":
+        self._conn = psycopg.connect(self._dsn, row_factory=dict_row, autocommit=True)
+        with self._conn.cursor() as cur:
+            # Session layer: belt-and-suspenders even if a write somehow
+            # reached Postgres despite the protocol and pre-flight layers.
+            cur.execute("SET default_transaction_read_only = on", ())
+            # SET does not accept a bind parameter for this value (it's a
+            # utility statement, not a regular query) — coerce through
+            # int() first so there's nothing to inject even though this
+            # particular string is built with an f-string.
+            cur.execute(f"SET statement_timeout = {int(self._statement_timeout_ms)}", ())
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+
+    def query(self, sql: str, params: tuple | list | None = None) -> list[dict]:
+        if self._conn is None:
+            raise RuntimeError("ReadOnlyConnection must be used as a context manager")
+
+        # Pre-flight layer: reject obviously dangerous input with a clear,
+        # specific error before it ever reaches the network.
+        validate_readonly_query(sql)
+
+        with self._conn.cursor() as cur:
+            # Protocol layer — the real boundary. Passing params, even an
+            # empty tuple, forces psycopg's extended query protocol, which
+            # Postgres refuses to run more than one statement through.
+            cur.execute(sql, params if params is not None else ())
+            return cur.fetchmany(self._row_limit)

pg_guard_mcp-0.1.0/src/pg_guard_mcp/safety.py

@@ -0,0 +1,197 @@
+"""
+Pre-flight read-only query validation.
+
+IMPORTANT: this module is the *second* line of defense, not the first. The
+real boundary is that db.py always executes queries through Postgres's
+extended query protocol, which cannot run more than one statement per call
+no matter what string is submitted — see db.py's module docstring. This
+module exists to reject an obviously dangerous query early, with a clear
+error, instead of relying on protocol behaviour alone.
+
+It works by masking out everything that cannot affect statement structure —
+string literals, quoted identifiers, dollar-quoted bodies, and comments —
+before looking for statement-separating semicolons or dangerous keywords.
+This is what lets `WHERE message = 'a; DROP TABLE users;'` pass (the
+semicolons are just data) while `SELECT 1; DROP TABLE users;` is rejected
+(that semicolon is a real statement boundary).
+"""
+
+from __future__ import annotations
+
+import re
+
+__all__ = ["UnsafeQueryError", "validate_readonly_query"]
+
+
+class UnsafeQueryError(ValueError):
+    """Raised when a query is not a single, plain read-only statement."""
+
+
+# A statement is only accepted if it opens with one of these. This alone
+# rejects every bare write/DDL/admin statement (INSERT, DROP, GRANT, ...)
+# without needing to name each one individually.
+_ALLOWED_START_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "SHOW"}
+
+# Extra defense for dangerous keywords that can appear *inside* an
+# otherwise SELECT/WITH-shaped statement: transaction control smuggled
+# past the outer wrapper, data-modifying CTEs (`WITH d AS (DELETE ...)`),
+# and admin functions callable from a plain SELECT list.
+_DANGEROUS_KEYWORDS = [
+    "COMMIT", "ROLLBACK", "BEGIN", "SAVEPOINT", "RELEASE",
+    "SET", "INSERT", "UPDATE", "DELETE", "DROP", "TRUNCATE",
+    "ALTER", "CREATE", "GRANT", "REVOKE", "VACUUM", "CALL",
+    "COPY", "MERGE", "REINDEX", "CLUSTER", "LOCK", "DO",
+    "EXECUTE", "PREPARE", "DEALLOCATE", "LISTEN", "NOTIFY",
+    "PG_TERMINATE_BACKEND", "PG_CANCEL_BACKEND", "PG_RELOAD_CONF",
+]
+
+_DOLLAR_TAG_RE = re.compile(r"\$[A-Za-z_]*\$")
+_LEADING_WORD_RE = re.compile(r"\s*([A-Za-z_][A-Za-z_0-9]*)")
+_KEYWORD_RE_CACHE = {
+    kw: re.compile(rf"(?<![A-Za-z0-9_]){re.escape(kw)}(?![A-Za-z0-9_])", re.IGNORECASE)
+    for kw in _DANGEROUS_KEYWORDS
+}
+
+
+def _mask(sql: str) -> str:
+    """Return a string the same length as `sql` with the contents of every
+    string literal, quoted identifier, dollar-quoted body, and comment
+    replaced by spaces. Everything else is left untouched."""
+    out: list[str] = []
+    i = 0
+    n = len(sql)
+    while i < n:
+        ch = sql[i]
+
+        if ch == "'":
+            # An E'...' / e'...' string uses backslash escapes (\' , \\, ...)
+            # in addition to the standard '' doubling. A plain '...' string
+            # does not treat backslash specially at all under
+            # standard_conforming_strings=on, which has been Postgres's
+            # default since 9.1 — so backslash is ignored for those, and a
+            # lone "'" always closes them, exactly like Postgres parses it.
+            prev = sql[i - 1] if i > 0 else ""
+            prev_prev = sql[i - 2] if i > 1 else ""
+            is_e_string = prev in ("E", "e") and not (prev_prev.isalnum() or prev_prev == "_")
+
+            out.append(" ")
+            i += 1
+            while i < n:
+                if is_e_string and sql[i] == "\\" and i + 1 < n:
+                    out.append("  ")
+                    i += 2
+                    continue
+                if sql[i] == "'" and i + 1 < n and sql[i + 1] == "'":
+                    out.append("  ")
+                    i += 2
+                    continue
+                is_closing = sql[i] == "'"
+                out.append(" ")
+                i += 1
+                if is_closing:
+                    break
+            continue
+
+        if ch == '"':
+            out.append(" ")
+            i += 1
+            while i < n and sql[i] != '"':
+                out.append(" ")
+                i += 1
+            if i < n:
+                out.append(" ")
+                i += 1
+            continue
+
+        if sql[i:i + 2] == "--":
+            while i < n and sql[i] != "\n":
+                out.append(" ")
+                i += 1
+            continue
+
+        if sql[i:i + 2] == "/*":
+            out.append("  ")
+            i += 2
+            while i < n and sql[i:i + 2] != "*/":
+                out.append(" ")
+                i += 1
+            if i < n:
+                out.append("  ")
+                i += 2
+            continue
+
+        if ch == "$":
+            tag_match = _DOLLAR_TAG_RE.match(sql, i)
+            if tag_match:
+                tag = tag_match.group(0)
+                out.append(" " * len(tag))
+                i += len(tag)
+                end = sql.find(tag, i)
+                if end == -1:
+                    out.append(" " * (n - i))
+                    i = n
+                else:
+                    out.append(" " * (end - i))
+                    out.append(" " * len(tag))
+                    i = end + len(tag)
+                continue
+
+        out.append(ch)
+        i += 1
+
+    return "".join(out)
+
+
+def _split_statements(sql: str, masked: str) -> list[tuple[str, str]]:
+    """Split on semicolons that survived masking (i.e. real statement
+    separators), returning (original_text, masked_text) pairs so keyword
+    checks can run on the masked text while errors can quote the original."""
+    pairs: list[tuple[str, str]] = []
+    start = 0
+    for i, ch in enumerate(masked):
+        if ch == ";":
+            pairs.append((sql[start:i], masked[start:i]))
+            start = i + 1
+    pairs.append((sql[start:], masked[start:]))
+    return pairs
+
+
+def validate_readonly_query(sql: str) -> None:
+    """Raise UnsafeQueryError unless `sql` is exactly one plain read-only
+    statement (SELECT / WITH / EXPLAIN / SHOW) with no transaction-control
+    or write/DDL/admin keywords anywhere in it. Returns None if it's fine."""
+    if not sql or not sql.strip():
+        raise UnsafeQueryError("Empty query.")
+
+    masked = _mask(sql)
+    statements = [
+        (orig, masked_stmt)
+        for orig, masked_stmt in _split_statements(sql, masked)
+        if masked_stmt.strip()
+    ]
+
+    if not statements:
+        raise UnsafeQueryError("Empty query.")
+
+    if len(statements) > 1:
+        raise UnsafeQueryError(
+            f"Multiple statements are not allowed (found {len(statements)}). "
+            "Submit exactly one read-only statement per call."
+        )
+
+    original, masked_stmt = statements[0]
+
+    leading = _LEADING_WORD_RE.match(masked_stmt)
+    first_word = leading.group(1).upper() if leading else ""
+    if first_word not in _ALLOWED_START_KEYWORDS:
+        raise UnsafeQueryError(
+            f"Statement must start with SELECT, WITH, EXPLAIN, or SHOW — "
+            f"found '{first_word or original.strip()[:30]}'."
+        )
+
+    for keyword, pattern in _KEYWORD_RE_CACHE.items():
+        if pattern.search(masked_stmt):
+            raise UnsafeQueryError(
+                f"Query contains a disallowed keyword: {keyword}. "
+                "Only plain read-only SELECT/WITH/EXPLAIN/SHOW statements are permitted."
+            )

pg_guard_mcp-0.1.0/src/pg_guard_mcp/server.py

@@ -0,0 +1,173 @@
+"""
+MCP tool surface for pg-guard-mcp.
+
+Every tool here goes through ReadOnlyConnection (db.py), so every query —
+no matter which tool calls it — passes through the same three-layer
+defense (protocol / session / pre-flight) described there. There is
+deliberately no "run arbitrary write SQL" escape hatch and no shell-out to
+the `psql` binary: the only way this server talks to Postgres is through
+psycopg's extended query protocol.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+try:
+    # SDK >= 2.0: FastMCP was renamed to MCPServer, same .tool()/.run() API.
+    from mcp.server import MCPServer as _MCPServerImpl
+except ImportError:  # SDK < 2.0
+    from mcp.server.fastmcp import FastMCP as _MCPServerImpl  # type: ignore[no-redef]
+
+from .db import DEFAULT_ROW_LIMIT, DEFAULT_STATEMENT_TIMEOUT_MS, ReadOnlyConnection
+from .safety import UnsafeQueryError, validate_readonly_query
+
+mcp = _MCPServerImpl("pg-guard-mcp")
+
+
+def _dsn() -> str:
+    dsn = os.environ.get("PG_GUARD_DSN")
+    if dsn:
+        return dsn
+    # Fall back to standard libpq environment variables (PGHOST, PGPORT,
+    # PGDATABASE, PGUSER, PGPASSWORD, ...) — psycopg reads these itself
+    # when given an empty DSN, so this just documents the expectation.
+    return ""
+
+
+def _row_limit() -> int:
+    raw = os.environ.get("PG_GUARD_ROW_LIMIT")
+    return int(raw) if raw else DEFAULT_ROW_LIMIT
+
+
+def _statement_timeout_ms() -> int:
+    raw = os.environ.get("PG_GUARD_STATEMENT_TIMEOUT_MS")
+    return int(raw) if raw else DEFAULT_STATEMENT_TIMEOUT_MS
+
+
+def _connect() -> ReadOnlyConnection:
+    return ReadOnlyConnection(
+        _dsn(),
+        statement_timeout_ms=_statement_timeout_ms(),
+        row_limit=_row_limit(),
+    )
+
+
+def _run(sql: str, params: tuple | list | None = None) -> dict:
+    try:
+        with _connect() as conn:
+            rows = conn.query(sql, params)
+        return {"rows": rows, "row_count": len(rows)}
+    except UnsafeQueryError as e:
+        return {"error": "UnsafeQuery", "message": str(e)}
+    except Exception as e:  # psycopg errors, connection failures, etc.
+        return {"error": type(e).__name__, "message": str(e)}
+
+
+@mcp.tool()
+def pg_run_query(sql: str) -> dict:
+    """Run a single read-only SQL statement (SELECT / WITH / EXPLAIN / SHOW)
+    against the configured PostgreSQL database and return the rows.
+
+    Rejects anything that isn't exactly one plain read-only statement —
+    multiple statements, transaction-control keywords (COMMIT, ROLLBACK,
+    BEGIN, ...), and any write/DDL/admin keyword anywhere in the query are
+    all refused before the query is sent to the database. Results are
+    capped at PG_GUARD_ROW_LIMIT rows (default 1000).
+    """
+    return _run(sql)
+
+
+@mcp.tool()
+def pg_explain_query(sql: str) -> dict:
+    """Return the PostgreSQL query plan for a read-only SELECT/WITH
+    statement, without running it. Useful for checking whether a query
+    will be slow before running it for real."""
+    # Validate the user's actual input *before* wrapping it in EXPLAIN.
+    # This is deliberate: wrapping first and validating the composed
+    # string would still catch anything on the dangerous-keyword
+    # blocklist, but it makes the error about what the caller typed, and
+    # it means a change to that composition logic can never quietly
+    # start validating something other than the real input.
+    try:
+        validate_readonly_query(sql)
+    except UnsafeQueryError as e:
+        return {"error": "UnsafeQuery", "message": str(e)}
+
+    stripped = sql.strip().rstrip(";")
+    return _run(f"EXPLAIN {stripped}")
+
+
+@mcp.tool()
+def pg_list_tables(schema: str = "public") -> dict:
+    """List the tables and views visible to the connected role in the
+    given schema (default: public)."""
+    return _run(
+        "SELECT table_name, table_type FROM information_schema.tables "
+        "WHERE table_schema = %s ORDER BY table_name",
+        (schema,),
+    )
+
+
+@mcp.tool()
+def pg_check_privileges() -> dict:
+    """Report any write privilege (INSERT/UPDATE/DELETE/TRUNCATE) the
+    connected role actually holds on any table. An empty list is the
+    expected, safe result — anything else means the privilege layer of
+    defense is missing and the role should be locked down (see README.md),
+    even though the protocol and session layers still hold on their own."""
+    return _run(
+        "SELECT table_schema, table_name, privilege_type "
+        "FROM information_schema.role_table_grants "
+        "WHERE grantee = current_user "
+        "AND privilege_type IN ('INSERT', 'UPDATE', 'DELETE', 'TRUNCATE') "
+        "ORDER BY table_schema, table_name, privilege_type"
+    )
+
+
+@mcp.tool()
+def pg_describe_table(table_name: str, schema: str = "public") -> dict:
+    """List the columns of a table: name, data type, nullability, and
+    default, in column order."""
+    return _run(
+        "SELECT column_name, data_type, is_nullable, column_default "
+        "FROM information_schema.columns "
+        "WHERE table_schema = %s AND table_name = %s "
+        "ORDER BY ordinal_position",
+        (schema, table_name),
+    )
+
+
+def _warn_if_role_can_write() -> None:
+    """Best-effort startup check: if the connected role holds any write
+    privilege, say so loudly. Never fails startup — this is advisory."""
+    try:
+        result = pg_check_privileges()
+    except Exception:
+        return
+    grants = result.get("rows") or []
+    if grants:
+        print(
+            f"pg-guard-mcp: WARNING — the connected role has {len(grants)} write "
+            "grant(s) (see pg_check_privileges). The protocol and session layers "
+            "still block writes, but a role with write access revoked is the "
+            "recommended setup. See README.md.",
+            file=sys.stderr,
+        )
+
+
+def main() -> None:
+    if not os.environ.get("PG_GUARD_DSN") and not os.environ.get("PGDATABASE"):
+        print(
+            "pg-guard-mcp: no PG_GUARD_DSN or PGDATABASE set — connections will fail "
+            "until one is configured. See README.md for setup.",
+            file=sys.stderr,
+        )
+    else:
+        _warn_if_role_can_write()
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

pg_guard_mcp-0.1.0/tests/test_db.py

@@ -0,0 +1,110 @@
+"""
+Integration tests against a real local PostgreSQL instance.
+
+These prove the actual security boundary: that pg-guard-mcp's connection
+layer executes every query through the Postgres *extended* query protocol,
+which structurally refuses to run more than one statement per call — the
+exact mechanism the official @modelcontextprotocol/server-postgres lacked,
+which is what let a single `COMMIT;` bypass its read-only wrapper.
+
+Skipped automatically if PG_GUARD_TEST_DSN is not set, so the fast unit
+suite (test_safety.py) still runs everywhere with no external dependency.
+"""
+
+import os
+
+import psycopg
+import pytest
+
+from pg_guard_mcp.db import ReadOnlyConnection
+
+TEST_DSN = os.environ.get(
+    "PG_GUARD_TEST_DSN",
+    "host=127.0.0.1 dbname=pgguard_test user=pgguard_readonly password=pgguard_readonly_dev_pw",
+)
+
+
+def _connectable() -> bool:
+    try:
+        with psycopg.connect(TEST_DSN, connect_timeout=3):
+            return True
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _connectable(),
+    reason="no local pgguard_test database reachable — set PG_GUARD_TEST_DSN or run the dev setup script",
+)
+
+
+class TestOrdinaryReadsWork:
+    def test_select_returns_rows(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT id, email FROM users ORDER BY id")
+        assert len(rows) == 2
+        assert rows[0]["email"] == "a@example.com"
+
+    def test_parameterized_select_works(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT email FROM users WHERE id = %s", (1,))
+        assert rows == [{"email": "a@example.com"}]
+
+
+class TestTheDatadogExploitFailsAgainstARealServer:
+    """The exact payload that deprecated the official Postgres MCP server,
+    run against a real connection, to prove the defense is structural and
+    not just a string check that could itself have a bypass."""
+
+    def test_commit_bypass_drop_table_is_refused_by_postgres_itself(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            with pytest.raises(Exception):
+                conn.query("SELECT 1; COMMIT; DROP TABLE users;")
+
+        # Prove the table really does still exist and still has its rows —
+        # not just that *an* exception was raised.
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT count(*) AS n FROM users")
+        assert rows[0]["n"] == 2
+
+    def test_multi_statement_select_select_is_refused(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            with pytest.raises(Exception):
+                conn.query("SELECT 1; SELECT 2;")
+
+
+class TestWritesAreRefusedEvenIfSomehowSubmittedAlone:
+    """Belt-and-suspenders: even if the pre-flight keyword check in
+    safety.py were somehow bypassed, the role itself has no write grant
+    and the session is read-only, so Postgres refuses the write."""
+
+    def test_insert_is_refused_by_the_database(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            with pytest.raises(Exception):
+                conn.query("INSERT INTO users (email) VALUES ('hacker@example.com')")
+
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT count(*) AS n FROM users")
+        assert rows[0]["n"] == 2
+
+    def test_delete_is_refused_by_the_database(self):
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            with pytest.raises(Exception):
+                conn.query("DELETE FROM users")
+
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT count(*) AS n FROM users")
+        assert rows[0]["n"] == 2
+
+
+class TestConnectionRefusesUnsafeQueriesBeforeSendingThem:
+    def test_dangerous_query_never_reaches_the_network(self):
+        # safety.py should reject this before ReadOnlyConnection even opens
+        # a cursor — confirmed indirectly: it still raises, and the data is
+        # untouched, whether or not a network round-trip happened.
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            with pytest.raises(Exception):
+                conn.query("DROP TABLE users")
+        with ReadOnlyConnection(TEST_DSN) as conn:
+            rows = conn.query("SELECT count(*) AS n FROM users")
+        assert rows[0]["n"] == 2

pg_guard_mcp-0.1.0/tests/test_safety.py

@@ -0,0 +1,191 @@
+"""
+Tests for the pre-execution query safety check.
+
+This is deliberately layer 2 of the defense, not layer 1. Layer 1 (the real
+boundary) is that every query is executed through the Postgres *extended*
+query protocol, which structurally cannot run more than one statement per
+call — see db.py and test_db.py for that proof. This module exists so a
+dangerous query is rejected with a clear error *before* it ever reaches the
+network, and so an obviously-malicious payload never depends on protocol
+subtlety alone to be stopped.
+
+The exact exploit these tests are written against is the one that got the
+official @modelcontextprotocol/server-postgres deprecated: wrap the query in
+`BEGIN TRANSACTION READ ONLY`, then let the attacker's string smuggle a
+`COMMIT;` followed by a write statement, closing the read-only transaction
+early. See: https://securitylabs.datadoghq.com/articles/mcp-vulnerability-case-study-SQL-injection-in-the-postgresql-mcp-server/
+"""
+
+import pytest
+
+from pg_guard_mcp.safety import UnsafeQueryError, validate_readonly_query
+
+
+class TestObviouslySafeQueries:
+    def test_simple_select_is_allowed(self):
+        validate_readonly_query("SELECT * FROM users")
+
+    def test_select_with_where_clause_is_allowed(self):
+        validate_readonly_query("SELECT id, email FROM users WHERE active = true")
+
+    def test_select_with_join_is_allowed(self):
+        validate_readonly_query(
+            "SELECT u.id, o.total FROM users u JOIN orders o ON o.user_id = u.id"
+        )
+
+    def test_select_with_trailing_semicolon_is_allowed(self):
+        validate_readonly_query("SELECT * FROM users;")
+
+    def test_select_with_trailing_semicolon_and_whitespace_is_allowed(self):
+        validate_readonly_query("SELECT * FROM users;   \n")
+
+    def test_cte_select_is_allowed(self):
+        validate_readonly_query(
+            "WITH recent AS (SELECT * FROM orders WHERE created_at > now() - interval '1 day') "
+            "SELECT * FROM recent"
+        )
+
+    def test_explain_select_is_allowed(self):
+        validate_readonly_query("EXPLAIN SELECT * FROM users")
+
+
+class TestTheDatadogExploit:
+    """The exact real-world attack that deprecated the official Postgres MCP server."""
+
+    def test_commit_bypass_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT 1; COMMIT; DROP TABLE users;")
+
+    def test_commit_bypass_with_schema_drop_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT 1; COMMIT; DROP SCHEMA public CASCADE;")
+
+    def test_commit_alone_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("COMMIT;")
+
+    def test_rollback_alone_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("ROLLBACK;")
+
+    def test_begin_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("BEGIN; SELECT * FROM users;")
+
+    def test_savepoint_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SAVEPOINT sp1; SELECT * FROM users;")
+
+    def test_release_savepoint_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("RELEASE SAVEPOINT sp1;")
+
+    def test_set_transaction_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SET TRANSACTION READ WRITE; SELECT 1;")
+
+
+class TestMultiStatementPayloads:
+    """Even without transaction-control keywords, a second statement is refused."""
+
+    def test_two_selects_are_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT 1; SELECT 2;")
+
+    def test_select_then_insert_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT * FROM users; INSERT INTO users (id) VALUES (1);")
+
+    def test_select_then_delete_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT * FROM users; DELETE FROM users;")
+
+    def test_semicolon_inside_a_string_literal_is_still_allowed(self):
+        # A literal semicolon inside quotes is not a statement boundary.
+        validate_readonly_query("SELECT * FROM logs WHERE message = 'a; b; c'")
+
+    def test_semicolon_inside_a_string_literal_followed_by_real_second_statement_is_rejected(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query(
+                "SELECT * FROM logs WHERE message = 'a; b'; DROP TABLE logs;"
+            )
+
+
+class TestWriteStatementsAreRejectedEvenAsTheOnlyStatement:
+    @pytest.mark.parametrize(
+        "sql",
+        [
+            "INSERT INTO users (id) VALUES (1)",
+            "UPDATE users SET active = false",
+            "DELETE FROM users",
+            "DROP TABLE users",
+            "DROP SCHEMA public CASCADE",
+            "TRUNCATE users",
+            "ALTER TABLE users ADD COLUMN x int",
+            "CREATE TABLE evil (id int)",
+            "GRANT ALL ON users TO public",
+            "REVOKE ALL ON users FROM public",
+            "VACUUM users",
+            "CALL some_procedure()",
+            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity",
+        ],
+    )
+    def test_write_or_admin_statement_is_rejected(self, sql):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query(sql)
+
+
+class TestEscapeStringEdgeCases:
+    """Postgres has two incompatible quoting dialects live at once:
+    standard '...' strings (backslash is NOT special — the default since
+    PG 9.1, standard_conforming_strings=on) and E'...' strings (backslash
+    IS an escape character). Getting this wrong in either direction is
+    the classic SQL-parser-mismatch bug class. These pin down that the
+    masker matches Postgres's real behavior for both."""
+
+    def test_standard_string_backslash_is_not_an_escape(self):
+        # Under standard_conforming_strings=on (the default), '...' ends
+        # at the first single quote no matter what precedes it — a
+        # trailing backslash does NOT protect it. This mirrors real
+        # Postgres parsing, so a payload relying on backslash-escaping a
+        # standard string is correctly seen as closing early, exposing
+        # the rest as real (and here, dangerous) SQL.
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("SELECT * FROM t WHERE x = 'abc\\'; DROP TABLE users; --'")
+
+    def test_e_string_backslash_escaped_quote_stays_inside_the_string(self):
+        # E'...' DOES treat backslash as an escape. \' must not be read
+        # as the closing quote, or a real semicolon later in the same
+        # E-string would be wrongly treated as data by Postgres but as a
+        # statement boundary by an unaware masker (a false positive, not
+        # a bypass — but still wrong).
+        validate_readonly_query(r"SELECT * FROM logs WHERE msg = E'it\'s fine; still one row'")
+
+    def test_e_string_is_case_insensitive(self):
+        validate_readonly_query(r"SELECT * FROM logs WHERE msg = e'it\'s fine too'")
+
+    def test_e_string_does_not_falsely_trigger_on_a_column_named_ending_in_e(self):
+        # The "is this an E-string" check must not fire just because some
+        # unrelated identifier happens to end in e/E right before a
+        # genuinely standard string starts.
+        validate_readonly_query("SELECT * FROM t WHERE name = 'value'")
+
+    def test_dollar_quoted_body_is_masked(self):
+        validate_readonly_query("SELECT $tag$anything; DROP TABLE users;$tag$ AS literal")
+
+    def test_dollar_quoted_body_does_not_end_on_a_different_tag(self):
+        validate_readonly_query("SELECT $a$ text with $b$ inside $a$ AS literal")
+
+
+class TestErrorMessageQuality:
+    def test_error_names_the_offending_keyword(self):
+        with pytest.raises(UnsafeQueryError, match="DROP"):
+            validate_readonly_query("DROP TABLE users")
+
+    def test_error_on_empty_query_is_clear(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("")
+
+    def test_error_on_whitespace_only_query_is_clear(self):
+        with pytest.raises(UnsafeQueryError):
+            validate_readonly_query("   \n\t  ")

pg_guard_mcp-0.1.0/tests/test_server.py

@@ -0,0 +1,109 @@
+"""Integration tests for the MCP tool layer, against the same local
+pgguard_test database used by test_db.py."""
+
+import os
+
+import psycopg
+import pytest
+
+TEST_DSN = os.environ.get(
+    "PG_GUARD_TEST_DSN",
+    "host=127.0.0.1 dbname=pgguard_test user=pgguard_readonly password=pgguard_readonly_dev_pw",
+)
+
+
+def _connectable() -> bool:
+    try:
+        with psycopg.connect(TEST_DSN, connect_timeout=3):
+            return True
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _connectable(),
+    reason="no local pgguard_test database reachable — set PG_GUARD_TEST_DSN or run the dev setup script",
+)
+
+
+@pytest.fixture(autouse=True)
+def _configure_dsn(monkeypatch):
+    monkeypatch.setenv("PG_GUARD_DSN", TEST_DSN)
+    yield
+
+
+class TestPgRunQuery:
+    def test_returns_rows_for_a_plain_select(self):
+        from pg_guard_mcp.server import pg_run_query
+
+        result = pg_run_query("SELECT id, email FROM users ORDER BY id")
+        assert result["row_count"] == 2
+        assert result["rows"][0]["email"] == "a@example.com"
+
+    def test_returns_a_typed_error_instead_of_raising_for_unsafe_query(self):
+        from pg_guard_mcp.server import pg_run_query
+
+        result = pg_run_query("DROP TABLE users")
+        assert result["error"] == "UnsafeQuery"
+        assert "DROP" in result["message"]
+
+    def test_returns_a_typed_error_for_the_datadog_exploit(self):
+        from pg_guard_mcp.server import pg_run_query
+
+        result = pg_run_query("SELECT 1; COMMIT; DROP TABLE users;")
+        assert "error" in result
+
+        # and the table really is untouched
+        followup = pg_run_query("SELECT count(*) AS n FROM users")
+        assert followup["rows"][0]["n"] == 2
+
+
+class TestPgListTables:
+    def test_finds_the_users_table(self):
+        from pg_guard_mcp.server import pg_list_tables
+
+        result = pg_list_tables()
+        names = [row["table_name"] for row in result["rows"]]
+        assert "users" in names
+
+
+class TestPgDescribeTable:
+    def test_lists_columns_of_users(self):
+        from pg_guard_mcp.server import pg_describe_table
+
+        result = pg_describe_table("users")
+        names = [row["column_name"] for row in result["rows"]]
+        assert "id" in names
+        assert "email" in names
+
+
+class TestPgExplainQuery:
+    def test_returns_a_plan_without_error(self):
+        from pg_guard_mcp.server import pg_explain_query
+
+        result = pg_explain_query("SELECT * FROM users")
+        assert "error" not in result
+        assert result["row_count"] > 0
+
+    def test_rejects_a_write_statement_instead_of_explaining_it(self):
+        from pg_guard_mcp.server import pg_explain_query
+
+        # EXPLAIN ANALYZE of a write statement actually *runs* the write
+        # in real Postgres — this must never reach the database at all.
+        result = pg_explain_query("ANALYZE DELETE FROM users")
+        assert result.get("error") == "UnsafeQuery"
+
+    def test_table_is_untouched_after_a_rejected_explain_analyze_write(self):
+        from pg_guard_mcp.server import pg_explain_query, pg_run_query
+
+        pg_explain_query("ANALYZE DELETE FROM users")
+        followup = pg_run_query("SELECT count(*) AS n FROM users")
+        assert followup["rows"][0]["n"] == 2
+
+
+class TestPgCheckPrivileges:
+    def test_readonly_role_has_no_write_grants(self):
+        from pg_guard_mcp.server import pg_check_privileges
+
+        result = pg_check_privileges()
+        assert result["rows"] == []