pgbrake 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,5 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.egg-info/
4
+ dist/
5
+ .pytest_cache/
pgbrake-0.1.0/LICENSE ADDED
@@ -0,0 +1,9 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Zhang Lixiang
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
6
+
7
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
8
+
9
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
pgbrake-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,116 @@
1
+ Metadata-Version: 2.5
2
+ Name: pgbrake
3
+ Version: 0.1.0
4
+ Summary: A local safety brake between AI coding agents and your Postgres: approval for risky SQL, row limits, one-command undo, audit log.
5
+ Project-URL: Homepage, https://github.com/lexang/pgbrake
6
+ Project-URL: Issues, https://github.com/lexang/pgbrake/issues
7
+ Author: Zhang Lixiang
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: ai-agent,claude-code,codex,mcp,postgres,safety,undo
11
+ Requires-Python: >=3.10
12
+ Requires-Dist: mcp>=1.10
13
+ Requires-Dist: pglast>=6
14
+ Requires-Dist: psycopg[binary]>=3.1
15
+ Requires-Dist: pydantic>=2
16
+ Provides-Extra: dev
17
+ Requires-Dist: pytest>=8; extra == 'dev'
18
+ Description-Content-Type: text/markdown
19
+
20
+ # pgbrake
21
+
22
+ **An undo button for your database when an AI agent is driving.**
23
+
24
+ pgbrake is a local MCP server that sits between Claude Code / Codex / Cursor and your Postgres.
25
+ Every write the agent makes gets a change id you can reverse with one command. Risky SQL waits
26
+ for a human. Everything is logged. Nothing leaves your machine.
27
+
28
+ ```
29
+ agent> DELETE FROM users WHERE created_at < '2025-01-01'
30
+ pgbrake> held: DELETE would touch 4,812 rows (limit 100) — approval apv_3f9c21aa
31
+
32
+ agent> DELETE FROM users WHERE id = 42
33
+ pgbrake> 1 row affected; undo with change id chg_8e01b7d2
34
+
35
+ you$ pgbrake undo chg_8e01b7d2
36
+ undid delete on "users": 1 row(s) restored
37
+ ```
38
+
39
+ ## Why
40
+
41
+ Claude Code checkpoints rewind files, not databases. Client-side permission prompts and hooks
42
+ can be skipped or misfire for MCP calls, and a tool's own "destructive" label is just a hint.
43
+ pgbrake enforces its rules **inside the server**, so they hold no matter which agent or client calls it.
44
+
45
+ ## What it does (free, MIT)
46
+
47
+ - **Real SQL parsing** (the Postgres parser via pglast), not regex.
48
+ - **Held for approval:** DDL (`DROP`, `ALTER`, `TRUNCATE`...), `UPDATE`/`DELETE` without `WHERE`,
49
+ any write over the row limit (default 100), upserts and `MERGE`.
50
+ - **Never allowed:** `ALTER SYSTEM`, role/database creation and drops, `COPY ... PROGRAM`, `DO` blocks.
51
+ - **One-command undo:** before each `UPDATE`/`DELETE`, the affected rows are saved locally; inserts are
52
+ tracked by primary key. Undo refuses if the rows changed since, unless you `--force`.
53
+ - **Approval only by a human:** in the client's own prompt (MCP elicitation) or `pgbrake approve <id>`
54
+ in an interactive terminal. Approvals are one-shot, expire in 15 minutes, and are bound to the exact SQL.
55
+ - **Audit log** of every statement, decision and row count, in `~/.pgbrake/journal.db`.
56
+ - Multi-statement calls run in one transaction: all or nothing.
57
+
58
+ ## Install
59
+
60
+ ```bash
61
+ pip install pgbrake # or: uvx pgbrake serve
62
+ # before the PyPI release: pip install git+https://github.com/lexang/pgbrake
63
+ ```
64
+
65
+ Claude Code:
66
+
67
+ ```bash
68
+ claude mcp add pgbrake --env PGBRAKE_DSN="postgresql://me@localhost/app" -- uvx pgbrake serve
69
+ ```
70
+
71
+ Codex (`~/.codex/config.toml`):
72
+
73
+ ```toml
74
+ [mcp_servers.pgbrake]
75
+ command = "uvx"
76
+ args = ["pgbrake", "serve"]
77
+ env = { PGBRAKE_DSN = "postgresql://me@localhost/app" }
78
+ ```
79
+
80
+ Then remove any other database MCP / `psql` access from the agent, so pgbrake is the only way in.
81
+ In Claude Code, add `"Bash(psql:*)"` and `"Bash(pgbrake approve:*)"` to `permissions.deny`.
82
+
83
+ ## Commands
84
+
85
+ ```
86
+ pgbrake pending # SQL waiting for you
87
+ pgbrake approve <id> # approve (interactive terminal only)
88
+ pgbrake deny <id>
89
+ pgbrake changes # undoable changes
90
+ pgbrake undo <change-id> # reverse one
91
+ pgbrake log # audit log
92
+ ```
93
+
94
+ ## Settings (env)
95
+
96
+ | Variable | Default | |
97
+ |---|---|---|
98
+ | `PGBRAKE_DSN` | — | connection string (falls back to `DATABASE_URL`) |
99
+ | `PGBRAKE_MAX_ROWS` | 100 | writes touching more rows need approval |
100
+ | `PGBRAKE_ALLOW_DDL` | 0 | set 1 to let schema changes through (still logged) |
101
+ | `PGBRAKE_READ_ONLY` | 0 | set 1 to refuse every write |
102
+ | `PGBRAKE_STATEMENT_TIMEOUT_MS` | 15000 | |
103
+ | `PGBRAKE_JOURNAL` | `~/.pgbrake/journal.db` | |
104
+
105
+ ## Limits
106
+
107
+ - Undo covers rows changed through pgbrake. It is not a backup: keep real backups.
108
+ - `TRUNCATE`, DDL, upserts and `MERGE` are held for approval but cannot be undone.
109
+ - Updating a table with no primary key needs approval and cannot be undone.
110
+ - The free version keeps the last 20 undoable changes.
111
+
112
+ ## Pro (coming)
113
+
114
+ Unlimited undo history, a migration guard for agent terminal commands
115
+ (`prisma migrate reset --force`, `drizzle-kit push --force`, `alembic downgrade base`...: snapshot first,
116
+ then ask), per-project policy files, and audit export. One-time license.
@@ -0,0 +1,97 @@
1
+ # pgbrake
2
+
3
+ **An undo button for your database when an AI agent is driving.**
4
+
5
+ pgbrake is a local MCP server that sits between Claude Code / Codex / Cursor and your Postgres.
6
+ Every write the agent makes gets a change id you can reverse with one command. Risky SQL waits
7
+ for a human. Everything is logged. Nothing leaves your machine.
8
+
9
+ ```
10
+ agent> DELETE FROM users WHERE created_at < '2025-01-01'
11
+ pgbrake> held: DELETE would touch 4,812 rows (limit 100) — approval apv_3f9c21aa
12
+
13
+ agent> DELETE FROM users WHERE id = 42
14
+ pgbrake> 1 row affected; undo with change id chg_8e01b7d2
15
+
16
+ you$ pgbrake undo chg_8e01b7d2
17
+ undid delete on "users": 1 row(s) restored
18
+ ```
19
+
20
+ ## Why
21
+
22
+ Claude Code checkpoints rewind files, not databases. Client-side permission prompts and hooks
23
+ can be skipped or misfire for MCP calls, and a tool's own "destructive" label is just a hint.
24
+ pgbrake enforces its rules **inside the server**, so they hold no matter which agent or client calls it.
25
+
26
+ ## What it does (free, MIT)
27
+
28
+ - **Real SQL parsing** (the Postgres parser via pglast), not regex.
29
+ - **Held for approval:** DDL (`DROP`, `ALTER`, `TRUNCATE`...), `UPDATE`/`DELETE` without `WHERE`,
30
+ any write over the row limit (default 100), upserts and `MERGE`.
31
+ - **Never allowed:** `ALTER SYSTEM`, role/database creation and drops, `COPY ... PROGRAM`, `DO` blocks.
32
+ - **One-command undo:** before each `UPDATE`/`DELETE`, the affected rows are saved locally; inserts are
33
+ tracked by primary key. Undo refuses if the rows changed since, unless you `--force`.
34
+ - **Approval only by a human:** in the client's own prompt (MCP elicitation) or `pgbrake approve <id>`
35
+ in an interactive terminal. Approvals are one-shot, expire in 15 minutes, and are bound to the exact SQL.
36
+ - **Audit log** of every statement, decision and row count, in `~/.pgbrake/journal.db`.
37
+ - Multi-statement calls run in one transaction: all or nothing.
38
+
39
+ ## Install
40
+
41
+ ```bash
42
+ pip install pgbrake # or: uvx pgbrake serve
43
+ # before the PyPI release: pip install git+https://github.com/lexang/pgbrake
44
+ ```
45
+
46
+ Claude Code:
47
+
48
+ ```bash
49
+ claude mcp add pgbrake --env PGBRAKE_DSN="postgresql://me@localhost/app" -- uvx pgbrake serve
50
+ ```
51
+
52
+ Codex (`~/.codex/config.toml`):
53
+
54
+ ```toml
55
+ [mcp_servers.pgbrake]
56
+ command = "uvx"
57
+ args = ["pgbrake", "serve"]
58
+ env = { PGBRAKE_DSN = "postgresql://me@localhost/app" }
59
+ ```
60
+
61
+ Then remove any other database MCP / `psql` access from the agent, so pgbrake is the only way in.
62
+ In Claude Code, add `"Bash(psql:*)"` and `"Bash(pgbrake approve:*)"` to `permissions.deny`.
63
+
64
+ ## Commands
65
+
66
+ ```
67
+ pgbrake pending # SQL waiting for you
68
+ pgbrake approve <id> # approve (interactive terminal only)
69
+ pgbrake deny <id>
70
+ pgbrake changes # undoable changes
71
+ pgbrake undo <change-id> # reverse one
72
+ pgbrake log # audit log
73
+ ```
74
+
75
+ ## Settings (env)
76
+
77
+ | Variable | Default | |
78
+ |---|---|---|
79
+ | `PGBRAKE_DSN` | — | connection string (falls back to `DATABASE_URL`) |
80
+ | `PGBRAKE_MAX_ROWS` | 100 | writes touching more rows need approval |
81
+ | `PGBRAKE_ALLOW_DDL` | 0 | set 1 to let schema changes through (still logged) |
82
+ | `PGBRAKE_READ_ONLY` | 0 | set 1 to refuse every write |
83
+ | `PGBRAKE_STATEMENT_TIMEOUT_MS` | 15000 | |
84
+ | `PGBRAKE_JOURNAL` | `~/.pgbrake/journal.db` | |
85
+
86
+ ## Limits
87
+
88
+ - Undo covers rows changed through pgbrake. It is not a backup: keep real backups.
89
+ - `TRUNCATE`, DDL, upserts and `MERGE` are held for approval but cannot be undone.
90
+ - Updating a table with no primary key needs approval and cannot be undone.
91
+ - The free version keeps the last 20 undoable changes.
92
+
93
+ ## Pro (coming)
94
+
95
+ Unlimited undo history, a migration guard for agent terminal commands
96
+ (`prisma migrate reset --force`, `drizzle-kit push --force`, `alembic downgrade base`...: snapshot first,
97
+ then ask), per-project policy files, and audit export. One-time license.
@@ -0,0 +1,32 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "pgbrake"
7
+ version = "0.1.0"
8
+ description = "A local safety brake between AI coding agents and your Postgres: approval for risky SQL, row limits, one-command undo, audit log."
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "MIT"
12
+ authors = [{ name = "Zhang Lixiang" }]
13
+ keywords = ["mcp", "postgres", "claude-code", "codex", "ai-agent", "safety", "undo"]
14
+ dependencies = [
15
+ "mcp>=1.10",
16
+ "psycopg[binary]>=3.1",
17
+ "pglast>=6",
18
+ "pydantic>=2",
19
+ ]
20
+
21
+ [project.urls]
22
+ Homepage = "https://github.com/lexang/pgbrake"
23
+ Issues = "https://github.com/lexang/pgbrake/issues"
24
+
25
+ [project.optional-dependencies]
26
+ dev = ["pytest>=8"]
27
+
28
+ [project.scripts]
29
+ pgbrake = "pgbrake.cli:main"
30
+
31
+ [tool.hatch.build.targets.wheel]
32
+ packages = ["src/pgbrake"]
@@ -0,0 +1,3 @@
1
+ """pgbrake: a local safety brake between AI coding agents and your Postgres."""
2
+ __version__ = "0.1.0"
3
+ from .engine import Brake, Policy, Result # noqa: F401
@@ -0,0 +1,132 @@
1
+ """Classify SQL statements by risk using the real Postgres parser (pglast)."""
2
+ from __future__ import annotations
3
+
4
+ from dataclasses import dataclass, field
5
+ from enum import Enum
6
+
7
+ from pglast import parse_sql
8
+ from pglast.stream import RawStream
9
+
10
+
11
+ class Kind(str, Enum):
12
+ READ = "read" # SELECT / SHOW / EXPLAIN without ANALYZE
13
+ WRITE = "write" # INSERT / UPDATE / DELETE with a WHERE clause
14
+ BULK_WRITE = "bulk" # UPDATE / DELETE without WHERE, TRUNCATE
15
+ DDL = "ddl" # CREATE / ALTER / DROP / GRANT ...
16
+ BLOCKED = "blocked" # things we never pass through (COPY TO PROGRAM, ALTER SYSTEM...)
17
+ OTHER = "other" # SET, BEGIN, unknown -> treated as needing approval
18
+
19
+
20
+ READ_NODES = {"SelectStmt", "VariableShowStmt", "ExplainStmt"}
21
+ WRITE_NODES = {"InsertStmt", "UpdateStmt", "DeleteStmt", "MergeStmt"}
22
+ BLOCKED_NODES = {"AlterSystemStmt", "CreateRoleStmt", "AlterRoleStmt", "DropRoleStmt",
23
+ "DropdbStmt", "CreatedbStmt", "LoadStmt", "DoStmt"}
24
+ DDL_PREFIXES = ("Create", "Alter", "Drop", "Grant", "Rename", "Comment", "Truncate",
25
+ "Reindex", "Cluster", "Vacuum", "Refresh", "Lock")
26
+
27
+
28
+ @dataclass
29
+ class Statement:
30
+ sql: str
31
+ node: str
32
+ kind: Kind
33
+ table: str | None = None
34
+ schema: str | None = None
35
+ where_sql: str | None = None
36
+ has_returning: bool = False
37
+ alias: str | None = None
38
+ from_sql: str | None = None # extra FROM/USING items for join-based writes
39
+ undo_supported: bool = True
40
+ reasons: list[str] = field(default_factory=list)
41
+
42
+ @property
43
+ def qualified_table(self) -> str | None:
44
+ if not self.table:
45
+ return None
46
+ return f'"{self.schema}"."{self.table}"' if self.schema else f'"{self.table}"'
47
+
48
+
49
+ def _rangevar(rv):
50
+ if rv is None:
51
+ return None, None
52
+ return rv.relname, rv.schemaname
53
+
54
+
55
+ def _select_writes(stmt) -> bool:
56
+ """SELECT ... INTO, or data-modifying CTEs inside a SELECT."""
57
+ if getattr(stmt, "intoClause", None) is not None:
58
+ return True
59
+ wc = getattr(stmt, "withClause", None)
60
+ if wc and wc.ctes:
61
+ for cte in wc.ctes:
62
+ if type(cte.ctequery).__name__ in WRITE_NODES:
63
+ return True
64
+ return False
65
+
66
+
67
+ def classify(sql: str) -> list[Statement]:
68
+ out: list[Statement] = []
69
+ for raw in parse_sql(sql):
70
+ node = raw.stmt
71
+ name = type(node).__name__
72
+ text = RawStream()(node)
73
+ st = Statement(sql=text, node=name, kind=Kind.OTHER)
74
+
75
+ if name in READ_NODES:
76
+ st.kind = Kind.READ
77
+ if name == "SelectStmt" and _select_writes(node):
78
+ st.kind = Kind.WRITE
79
+ st.reasons.append("SELECT with data-modifying CTE / INTO")
80
+ if name == "ExplainStmt":
81
+ opts = [o.defname.lower() for o in (node.options or [])]
82
+ if "analyze" in opts:
83
+ inner = classify(RawStream()(node.query))[0]
84
+ st.kind = inner.kind if inner.kind != Kind.READ else Kind.READ
85
+ st.reasons.append("EXPLAIN ANALYZE executes the statement")
86
+ elif name in WRITE_NODES:
87
+ st.table, st.schema = _rangevar(node.relation)
88
+ if node.relation is not None and node.relation.alias is not None:
89
+ st.alias = node.relation.alias.aliasname
90
+ st.has_returning = bool(getattr(node, "returningList", None))
91
+ if name in ("UpdateStmt", "DeleteStmt"):
92
+ where = node.whereClause
93
+ if where is None:
94
+ st.kind = Kind.BULK_WRITE
95
+ st.reasons.append(f"{name[:-4].upper()} without WHERE affects every row")
96
+ else:
97
+ st.kind = Kind.WRITE
98
+ st.where_sql = RawStream()(where)
99
+ extra = getattr(node, "fromClause", None) or getattr(node, "usingClause", None)
100
+ if extra:
101
+ st.from_sql = ", ".join(RawStream()(x) for x in extra)
102
+ elif name == "MergeStmt":
103
+ st.kind = Kind.BULK_WRITE
104
+ st.undo_supported = False
105
+ st.reasons.append("MERGE can touch many rows and is not undoable")
106
+ else:
107
+ st.kind = Kind.WRITE
108
+ occ = getattr(node, "onConflictClause", None)
109
+ if occ is not None and int(occ.action) == 2:
110
+ st.undo_supported = False
111
+ st.reasons.append("upsert (ON CONFLICT DO UPDATE) is not undoable")
112
+ elif name == "CopyStmt":
113
+ if getattr(node, "is_program", False) or (not node.is_from and node.filename):
114
+ st.kind = Kind.BLOCKED
115
+ st.reasons.append("COPY to/from server file or program")
116
+ else:
117
+ st.kind = Kind.BULK_WRITE if node.is_from else Kind.READ
118
+ elif name in BLOCKED_NODES:
119
+ st.kind = Kind.BLOCKED
120
+ st.reasons.append(f"{name} is never allowed through pgbrake")
121
+ elif name == "TruncateStmt":
122
+ st.kind = Kind.BULK_WRITE
123
+ st.undo_supported = False
124
+ st.reasons.append("TRUNCATE removes every row and cannot be undone by pgbrake")
125
+ elif name.startswith(DDL_PREFIXES):
126
+ st.kind = Kind.DDL
127
+ st.reasons.append(f"schema change ({name})")
128
+ elif name == "TransactionStmt":
129
+ st.kind = Kind.BLOCKED
130
+ st.reasons.append("pgbrake manages transactions itself; send one statement group per call")
131
+ out.append(st)
132
+ return out
@@ -0,0 +1,85 @@
1
+ """pgbrake command line: serve the MCP server, approve held SQL, inspect and undo changes."""
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import datetime as dt
6
+ import json
7
+ import os
8
+ import sys
9
+
10
+ from .journal import Journal
11
+
12
+
13
+ def _ts(t):
14
+ return dt.datetime.fromtimestamp(t).strftime("%Y-%m-%d %H:%M:%S") if t else "-"
15
+
16
+
17
+ def main(argv=None):
18
+ ap = argparse.ArgumentParser(prog="pgbrake", description="Safety brake for AI agents on Postgres")
19
+ sub = ap.add_subparsers(dest="cmd", required=True)
20
+ s = sub.add_parser("serve", help="run the MCP server (stdio)")
21
+ s.add_argument("--dsn")
22
+ sub.add_parser("pending", help="list SQL waiting for approval")
23
+ a = sub.add_parser("approve", help="approve held SQL")
24
+ a.add_argument("id")
25
+ d = sub.add_parser("deny", help="deny held SQL")
26
+ d.add_argument("id")
27
+ c = sub.add_parser("changes", help="list undoable changes")
28
+ c.add_argument("-n", type=int, default=20)
29
+ u = sub.add_parser("undo", help="undo a change")
30
+ u.add_argument("id")
31
+ u.add_argument("--dsn")
32
+ u.add_argument("--force", action="store_true")
33
+ lg = sub.add_parser("log", help="show the audit log")
34
+ lg.add_argument("-n", type=int, default=30)
35
+ lg.add_argument("--json", action="store_true")
36
+ args = ap.parse_args(argv)
37
+
38
+ if args.cmd == "serve":
39
+ from .server import main as serve
40
+ return serve(args.dsn)
41
+
42
+ j = Journal()
43
+ if args.cmd == "pending":
44
+ rows = j.list_pending()
45
+ if not rows:
46
+ print("nothing waiting")
47
+ for p in rows:
48
+ print(f"{p['id']} {_ts(p['ts'])} {', '.join(json.loads(p['reasons']))}\n {p['sql']}\n")
49
+ elif args.cmd in ("approve", "deny"):
50
+ p = j.get_pending(args.id)
51
+ if not p:
52
+ sys.exit(f"unknown approval id {args.id}")
53
+ print(f"SQL:\n {p['sql']}\nReasons: {', '.join(json.loads(p['reasons']))}")
54
+ if args.cmd == "approve":
55
+ # Agents run shell commands without a TTY. Requiring one keeps an agent
56
+ # from approving its own SQL via `pgbrake approve`.
57
+ if not (sys.stdin.isatty() and sys.stdout.isatty()):
58
+ sys.exit("pgbrake approve must be run by a human in an interactive terminal")
59
+ if input("Approve? [y/N] ").strip().lower() != "y":
60
+ sys.exit("not approved")
61
+ ok = j.decide(args.id, "approved" if args.cmd == "approve" else "denied")
62
+ print(("approved" if args.cmd == "approve" else "denied") if ok else "already decided")
63
+ elif args.cmd == "changes":
64
+ for ch in j.changes(args.n):
65
+ state = f"undone {_ts(ch['undone_at'])}" if ch["undone_at"] else "undoable"
66
+ print(f"{ch['id']} {_ts(ch['ts'])} {ch['op']:6} {ch['tbl']} {ch['rows']} row(s) [{state}]\n {ch['sql']}")
67
+ elif args.cmd == "undo":
68
+ from .engine import Brake
69
+ dsn = args.dsn or os.environ.get("PGBRAKE_DSN") or os.environ.get("DATABASE_URL")
70
+ if not dsn:
71
+ sys.exit("set PGBRAKE_DSN or pass --dsn")
72
+ r = Brake(dsn, journal=j, client="cli").undo(args.id, force=args.force)
73
+ print(r.message)
74
+ sys.exit(0 if r.status == "ok" else 1)
75
+ elif args.cmd == "log":
76
+ rows = j.audit(args.n)
77
+ if args.json:
78
+ print(json.dumps(rows, indent=1, default=str))
79
+ else:
80
+ for r in reversed(rows):
81
+ print(f"{_ts(r['ts'])} {r['decision']:14} {r['kind'] or '':6} {r['rowcount'] if r['rowcount'] is not None else '':>5} {r['sql'][:120]}")
82
+
83
+
84
+ if __name__ == "__main__":
85
+ main()
@@ -0,0 +1,273 @@
1
+ """The brake: decide, capture before-images, execute, undo.
2
+
3
+ All enforcement lives here (server side), so it holds even when the MCP
4
+ client's own permission prompts or hooks are skipped.
5
+ """
6
+ from __future__ import annotations
7
+
8
+ import json
9
+ import os
10
+ import time
11
+ from dataclasses import dataclass, field
12
+ from typing import Any
13
+
14
+ import psycopg
15
+ from psycopg import sql as S
16
+ from psycopg.rows import dict_row
17
+ from psycopg.types.json import Jsonb
18
+
19
+ from .classify import Kind, Statement, classify
20
+ from .journal import Journal, sql_hash
21
+
22
+
23
+ @dataclass
24
+ class Policy:
25
+ max_rows: int = int(os.environ.get("PGBRAKE_MAX_ROWS", "100"))
26
+ allow_ddl: bool = os.environ.get("PGBRAKE_ALLOW_DDL", "0") == "1"
27
+ read_only: bool = os.environ.get("PGBRAKE_READ_ONLY", "0") == "1"
28
+ statement_timeout_ms: int = int(os.environ.get("PGBRAKE_STATEMENT_TIMEOUT_MS", "15000"))
29
+ max_result_rows: int = 500
30
+
31
+
32
+ @dataclass
33
+ class Result:
34
+ status: str # ok | needs_approval | denied | error
35
+ message: str = ""
36
+ rows: list[dict] = field(default_factory=list)
37
+ rowcount: int | None = None
38
+ change_ids: list[str] = field(default_factory=list)
39
+ pending_id: str | None = None
40
+ reasons: list[str] = field(default_factory=list)
41
+
42
+ def to_dict(self) -> dict[str, Any]:
43
+ d = {k: v for k, v in self.__dict__.items() if v not in (None, [], "")}
44
+ return json.loads(json.dumps(d, default=str))
45
+
46
+
47
+ class NeedsApproval(Exception):
48
+ def __init__(self, reasons):
49
+ self.reasons = reasons
50
+
51
+
52
+ class Brake:
53
+ def __init__(self, dsn: str, journal: Journal | None = None, policy: Policy | None = None,
54
+ client: str | None = None):
55
+ self.dsn = dsn
56
+ self.journal = journal or Journal()
57
+ self.policy = policy or Policy()
58
+ self.client = client
59
+ self.db_label = psycopg.conninfo.conninfo_to_dict(dsn).get("dbname") if "=" in dsn or "://" in dsn else dsn
60
+
61
+ def _connect(self):
62
+ return psycopg.connect(self.dsn, row_factory=dict_row, autocommit=False,
63
+ options=f"-c statement_timeout={self.policy.statement_timeout_ms}")
64
+
65
+ # ---- decision ----------------------------------------------------
66
+ def _static_reasons(self, stmts: list[Statement]) -> tuple[list[str], list[str]]:
67
+ deny, approve = [], []
68
+ for st in stmts:
69
+ if st.kind == Kind.BLOCKED:
70
+ deny += st.reasons or [f"{st.node} blocked"]
71
+ elif self.policy.read_only and st.kind != Kind.READ:
72
+ deny.append(f"read-only mode: {st.node} refused")
73
+ elif st.kind == Kind.DDL and not self.policy.allow_ddl:
74
+ approve += st.reasons
75
+ elif st.kind in (Kind.BULK_WRITE, Kind.OTHER):
76
+ approve += st.reasons or [f"{st.node} needs approval"]
77
+ elif st.kind == Kind.WRITE and not st.undo_supported:
78
+ approve += st.reasons
79
+ return deny, approve
80
+
81
+ def run(self, sql: str, approval_id: str | None = None) -> Result:
82
+ try:
83
+ stmts = classify(sql)
84
+ except Exception as e: # parse error: refuse rather than guess
85
+ r = Result("denied", f"could not parse SQL: {e}")
86
+ self.journal.log(sql=sql, decision="denied", message=r.message, client=self.client, db=self.db_label)
87
+ return r
88
+ if not stmts:
89
+ return Result("error", "empty statement")
90
+ kind = max((s.kind.value for s in stmts), key=lambda k: list(Kind._value2member_map_).index(k))
91
+
92
+ deny, approve = self._static_reasons(stmts)
93
+ if deny:
94
+ r = Result("denied", "refused by pgbrake policy", reasons=deny)
95
+ self.journal.log(sql=sql, kind=kind, decision="denied", message="; ".join(deny),
96
+ client=self.client, db=self.db_label)
97
+ return r
98
+
99
+ approved = False
100
+ if approval_id:
101
+ p = self.journal.get_pending(approval_id)
102
+ if not p or p["sql_hash"] != sql_hash(sql):
103
+ return Result("denied", "approval id does not match this exact SQL")
104
+ if p["status"] != "approved":
105
+ return Result("needs_approval", f"approval {approval_id} is {p['status']}; a human must approve it",
106
+ pending_id=approval_id)
107
+ if p["expires"] < time.time():
108
+ return Result("denied", "approval expired; request a new one")
109
+ approved = True
110
+ elif approve:
111
+ return self._pend(sql, kind, approve)
112
+
113
+ try:
114
+ return self._execute(sql, stmts, kind, approved, approval_id)
115
+ except NeedsApproval as na:
116
+ return self._pend(sql, kind, na.reasons)
117
+ except psycopg.Error as e:
118
+ msg = str(e).strip()
119
+ self.journal.log(sql=sql, kind=kind, decision="error", message=msg, client=self.client, db=self.db_label)
120
+ return Result("error", msg)
121
+
122
+ def _pend(self, sql, kind, reasons) -> Result:
123
+ pid = self.journal.add_pending(db=self.db_label, sql=sql, reasons=reasons)
124
+ self.journal.log(sql=sql, kind=kind, decision="needs_approval", pending_id=pid,
125
+ message="; ".join(reasons), client=self.client, db=self.db_label)
126
+ return Result("needs_approval",
127
+ f"Held for human approval. A person must run `pgbrake approve {pid}` "
128
+ f"(or accept the prompt in their client), then call execute again with approval_id.",
129
+ pending_id=pid, reasons=reasons)
130
+
131
+ # ---- execution ---------------------------------------------------
132
+ def _pk(self, conn, st: Statement) -> list[str]:
133
+ rows = conn.execute(
134
+ """SELECT a.attname FROM pg_index i
135
+ JOIN pg_attribute a ON a.attrelid=i.indrelid AND a.attnum = ANY(i.indkey)
136
+ WHERE i.indrelid = %s::regclass AND i.indisprimary ORDER BY a.attnum""",
137
+ (st.qualified_table,)).fetchall()
138
+ return [r["attname"] for r in rows]
139
+
140
+ def _capture(self, conn, st: Statement) -> list[dict]:
141
+ target = S.SQL(st.qualified_table)
142
+ ref = S.Identifier(st.alias) if st.alias else target
143
+ from_ = S.SQL("{} AS {}").format(target, S.Identifier(st.alias)) if st.alias else target
144
+ if st.from_sql:
145
+ from_ = S.SQL("{}, {}").format(from_, S.SQL(st.from_sql))
146
+ q = S.SQL("SELECT to_jsonb({ref}) AS r FROM {from_} WHERE {where} FOR UPDATE OF {lock}").format(
147
+ ref=ref, from_=from_, where=S.SQL(st.where_sql),
148
+ lock=S.Identifier(st.alias) if st.alias else S.Identifier(st.table))
149
+ seen, out = set(), []
150
+ for r in conn.execute(q).fetchall(): # join-based writes can repeat target rows
151
+ key = json.dumps(r["r"], sort_keys=True, default=str)
152
+ if key not in seen:
153
+ seen.add(key)
154
+ out.append(r["r"])
155
+ return out
156
+
157
+ def _refetch(self, conn, st: Statement, pk: list[str], rows: list[dict]) -> list[dict]:
158
+ if not rows or not pk:
159
+ return []
160
+ q = S.SQL("SELECT to_jsonb(t) AS r FROM {tbl} t JOIN jsonb_populate_recordset(NULL::{tbl}, %s) k "
161
+ "USING ({pk})").format(tbl=S.SQL(st.qualified_table), pk=S.SQL(", ").join(map(S.Identifier, pk)))
162
+ return [r["r"] for r in conn.execute(q, (Jsonb(rows),)).fetchall()]
163
+
164
+ def _execute(self, sql, stmts, kind, approved, approval_id) -> Result:
165
+ change_ids: list[str] = []
166
+ total = 0
167
+ out_rows: list[dict] = []
168
+ with self._connect() as conn:
169
+ if kind == Kind.READ.value:
170
+ conn.read_only = True
171
+ pending_changes = []
172
+ for st in stmts:
173
+ before, pk = [], []
174
+ if st.kind in (Kind.WRITE, Kind.BULK_WRITE) and st.table and st.undo_supported:
175
+ pk = self._pk(conn, st)
176
+ if st.node in ("UpdateStmt", "DeleteStmt"):
177
+ if st.where_sql is None: # bulk, approved: capture whole table
178
+ st.where_sql = "true"
179
+ before = self._capture(conn, st)
180
+ if len(before) > self.policy.max_rows and not approved:
181
+ conn.rollback()
182
+ raise NeedsApproval([f"{st.node[:-4].upper()} would touch {len(before)} rows "
183
+ f"(limit {self.policy.max_rows})"])
184
+ if st.node == "UpdateStmt" and not pk:
185
+ if not approved:
186
+ conn.rollback()
187
+ raise NeedsApproval([f"{st.table} has no primary key; UPDATE cannot be undone"])
188
+ exec_sql = st.sql
189
+ if st.node == "InsertStmt" and st.undo_supported and st.table:
190
+ exec_sql = f"WITH _pgbrake AS ({st.sql}{'' if st.has_returning else ' RETURNING *'}) " \
191
+ f"SELECT to_jsonb(_pgbrake) AS r FROM _pgbrake"
192
+ cur = conn.execute(exec_sql)
193
+ rc = cur.rowcount
194
+ rows = cur.fetchall() if cur.description else []
195
+ if st.node == "InsertStmt" and st.undo_supported and st.table:
196
+ inserted = [r["r"] for r in rows]
197
+ rows = inserted if st.has_returning else []
198
+ rc = len(inserted)
199
+ if rc > self.policy.max_rows and not approved:
200
+ conn.rollback()
201
+ raise NeedsApproval([f"INSERT of {rc} rows (limit {self.policy.max_rows})"])
202
+ if pk and inserted and all(k in inserted[0] for k in pk):
203
+ pending_changes.append(dict(tbl=st.qualified_table, op="insert", pk_cols=pk,
204
+ before=[], after=[{k: r[k] for k in pk} for r in inserted],
205
+ sql=st.sql))
206
+ elif st.node == "UpdateStmt" and before:
207
+ after = self._refetch(conn, st, pk, [{k: r[k] for k in pk} for r in before])
208
+ pending_changes.append(dict(tbl=st.qualified_table, op="update", pk_cols=pk,
209
+ before=before, after=after, sql=st.sql))
210
+ elif st.node == "DeleteStmt" and before:
211
+ pending_changes.append(dict(tbl=st.qualified_table, op="delete", pk_cols=pk,
212
+ before=before, after=[], sql=st.sql))
213
+ if st.kind in (Kind.WRITE, Kind.BULK_WRITE) and rc and rc > self.policy.max_rows and not approved:
214
+ conn.rollback()
215
+ raise NeedsApproval([f"{st.node} affected {rc} rows (limit {self.policy.max_rows})"])
216
+ total += max(rc, 0) if rc is not None else 0
217
+ out_rows = rows[: self.policy.max_result_rows]
218
+ conn.commit()
219
+ for ch in pending_changes:
220
+ change_ids.append(self.journal.add_change(db=self.db_label, **ch))
221
+ if approval_id:
222
+ self.journal.decide(approval_id, "used")
223
+ self.journal.log(sql=sql, kind=kind, decision="approved" if approved else "allowed",
224
+ rowcount=total, change_ids=change_ids, client=self.client, db=self.db_label)
225
+ msg = f"{total} row(s) affected" if kind != Kind.READ.value else f"{len(out_rows)} row(s)"
226
+ if change_ids:
227
+ msg += f"; undo with change id {', '.join(change_ids)}"
228
+ return Result("ok", msg, rows=out_rows, rowcount=total, change_ids=change_ids)
229
+
230
+ # ---- undo ----------------------------------------------------------
231
+ def undo(self, change_id: str, force: bool = False) -> Result:
232
+ ch = self.journal.get_change(change_id)
233
+ if not ch:
234
+ return Result("error", f"unknown change {change_id} (free tier keeps the last 20 changes)")
235
+ if ch["undone_at"]:
236
+ return Result("error", f"{change_id} was already undone")
237
+ tbl, pk = S.SQL(ch["tbl"]), ch["pk_cols"]
238
+ pkl = S.SQL(", ").join(map(S.Identifier, pk))
239
+ with self._connect() as conn:
240
+ if ch["op"] in ("update", "insert") and not force and ch["after"]:
241
+ # conflict check: rows must still look the way pgbrake left them
242
+ cur = conn.execute(S.SQL(
243
+ "SELECT count(*) AS n FROM {t} t JOIN jsonb_populate_recordset(NULL::{t}, %s) k USING ({pk})"
244
+ ).format(t=tbl, pk=pkl), (Jsonb(ch["after"]),)).fetchone()["n"]
245
+ if cur != len(ch["after"]):
246
+ return Result("error", f"{len(ch['after']) - cur} row(s) changed or vanished since; "
247
+ f"re-run with force=true to undo anyway")
248
+ cols = [r["attname"] for r in conn.execute(
249
+ "SELECT attname FROM pg_attribute WHERE attrelid=%s::regclass AND attnum>0 AND NOT attisdropped "
250
+ "AND attgenerated='' ORDER BY attnum", (ch["tbl"],)).fetchall()]
251
+ cl = S.SQL(", ").join(map(S.Identifier, cols))
252
+ if ch["op"] == "delete":
253
+ q = S.SQL("INSERT INTO {t} ({c}) OVERRIDING SYSTEM VALUE SELECT {c} FROM "
254
+ "jsonb_populate_recordset(NULL::{t}, %s)").format(t=tbl, c=cl)
255
+ n = conn.execute(q, (Jsonb(ch["before"]),)).rowcount
256
+ elif ch["op"] == "update":
257
+ setcols = [c for c in cols if c not in pk]
258
+ q = S.SQL("UPDATE {t} AS t SET ({c}) = ({rc}) FROM jsonb_populate_recordset(NULL::{t}, %s) AS r "
259
+ "WHERE {on}").format(
260
+ t=tbl, c=S.SQL(", ").join(map(S.Identifier, setcols)),
261
+ rc=S.SQL(", ").join(S.SQL("r.{}").format(S.Identifier(c)) for c in setcols),
262
+ on=S.SQL(" AND ").join(S.SQL("t.{0} = r.{0}").format(S.Identifier(k)) for k in pk))
263
+ n = conn.execute(q, (Jsonb(ch["before"]),)).rowcount
264
+ else: # insert
265
+ q = S.SQL("DELETE FROM {t} AS t USING jsonb_populate_recordset(NULL::{t}, %s) AS r WHERE {on}").format(
266
+ t=tbl, on=S.SQL(" AND ").join(S.SQL("t.{0} = r.{0}").format(S.Identifier(k)) for k in pk))
267
+ n = conn.execute(q, (Jsonb(ch["after"]),)).rowcount
268
+ conn.commit()
269
+ self.journal.mark_undone(change_id)
270
+ self.journal.log(sql=f"-- undo {change_id}", kind="undo", decision="allowed", rowcount=n,
271
+ change_ids=[change_id], client=self.client, db=self.db_label)
272
+ return Result("ok", f"undid {ch['op']} on {ch['tbl']}: {n} row(s) restored", rowcount=n,
273
+ change_ids=[change_id])
@@ -0,0 +1,125 @@
1
+ """Local SQLite journal: audit log, undoable changes, pending approvals.
2
+
3
+ Everything stays on the developer's machine (default ~/.pgbrake/journal.db).
4
+ """
5
+ from __future__ import annotations
6
+
7
+ import hashlib
8
+ import json
9
+ import os
10
+ import secrets
11
+ import sqlite3
12
+ import time
13
+ from pathlib import Path
14
+
15
+ FREE_UNDO_HISTORY = 20 # free tier keeps the last N undoable changes
16
+ APPROVAL_TTL_SECONDS = 15 * 60
17
+
18
+ SCHEMA = """
19
+ CREATE TABLE IF NOT EXISTS audit (
20
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
21
+ ts REAL NOT NULL, client TEXT, db TEXT, sql TEXT NOT NULL, kind TEXT,
22
+ decision TEXT NOT NULL, rowcount INTEGER, change_ids TEXT, pending_id TEXT, message TEXT
23
+ );
24
+ CREATE TABLE IF NOT EXISTS changes (
25
+ id TEXT PRIMARY KEY, ts REAL NOT NULL, db TEXT, tbl TEXT NOT NULL, op TEXT NOT NULL,
26
+ pk_cols TEXT NOT NULL, before TEXT, after TEXT, sql TEXT, undone_at REAL
27
+ );
28
+ CREATE TABLE IF NOT EXISTS pending (
29
+ id TEXT PRIMARY KEY, ts REAL NOT NULL, db TEXT, sql TEXT NOT NULL, sql_hash TEXT NOT NULL,
30
+ reasons TEXT, status TEXT NOT NULL DEFAULT 'pending', decided_at REAL, expires REAL NOT NULL
31
+ );
32
+ """
33
+
34
+
35
+ def default_path() -> Path:
36
+ return Path(os.environ.get("PGBRAKE_JOURNAL", Path.home() / ".pgbrake" / "journal.db"))
37
+
38
+
39
+ def sql_hash(sql: str) -> str:
40
+ return hashlib.sha256(sql.strip().encode()).hexdigest()
41
+
42
+
43
+ def new_id(prefix: str) -> str:
44
+ return f"{prefix}_{secrets.token_hex(4)}"
45
+
46
+
47
+ class Journal:
48
+ def __init__(self, path: str | Path | None = None, undo_history: int | None = FREE_UNDO_HISTORY):
49
+ self.path = Path(path) if path else default_path()
50
+ self.path.parent.mkdir(parents=True, exist_ok=True)
51
+ self.undo_history = undo_history
52
+ self.db = sqlite3.connect(self.path, isolation_level=None, check_same_thread=False)
53
+ self.db.row_factory = sqlite3.Row
54
+ self.db.executescript(SCHEMA)
55
+
56
+ # ---- audit -------------------------------------------------------
57
+ def log(self, *, sql, decision, kind=None, rowcount=None, change_ids=None,
58
+ pending_id=None, message=None, client=None, db=None):
59
+ self.db.execute(
60
+ "INSERT INTO audit(ts,client,db,sql,kind,decision,rowcount,change_ids,pending_id,message)"
61
+ " VALUES (?,?,?,?,?,?,?,?,?,?)",
62
+ (time.time(), client, db, sql, kind, decision, rowcount,
63
+ json.dumps(change_ids or []), pending_id, message))
64
+
65
+ def audit(self, limit=50):
66
+ return [dict(r) for r in self.db.execute(
67
+ "SELECT * FROM audit ORDER BY id DESC LIMIT ?", (limit,))]
68
+
69
+ # ---- changes -----------------------------------------------------
70
+ def add_change(self, *, db, tbl, op, pk_cols, before, after, sql) -> str:
71
+ cid = new_id("chg")
72
+ self.db.execute(
73
+ "INSERT INTO changes(id,ts,db,tbl,op,pk_cols,before,after,sql) VALUES (?,?,?,?,?,?,?,?,?)",
74
+ (cid, time.time(), db, tbl, op, json.dumps(pk_cols),
75
+ json.dumps(before, default=str), json.dumps(after, default=str), sql))
76
+ self._prune()
77
+ return cid
78
+
79
+ def _prune(self):
80
+ if self.undo_history is None:
81
+ return
82
+ self.db.execute(
83
+ "DELETE FROM changes WHERE id NOT IN (SELECT id FROM changes ORDER BY ts DESC LIMIT ?)",
84
+ (self.undo_history,))
85
+
86
+ def get_change(self, cid):
87
+ r = self.db.execute("SELECT * FROM changes WHERE id=?", (cid,)).fetchone()
88
+ if not r:
89
+ return None
90
+ d = dict(r)
91
+ for k in ("pk_cols", "before", "after"):
92
+ d[k] = json.loads(d[k]) if d[k] else None
93
+ return d
94
+
95
+ def changes(self, limit=20):
96
+ return [{k: r[k] for k in ("id", "ts", "tbl", "op", "sql", "undone_at")} |
97
+ {"rows": len(json.loads(r["before"] or "[]") or json.loads(r["after"] or "[]"))}
98
+ for r in self.db.execute("SELECT * FROM changes ORDER BY ts DESC LIMIT ?", (limit,))]
99
+
100
+ def mark_undone(self, cid):
101
+ self.db.execute("UPDATE changes SET undone_at=? WHERE id=?", (time.time(), cid))
102
+
103
+ # ---- approvals ---------------------------------------------------
104
+ def add_pending(self, *, db, sql, reasons) -> str:
105
+ pid = new_id("apv")
106
+ now = time.time()
107
+ self.db.execute(
108
+ "INSERT INTO pending(id,ts,db,sql,sql_hash,reasons,expires) VALUES (?,?,?,?,?,?,?)",
109
+ (pid, now, db, sql, sql_hash(sql), json.dumps(reasons), now + APPROVAL_TTL_SECONDS))
110
+ return pid
111
+
112
+ def get_pending(self, pid):
113
+ r = self.db.execute("SELECT * FROM pending WHERE id=?", (pid,)).fetchone()
114
+ return dict(r) if r else None
115
+
116
+ def list_pending(self):
117
+ return [dict(r) for r in self.db.execute(
118
+ "SELECT * FROM pending WHERE status='pending' AND expires>? ORDER BY ts", (time.time(),))]
119
+
120
+ def decide(self, pid, status):
121
+ assert status in ("approved", "denied", "used")
122
+ cur = self.db.execute(
123
+ "UPDATE pending SET status=?, decided_at=? WHERE id=? AND status IN ('pending','approved')",
124
+ (status, time.time(), pid))
125
+ return cur.rowcount == 1
@@ -0,0 +1,79 @@
1
+ """MCP server exposing pgbrake tools over stdio."""
2
+ from __future__ import annotations
3
+
4
+ import json
5
+ import os
6
+
7
+ from pydantic import BaseModel, Field
8
+
9
+ try: # mcp >= 2
10
+ from mcp.server.mcpserver import Context, MCPServer as _Server
11
+ except ImportError: # mcp 1.x
12
+ from mcp.server.fastmcp import Context, FastMCP as _Server
13
+
14
+ from .engine import Brake
15
+
16
+ INSTRUCTIONS = """pgbrake guards a Postgres database. Use `query` for reads and `execute` for writes.
17
+ Writes return a change id that `undo` can reverse. Risky SQL (DDL, writes without WHERE,
18
+ writes over the row limit) is held for HUMAN approval: tell the user the approval id and
19
+ wait; never try to approve it yourself or work around the brake with other tools."""
20
+
21
+
22
+ class Approve(BaseModel):
23
+ approve: bool = Field(description="Run this SQL against the database?")
24
+
25
+
26
+ def build(dsn: str | None = None) -> _Server:
27
+ dsn = dsn or os.environ.get("PGBRAKE_DSN") or os.environ.get("DATABASE_URL")
28
+ if not dsn:
29
+ raise SystemExit("pgbrake: set PGBRAKE_DSN (or DATABASE_URL) or pass --dsn")
30
+ brake = Brake(dsn, client=os.environ.get("PGBRAKE_CLIENT", "mcp"))
31
+ mcp = _Server("pgbrake", instructions=INSTRUCTIONS)
32
+
33
+ def dump(r):
34
+ return json.dumps(r.to_dict(), ensure_ascii=False, indent=1)
35
+
36
+ @mcp.tool(description="Run a read-only SQL query (SELECT/EXPLAIN). Writes are rejected here.")
37
+ def query(sql: str) -> str:
38
+ brake.policy.read_only, prev = True, brake.policy.read_only
39
+ try:
40
+ return dump(brake.run(sql))
41
+ finally:
42
+ brake.policy.read_only = prev
43
+
44
+ @mcp.tool(description="Run SQL that may modify data. Returns rows, a change id for undo, or an "
45
+ "approval id when a human must approve first.")
46
+ async def execute(sql: str, ctx: Context, approval_id: str | None = None) -> str:
47
+ r = brake.run(sql, approval_id=approval_id)
48
+ if r.status == "needs_approval" and r.pending_id and not approval_id:
49
+ try: # ask the human directly if the client supports elicitation
50
+ ans = await ctx.elicit(
51
+ message=f"pgbrake is holding this SQL ({'; '.join(r.reasons)}):\n\n{sql}", schema=Approve)
52
+ if getattr(ans, "action", None) == "accept" and ans.data and ans.data.approve:
53
+ brake.journal.decide(r.pending_id, "approved")
54
+ return dump(brake.run(sql, approval_id=r.pending_id))
55
+ if getattr(ans, "action", None) in ("decline", "accept"):
56
+ brake.journal.decide(r.pending_id, "denied")
57
+ return dump(r.__class__("denied", "the user declined this SQL", pending_id=r.pending_id))
58
+ except Exception:
59
+ pass # client has no elicitation: fall back to `pgbrake approve <id>` in a terminal
60
+ return dump(r)
61
+
62
+ @mcp.tool(description="Reverse a change made through pgbrake, by change id.")
63
+ def undo(change_id: str, force: bool = False) -> str:
64
+ return dump(brake.undo(change_id, force=force))
65
+
66
+ @mcp.tool(description="List recent changes that can be undone.")
67
+ def list_changes(limit: int = 20) -> str:
68
+ return json.dumps(brake.journal.changes(limit), default=str, indent=1)
69
+
70
+ @mcp.tool(description="Check whether an approval id has been approved by a human yet.")
71
+ def approval_status(approval_id: str) -> str:
72
+ p = brake.journal.get_pending(approval_id)
73
+ return json.dumps({"approval_id": approval_id, "status": p["status"] if p else "unknown"})
74
+
75
+ return mcp
76
+
77
+
78
+ def main(dsn: str | None = None):
79
+ build(dsn).run()
@@ -0,0 +1,109 @@
1
+ import os
2
+
3
+ import psycopg
4
+ import pytest
5
+
6
+ from pgbrake.classify import Kind, classify
7
+ from pgbrake.engine import Brake, Policy
8
+ from pgbrake.journal import Journal
9
+
10
+ DSN = os.environ.get("PGBRAKE_TEST_DSN", "host=/tmp port=5544 user=postgres dbname=demo")
11
+
12
+
13
+ def test_classify():
14
+ k = lambda s: [x.kind for x in classify(s)]
15
+ assert k("select 1") == [Kind.READ]
16
+ assert k("delete from t where id=1") == [Kind.WRITE]
17
+ assert k("delete from t") == [Kind.BULK_WRITE]
18
+ assert k("update t set a=1") == [Kind.BULK_WRITE]
19
+ assert k("drop table t") == [Kind.DDL]
20
+ assert k("truncate t") == [Kind.BULK_WRITE]
21
+ assert k("alter system set x=1") == [Kind.BLOCKED]
22
+ assert k("copy t to program 'rm -rf /'") == [Kind.BLOCKED]
23
+ assert k("with d as (delete from t returning *) select * from d") == [Kind.WRITE]
24
+ assert k("begin") == [Kind.BLOCKED]
25
+
26
+
27
+ @pytest.fixture
28
+ def brake(tmp_path):
29
+ with psycopg.connect(DSN, autocommit=True) as c:
30
+ c.execute("drop table if exists users; create table users(id serial primary key, email text, plan text)")
31
+ c.execute("insert into users(email, plan) select 'u'||g||'@x.com', 'free' from generate_series(1,300) g")
32
+ return Brake(DSN, journal=Journal(tmp_path / "j.db"), policy=Policy(max_rows=100))
33
+
34
+
35
+ def count(where="true"):
36
+ with psycopg.connect(DSN) as c:
37
+ return c.execute(f"select count(*) from users where {where}").fetchone()[0]
38
+
39
+
40
+ def test_read(brake):
41
+ r = brake.run("select count(*) as n from users")
42
+ assert r.status == "ok" and r.rows[0]["n"] == 300
43
+
44
+
45
+ def test_delete_and_undo(brake):
46
+ r = brake.run("delete from users where id <= 10")
47
+ assert r.status == "ok" and r.rowcount == 10 and count() == 290
48
+ u = brake.undo(r.change_ids[0])
49
+ assert u.status == "ok" and count() == 300
50
+
51
+
52
+ def test_update_and_undo_with_alias(brake):
53
+ r = brake.run("update users as x set plan='pro' where x.id between 1 and 5")
54
+ assert r.status == "ok" and count("plan='pro'") == 5
55
+ assert brake.undo(r.change_ids[0]).status == "ok"
56
+ assert count("plan='pro'") == 0
57
+
58
+
59
+ def test_insert_and_undo(brake):
60
+ r = brake.run("insert into users(email, plan) values ('a@b.c','free'), ('d@e.f','free')")
61
+ assert r.status == "ok" and count() == 302
62
+ assert brake.undo(r.change_ids[0]).status == "ok" and count() == 300
63
+
64
+
65
+ def test_row_limit_holds_then_human_approves(brake):
66
+ r = brake.run("delete from users where id <= 200")
67
+ assert r.status == "needs_approval" and count() == 300
68
+ again = brake.run("delete from users where id <= 200", approval_id=r.pending_id)
69
+ assert again.status == "needs_approval" # agent cannot self-approve
70
+ brake.journal.decide(r.pending_id, "approved") # human: `pgbrake approve <id>`
71
+ ok = brake.run("delete from users where id <= 200", approval_id=r.pending_id)
72
+ assert ok.status == "ok" and count() == 100
73
+ assert brake.run("delete from users where id <= 200", approval_id=r.pending_id).status != "ok" # one-shot
74
+ assert brake.undo(ok.change_ids[0]).status == "ok" and count() == 300
75
+
76
+
77
+ def test_approval_bound_to_exact_sql(brake):
78
+ r = brake.run("drop table users")
79
+ brake.journal.decide(r.pending_id, "approved")
80
+ assert brake.run("drop table users cascade", approval_id=r.pending_id).status == "denied"
81
+
82
+
83
+ def test_bulk_and_ddl_held(brake):
84
+ assert brake.run("delete from users").status == "needs_approval"
85
+ assert brake.run("drop table users").status == "needs_approval"
86
+ assert count() == 300
87
+
88
+
89
+ def test_blocked(brake):
90
+ assert brake.run("alter system set work_mem='1GB'").status == "denied"
91
+
92
+
93
+ def test_undo_conflict_detection(brake):
94
+ r = brake.run("update users set plan='pro' where id=1")
95
+ with psycopg.connect(DSN, autocommit=True) as c:
96
+ c.execute("delete from users where id=1")
97
+ assert brake.undo(r.change_ids[0]).status == "error"
98
+
99
+
100
+ def test_multi_statement_atomic(brake):
101
+ r = brake.run("delete from users where id=1; delete from nope where id=1")
102
+ assert r.status == "error" and count() == 300
103
+
104
+
105
+ def test_audit_log(brake):
106
+ brake.run("select 1")
107
+ brake.run("delete from users")
108
+ decisions = [a["decision"] for a in brake.journal.audit()]
109
+ assert "needs_approval" in decisions and "allowed" in decisions