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.
- pgbrake-0.1.0/.gitignore +5 -0
- pgbrake-0.1.0/LICENSE +9 -0
- pgbrake-0.1.0/PKG-INFO +116 -0
- pgbrake-0.1.0/README.md +97 -0
- pgbrake-0.1.0/pyproject.toml +32 -0
- pgbrake-0.1.0/src/pgbrake/__init__.py +3 -0
- pgbrake-0.1.0/src/pgbrake/classify.py +132 -0
- pgbrake-0.1.0/src/pgbrake/cli.py +85 -0
- pgbrake-0.1.0/src/pgbrake/engine.py +273 -0
- pgbrake-0.1.0/src/pgbrake/journal.py +125 -0
- pgbrake-0.1.0/src/pgbrake/server.py +79 -0
- pgbrake-0.1.0/tests/test_brake.py +109 -0
pgbrake-0.1.0/.gitignore
ADDED
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.
|
pgbrake-0.1.0/README.md
ADDED
|
@@ -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,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
|