From cc4066b038359d52d2b1e955c160e0fd96b394f0 Mon Sep 17 00:00:00 2001 From: tmphost-dev Date: Sat, 24 Oct 2026 01:24:50 +0000 Subject: [PATCH 1/2] Let the owner delegate marking PII columns to the agent A database's agentMarksPii setting (PUT /databases/:id/pii/agent, a switch on the PII columns tab) offers the MCP tool a markPii action. It only adds: every column must exist, a marked column stays marked, and only the owner can unmark one. The setting is checked on every call. The PII editor now saves its changes onto the server's latest set, so a save doesn't drop columns the agent marked while it was open. PII columns are inserted in batches: the documented maximum of 10000 hit SQLite's bound-variable limit in one INSERT. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 4 + README.md | 5 +- backend/internal/app/app_test.go | 16 +++ backend/internal/controller/crud/crud.go | 103 ++++++++++++++-- backend/internal/controller/crud/crud_test.go | 113 ++++++++++++++++++ backend/internal/controller/crud/query.go | 7 +- .../internal/controller/mcpserver/server.go | 84 +++++++++++-- .../controller/mcpserver/server_test.go | 59 +++++++++ backend/internal/data/entity/entity.go | 16 ++- backend/internal/data/model/model.go | 7 +- backend/internal/data/view/view.go | 7 +- backend/internal/handler/api/api.go | 13 ++ backend/internal/handler/mcp/mcp.go | 2 +- backend/internal/repository/repository.go | 30 ++++- backend/internal/testutil/testutil.go | 8 +- docs/API.md | 15 ++- frontend/src/api.js | 6 + frontend/src/components/PiiEditor.vue | 47 +++++++- frontend/src/composables/usePiiSelection.js | 14 +++ frontend/src/views/DatabaseDetailView.vue | 8 +- frontend/test/PiiEditor.test.js | 54 ++++++++- frontend/test/api.test.js | 3 + frontend/test/usePiiSelection.test.js | 15 +++ 23 files changed, 598 insertions(+), 38 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dc2f1ca..2ebf10c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,6 +47,10 @@ app state). A handler never calls the repository directly. don't switch to the visitor. The guard parses as the server reads only because the driver forces `sql_mode` (no `ANSI_QUOTES`, `NO_BACKSLASH_ESCAPES`) and `utf8mb4`; a DSN must never loosen those. +- **An agent can only add PII marks** (`markPii`, when the owner turns on + `agentMarksPii`). Never give MCP a way to unmark; that would let an agent + lift its own guard. The editor saves its diff onto the server's latest set, + so a save doesn't drop what the agent marked meanwhile. - **OAuth registrations, consents and codes are signed tokens, not rows.** Deleting a database revokes its grants because every token is checked against the database's current owner. diff --git a/README.md b/README.md index 74a8e49..6c2298e 100644 --- a/README.md +++ b/README.md @@ -98,7 +98,9 @@ including the port, for example `https://sql.example.com:3443`. Its database is where a table with no schema is looked up, and a PII column marked with no schema lives there. 2. On the **PII columns** tab, tick the columns to protect. Columns whose - names look like PII (email, phone, ssn…) are suggested. + names look like PII (email, phone, ssn…) are suggested. Turn on **Let the + agent mark PII columns** to have MCP clients mark more as they find them; + an agent can only add marks, and only you can remove one. 3. On the **Connect** tab, copy the MCP URL (`/mcp/`) and add it to your client, for example `claude mcp add --transport http shop https://sql.example.com/mcp/`. @@ -114,6 +116,7 @@ The MCP server has one tool, `sql`, which takes `{action, payload}`: | `listTables` | none | every table and view, with its PII columns | | `describeTable` | `{table, schema?}` | columns, types, PII flags | | `query` | `{sql, store?}` | one read-only `SELECT` | +| `markPii` | `{table, schema?, columns}` | mark columns as PII; offered only when the owner allows it, and it never unmarks | | `help` | `{topic?}` | usage | The **Query** tab runs the same checks from the browser, and **Activity** logs diff --git a/backend/internal/app/app_test.go b/backend/internal/app/app_test.go index 6587715..9aeb94d 100644 --- a/backend/internal/app/app_test.go +++ b/backend/internal/app/app_test.go @@ -255,6 +255,10 @@ func TestDatabasesAPI(t *testing.T) { assert.Equal(t, http.StatusNotFound, s.do(other, "PUT", db+"/pii", map[string]any{"piiColumns": []any{}}).status) assert.Contains(t, s.do(c, "GET", db+"/pii", nil).body, `"column":"email"`) assert.Equal(t, http.StatusNotFound, s.do(other, "GET", db+"/pii", nil).status) + res = s.do(c, "PUT", db+"/pii/agent", map[string]bool{"agentMarksPii": true}) + assert.Equal(t, true, res.json(t)["database"].(map[string]any)["agentMarksPii"]) + assert.Equal(t, http.StatusBadRequest, s.do(c, "PUT", db+"/pii/agent", "{", "Content-Type", "application/json").status) + assert.Equal(t, http.StatusNotFound, s.do(other, "PUT", db+"/pii/agent", map[string]bool{"agentMarksPii": false}).status) assert.Contains(t, s.do(c, "GET", db+"/schema", nil).body, `"name":"email","dataType":"text","pii":true`) assert.Equal(t, http.StatusNotFound, s.do(other, "GET", db+"/schema", nil).status) @@ -432,6 +436,18 @@ func TestMCPOAuthFlow(t *testing.T) { require.Equal(t, http.StatusOK, res.status, res.body) assert.Contains(t, res.body, `rowCount: 3`) + // markPii works once the owner turns it on, and only adds. + mark := `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"sql","arguments":{"action":"markPii","payload":{"table":"orders","columns":["total"]}}}}` + mcpHeaders := []string{"Authorization", "Bearer " + access, "Content-Type", "application/json", "Accept", "application/json, text/event-stream", "Mcp-Protocol-Version", "2025-06-18"} + res = s.do(api, "POST", mcpURL, mark, mcpHeaders...) + assert.Contains(t, res.body, "the owner has not let agents mark PII columns") + require.Equal(t, http.StatusOK, s.do(c, "PUT", "/api/v1/databases/"+id+"/pii/agent", map[string]bool{"agentMarksPii": true}).status) + res = s.do(api, "POST", mcpURL, initBody, mcpHeaders...) + assert.Contains(t, res.body, "mark it with markPii") + res = s.do(api, "POST", mcpURL, mark, mcpHeaders...) + assert.Contains(t, res.body, `columns[1]{name,pii}:\n total,marked`) + assert.Contains(t, s.do(c, "GET", "/api/v1/databases/"+id+"/pii", nil).body, `"column":"total"`) + // Deleting the database revokes every grant. require.Equal(t, http.StatusNoContent, s.do(c, "DELETE", "/api/v1/databases/"+id, nil).status) assert.Equal(t, http.StatusNotFound, s.do(api, "POST", mcpURL, initBody, "Authorization", "Bearer "+access).status) diff --git a/backend/internal/controller/crud/crud.go b/backend/internal/controller/crud/crud.go index f52bbed..7e8d0b7 100644 --- a/backend/internal/controller/crud/crud.go +++ b/backend/internal/controller/crud/crud.go @@ -142,14 +142,15 @@ func (c *Controller) MCPURL(id int64) string { func (c *Controller) toDatabase(d *model.Database, piiCount int) *entity.Database { return &entity.Database{ - ID: FormatID(d.ID), - Name: d.Name, - Driver: d.Driver, - DSNMasked: d.DSNMasked, - PIICount: piiCount, - MCPURL: c.MCPURL(d.ID), - CreatedAt: d.CreatedAt, - UpdatedAt: d.UpdatedAt, + ID: FormatID(d.ID), + Name: d.Name, + Driver: d.Driver, + DSNMasked: d.DSNMasked, + PIICount: piiCount, + AgentMarksPII: d.AgentMarksPII, + MCPURL: c.MCPURL(d.ID), + CreatedAt: d.CreatedAt, + UpdatedAt: d.UpdatedAt, } } @@ -392,3 +393,89 @@ func (c *Controller) SetPII(ctx context.Context, userID, id int64, columns []ent } return c.ListPII(ctx, userID, id) } + +// SetAgentMarksPII lets, or stops, a database's MCP agent adding PII columns. +func (c *Controller) SetAgentMarksPII(ctx context.Context, userID, id int64, on bool) (*entity.Database, error) { + if err := c.repo.SetAgentMarksPII(ctx, userID, id, on); err != nil { + return nil, mapErr(err) + } + return c.GetDatabase(ctx, userID, id) +} + +// ErrAgentMarksPIIOff is an agent marking PII on a database whose owner has +// not let it. +var ErrAgentMarksPIIOff = entity.Invalid("the owner has not let agents mark PII columns on this database; ask them to turn it on in the PII columns tab") + +// AgentMarkPII marks columns of one table as PII for the database's MCP +// agent, when its owner allows it. It only ever adds: a column already +// marked stays marked, and nothing an agent sends unmarks one. Every column +// must exist, so an agent can't fill the set with names that aren't there. +func (c *Controller) AgentMarkPII(ctx context.Context, userID, id int64, schema, table string, columns []string) (*entity.PIIMarked, error) { + row, err := c.repo.GetDatabase(ctx, userID, id) + if err != nil { + return nil, mapErr(err) + } + if !row.AgentMarksPII { + return nil, ErrAgentMarksPIIOff + } + schema, table = strings.TrimSpace(schema), strings.TrimSpace(table) + if table == "" || len(columns) == 0 { + return nil, entity.Invalid("markPii needs a table and at least one column") + } + tables, err := c.Schema(ctx, userID, id) + if err != nil { + return nil, err + } + var found []entity.SchemaTable + for _, t := range tables { + if strings.EqualFold(t.Name, table) && (schema == "" || strings.EqualFold(t.Schema, schema)) { + found = append(found, t) + } + } + switch { + case len(found) == 0: + return nil, entity.Invalid("there is no table " + table + "; listTables lists them") + case len(found) > 1: + return nil, entity.Invalid(table + " is in more than one schema; name the schema") + } + t := found[0] + out := &entity.PIIMarked{Table: t.Schema + "." + t.Name, Marked: []string{}, Already: []string{}} + var rows []model.PIIColumn + seen := map[string]bool{} + for _, name := range columns { + name = strings.TrimSpace(name) + var col *entity.SchemaColumn + for i := range t.Columns { + if strings.EqualFold(t.Columns[i].Name, name) { + col = &t.Columns[i] + break + } + } + switch { + case col == nil: + return nil, entity.Invalid(out.Table + " has no column " + name + "; describeTable lists them") + case seen[strings.ToLower(col.Name)]: + continue + case col.PII: + out.Already = append(out.Already, col.Name) + default: + out.Marked = append(out.Marked, col.Name) + rows = append(rows, model.PIIColumn{Schema: t.Schema, Table: t.Name, Column: col.Name}) + } + seen[strings.ToLower(col.Name)] = true + } + if len(rows) == 0 { + return out, nil + } + counts, err := c.repo.CountPIIColumns(ctx, []int64{id}) + if err == nil && counts[id]+len(rows) > maxPIIColumns { + err = entity.Invalid("too many PII columns") + } + if err == nil { + err = c.repo.AddPIIColumns(ctx, id, rows) + } + if err != nil { + return nil, err + } + return out, nil +} diff --git a/backend/internal/controller/crud/crud_test.go b/backend/internal/controller/crud/crud_test.go index 31bb1a9..1da7ee0 100644 --- a/backend/internal/controller/crud/crud_test.go +++ b/backend/internal/controller/crud/crud_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "net/url" + "strconv" "strings" "testing" "time" @@ -17,6 +18,7 @@ import ( "github.com/agentrq/sql-mcp/backend/internal/controller/crud" "github.com/agentrq/sql-mcp/backend/internal/data/entity" "github.com/agentrq/sql-mcp/backend/internal/data/model" + "github.com/agentrq/sql-mcp/backend/internal/service/dbdriver" "github.com/agentrq/sql-mcp/backend/internal/testutil" ) @@ -198,6 +200,117 @@ func TestSchemaAndPII(t *testing.T) { assert.ErrorAs(t, err, &failed) } +func TestAgentMarksPII(t *testing.T) { + env := testutil.New(t) + alice := signIn(t, env, "alice@example.com") + bob := signIn(t, env, "bob@example.com") + id := newDB(t, env, alice) + + _, err := env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"email"}) + assert.ErrorIs(t, err, crud.ErrAgentMarksPIIOff, "off until the owner turns it on") + + _, err = env.Ctrl.SetAgentMarksPII(ctx, bob, id, true) + assert.ErrorIs(t, err, entity.ErrNotFound) + d, err := env.Ctrl.SetAgentMarksPII(ctx, alice, id, true) + require.NoError(t, err) + assert.True(t, d.AgentMarksPII) + owner, err := env.Ctrl.SystemDatabaseOwner(ctx, id) + require.NoError(t, err) + assert.True(t, owner.AgentMarksPII) + + _, err = env.Ctrl.AgentMarkPII(ctx, bob, id, "", "users", []string{"email"}) + assert.ErrorIs(t, err, entity.ErrNotFound) + + _, err = env.Ctrl.SetPII(ctx, alice, id, []entity.PIIColumn{{Table: "orders", Column: "total"}}) + require.NoError(t, err) + + got, err := env.Ctrl.AgentMarkPII(ctx, alice, id, "", "USERS", []string{" EMAIL ", "email"}) + require.NoError(t, err) + assert.Equal(t, &entity.PIIMarked{Table: "public.users", Marked: []string{"email"}, Already: []string{}}, got) + got, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "public", "orders", []string{"total", "id"}) + require.NoError(t, err) + assert.Equal(t, &entity.PIIMarked{Table: "public.orders", Marked: []string{"id"}, Already: []string{"total"}}, got) + got, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "orders", []string{"total"}) + require.NoError(t, err) + assert.Empty(t, got.Marked, "nothing new to mark") + + cols, err := env.Ctrl.ListPII(ctx, alice, id) + require.NoError(t, err) + assert.Equal(t, []entity.PIIColumn{ + {Schema: "public", Table: "orders", Column: "id"}, + {Schema: "public", Table: "orders", Column: "total"}, + {Schema: "public", Table: "users", Column: "email"}, + }, cols, "the agent only adds; the owner's column stays") + + env.Driver.Extra = []dbdriver.Table{{Schema: "billing", Name: "users", Kind: "table", Columns: []dbdriver.Column{{Name: "plan"}}}} + for name, tc := range map[string]struct { + schema, table string + columns []string + want string + }{ + "no table": {"", " ", []string{"email"}, "needs a table"}, + "no columns": {"", "users", nil, "at least one column"}, + "unknown": {"", "nope", []string{"email"}, "there is no table nope"}, + "wrong schema": {"billing", "orders", []string{"id"}, "there is no table orders"}, + "ambiguous": {"", "users", []string{"email"}, "more than one schema"}, + "no column": {"public", "users", []string{"email", "ssn"}, "public.users has no column ssn"}, + } { + _, err := env.Ctrl.AgentMarkPII(ctx, alice, id, tc.schema, tc.table, tc.columns) + var invalid *entity.InvalidError + require.ErrorAs(t, err, &invalid, name) + assert.Contains(t, invalid.Message, tc.want, name) + } + got, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "billing", "users", []string{"plan"}) + require.NoError(t, err) + assert.Equal(t, "billing.users", got.Table) + + env.Driver.SchemaErr = errors.New("permission denied") + _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "orders", []string{"id"}) + var failed *entity.QueryFailedError + assert.ErrorAs(t, err, &failed) + env.Driver.SchemaErr = nil + + d, err = env.Ctrl.SetAgentMarksPII(ctx, alice, id, false) + require.NoError(t, err) + assert.False(t, d.AgentMarksPII) + _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"id"}) + assert.ErrorIs(t, err, crud.ErrAgentMarksPIIOff) +} + +func TestAgentMarkPIIWhenTheStoreFails(t *testing.T) { + env := testutil.New(t) + alice := signIn(t, env, "alice@example.com") + id := newDB(t, env, alice) + _, err := env.Ctrl.SetAgentMarksPII(ctx, alice, id, true) + require.NoError(t, err) + + db, err := gorm.Open(sqlite.Open(env.DBPath), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.Exec("CREATE TRIGGER no_pii BEFORE INSERT ON pii_columns BEGIN SELECT RAISE(FAIL, 'read-only'); END").Error) + _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"email"}) + assert.ErrorContains(t, err, "read-only") + + require.NoError(t, db.Exec("DROP TABLE pii_columns").Error) + _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"email"}) + assert.Error(t, err) +} + +func TestTooManyAgentPIIColumns(t *testing.T) { + env := testutil.New(t) + alice := signIn(t, env, "alice@example.com") + id := newDB(t, env, alice) + _, err := env.Ctrl.SetAgentMarksPII(ctx, alice, id, true) + require.NoError(t, err) + full := make([]entity.PIIColumn, 10000) + for i := range full { + full[i] = entity.PIIColumn{Table: "t", Column: "c" + strconv.Itoa(i)} + } + _, err = env.Ctrl.SetPII(ctx, alice, id, full) + require.NoError(t, err) + _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"email"}) + assert.ErrorContains(t, err, "too many PII columns") +} + func TestRunQuery(t *testing.T) { env := testutil.New(t) alice := signIn(t, env, "alice@example.com") diff --git a/backend/internal/controller/crud/query.go b/backend/internal/controller/crud/query.go index f9506cc..3c9fe10 100644 --- a/backend/internal/controller/crud/query.go +++ b/backend/internal/controller/crud/query.go @@ -143,11 +143,14 @@ func queryLogEntity(q *model.QueryLog) entity.QueryLog { } } -// DatabaseOwner is who owns a database, its name and its engine. +// DatabaseOwner is who owns a database, its name, its engine and whether +// its agent may mark PII. type DatabaseOwner struct { UserID int64 Name string Driver string + // AgentMarksPII is whether the owner lets the MCP agent add PII columns. + AgentMarksPII bool } // SystemDatabaseOwner returns who owns a database, for the MCP server to @@ -157,7 +160,7 @@ func (c *Controller) SystemDatabaseOwner(ctx context.Context, id int64) (Databas if err != nil { return DatabaseOwner{}, mapErr(err) } - return DatabaseOwner{UserID: row.UserID, Name: row.Name, Driver: row.Driver}, nil + return DatabaseOwner{UserID: row.UserID, Name: row.Name, Driver: row.Driver, AgentMarksPII: row.AgentMarksPII}, nil } // ToCSV writes a result as CSV, the header first. diff --git a/backend/internal/controller/mcpserver/server.go b/backend/internal/controller/mcpserver/server.go index 009d6a0..8342273 100644 --- a/backend/internal/controller/mcpserver/server.go +++ b/backend/internal/controller/mcpserver/server.go @@ -29,6 +29,7 @@ import ( type Backend interface { Schema(ctx context.Context, userID, id int64) ([]entity.SchemaTable, error) RunQuery(ctx context.Context, rq crud.QueryRequest) (*entity.QueryResult, error) + AgentMarkPII(ctx context.Context, userID, id int64, schema, table string, columns []string) (*entity.PIIMarked, error) } // Caller is who a request is from and which database it is for. @@ -39,6 +40,9 @@ type Caller struct { // Driver is the engine's ID ("postgres", "mysql"), so the agent writes // the right dialect. Driver string + // AgentMarksPII offers the markPii action. The controller checks the + // database again on every call, so this only decides what is listed. + AgentMarksPII bool } type callerKey struct{} @@ -67,6 +71,21 @@ const ToolDescription = "Read-only SQL on this database. Takes {action, payload} "query {sql, store?} (one SELECT; PII columns, SELECT * over a table with PII, and writes are rejected), " + "help {topic?}. Answers are TOON: a table names its columns once, then one line per row. Start with listTables." +// MarkPIIToolDescription is the `sql` tool's description when the owner +// lets the agent mark PII columns. +const MarkPIIToolDescription = "Read-only SQL on this database. Takes {action, payload}. " + + "Actions: listTables (every table and view, and which columns are PII), " + + "describeTable {table, schema?} (columns and types), " + + "query {sql, store?} (one SELECT; PII columns, SELECT * over a table with PII, and writes are rejected), " + + "markPii {table, schema?, columns} (mark columns holding personal data as PII; a mark can't be undone), " + + "help {topic?}. Answers are TOON: a table names its columns once, then one line per row. Start with listTables." + +// markPIIInstructions is added to the server's instructions when the owner +// lets the agent mark PII columns. +const markPIIInstructions = "The owner lets you mark PII columns: when you find a column holding personal data " + + "(names, emails, phone numbers, addresses, birth dates, government IDs, IPs, free text about a person) that is not marked yet, " + + "mark it with markPii. Only the owner can unmark a column. " + // dialect tells the agent which SQL to write, by driver ID. The MySQL // driver also serves MariaDB; its sql_mode has no ANSI_QUOTES and sets // NO_BACKSLASH_ESCAPES (see sqlguard), which is what the hint describes. @@ -78,25 +97,31 @@ var dialect = map[string]string{ } type input struct { - Action string `json:"action" jsonschema:"listTables, describeTable, query or help"` + Action string `json:"action" jsonschema:"listTables, describeTable, query, markPii (when the owner allows it) or help"` Payload map[string]any `json:"payload,omitempty" jsonschema:"the action's fields; help prints them"` } // NewServer is the MCP server for one caller and database. func NewServer(backend Backend, caller Caller, version string) *mcp.Server { + description, marking := ToolDescription, "" + if caller.AgentMarksPII { + description, marking = MarkPIIToolDescription, markPIIInstructions + } srv := mcp.NewServer(&mcp.Implementation{Name: "sql-mcp", Title: "SQL MCP: " + caller.DatabaseName, Version: version}, &mcp.ServerOptions{ Instructions: fmt.Sprintf("SQL MCP for the database %q. %s"+ "Use the `sql` tool: call listTables first, then query with explicit column lists. "+ - "Columns marked PII cannot be selected, filtered, sorted or joined on; a query that touches one is rejected with the reason. "+ + "Columns marked PII cannot be selected, filtered, sorted or joined on; a query that touches one is rejected with the reason. %s"+ "Large results come back cut to a preview with a temporary tmphost.io URL holding every row. "+ - "Answers are TOON: `key: value` lines, and tables as `name[rows]{col,col}:` followed by one comma-separated line per row.", caller.DatabaseName, dialect[caller.Driver]), + "Answers are TOON: `key: value` lines, and tables as `name[rows]{col,col}:` followed by one comma-separated line per row.", caller.DatabaseName, dialect[caller.Driver], marking), }) t := &tool{backend: backend, caller: caller} mcp.AddTool(srv, &mcp.Tool{ Name: "sql", Title: "SQL", - Description: ToolDescription, - Annotations: &mcp.ToolAnnotations{Title: "Query " + caller.DatabaseName, ReadOnlyHint: true, IdempotentHint: true}, + Description: description, + // markPii changes the PII set, so the tool is no longer read-only; + // it never touches the database itself. + Annotations: &mcp.ToolAnnotations{Title: "Query " + caller.DatabaseName, ReadOnlyHint: !caller.AgentMarksPII, IdempotentHint: true}, }, t.handle) return srv } @@ -116,6 +141,12 @@ type queryPayload struct { Store bool `json:"store,omitempty"` } +type markPayload struct { + Table string `json:"table"` + Schema string `json:"schema,omitempty"` + Columns []string `json:"columns"` +} + type helpPayload struct { Topic string `json:"topic,omitempty"` } @@ -127,8 +158,14 @@ var help = map[string]string{ "query": "query {\"sql\": \"\", \"store\": } runs the query read-only. Refused: any statement but SELECT, more than one statement, " + "a PII column anywhere (select list, WHERE, JOIN, ORDER BY), SELECT * or t.* or whole-row references over a table with PII, and functions that run SQL from text. " + "The first rows come back inline as a TOON rows table after rowCount and the other totals; a larger result (or store: true) is stored whole as CSV on the owner's tmphost.io and its URL returned.", + "markPii": "markPii {\"table\": \"\", \"schema\": \"\", \"columns\": [\"email\", \"phone\"]} marks columns of one table as PII, " + + "so no query can read them from then on. It works only when the owner allows it. It only adds: a marked column stays marked, and only the owner can unmark one. " + + "Every column must exist. It returns a columns table: each column's name and whether it was marked now or was PII already.", } +// markPIIUsage is the help line for markPii, listed only when it is offered. +const markPIIUsage = "\n markPii {\"table\": \"users\", \"columns\": [\"email\"]} mark columns as PII; only the owner can unmark." + func (t *tool) handle(ctx context.Context, _ *mcp.CallToolRequest, in input) (*mcp.CallToolResult, any, error) { switch in.Action { case "help": @@ -138,7 +175,10 @@ func (t *tool) handle(ctx context.Context, _ *mcp.CallToolRequest, in input) (*m } text, ok := help[p.Topic] if !ok { - return errorResult("there is no action " + p.Topic + "; use listTables, describeTable, query or help"), nil, nil + return errorResult("there is no action " + p.Topic + "; use " + t.actions()), nil, nil + } + if p.Topic == "" && t.caller.AgentMarksPII { + text = strings.Replace(text, "\n help ", markPIIUsage+"\n help ", 1) } return textResult(text), nil, nil case "listTables": @@ -155,9 +195,22 @@ func (t *tool) handle(ctx context.Context, _ *mcp.CallToolRequest, in input) (*m return errorResult(err.Error()), nil, nil } return t.query(ctx, p) + case "markPii": + var p markPayload + if err := decode(in.Payload, &p); err != nil { + return errorResult(err.Error()), nil, nil + } + return t.markPII(ctx, p) default: - return errorResult(fmt.Sprintf("unknown action %q: use listTables, describeTable, query or help", in.Action)), nil, nil + return errorResult(fmt.Sprintf("unknown action %q: use %s", in.Action, t.actions())), nil, nil + } +} + +func (t *tool) actions() string { + if t.caller.AgentMarksPII { + return "listTables, describeTable, query, markPii or help" } + return "listTables, describeTable, query or help" } func (t *tool) schema(ctx context.Context) ([]entity.SchemaTable, error) { @@ -238,6 +291,23 @@ func (t *tool) query(ctx context.Context, p queryPayload) (*mcp.CallToolResult, return toonResult(out) } +func (t *tool) markPII(ctx context.Context, p markPayload) (*mcp.CallToolResult, any, error) { + res, err := t.backend.AgentMarkPII(ctx, t.caller.UserID, t.caller.DatabaseID, p.Schema, p.Table, p.Columns) + if err != nil { + return errorResult(err.Error()), nil, nil + } + cols := toon.Table{Columns: []string{"name", "pii"}} + for _, c := range res.Marked { + cols.Rows = append(cols.Rows, []any{c, "marked"}) + } + for _, c := range res.Already { + cols.Rows = append(cols.Rows, []any{c, "already"}) + } + return toonResult(toon.Object{ + {Key: "table", Value: res.Table}, {Key: "marked", Value: len(res.Marked)}, {Key: "columns", Value: cols}, + }) +} + func decode(payload map[string]any, out any) error { raw, err := json.Marshal(payload) if err != nil { diff --git a/backend/internal/controller/mcpserver/server_test.go b/backend/internal/controller/mcpserver/server_test.go index 3d1c5cb..62a5f2d 100644 --- a/backend/internal/controller/mcpserver/server_test.go +++ b/backend/internal/controller/mcpserver/server_test.go @@ -24,6 +24,9 @@ type backend struct { result *entity.QueryResult queryErr error last crud.QueryRequest + marked *entity.PIIMarked + markErr error + markArgs []any } func (b *backend) Schema(context.Context, int64, int64) ([]entity.SchemaTable, error) { @@ -43,6 +46,11 @@ func (b *backend) RunQuery(_ context.Context, rq crud.QueryRequest) (*entity.Que return b.result, b.queryErr } +func (b *backend) AgentMarkPII(_ context.Context, userID, id int64, schema, table string, columns []string) (*entity.PIIMarked, error) { + b.markArgs = []any{userID, id, schema, table, columns} + return b.marked, b.markErr +} + var caller = mcpserver.Caller{UserID: 7, DatabaseID: 9, DatabaseName: "shop", Driver: "postgres"} func session(t *testing.T, b *backend) *mcp.ClientSession { @@ -196,6 +204,57 @@ rows[1]{id,name}: 1,null`, text) } +func TestMarkPIIIsOfferedOnlyWhenAllowed(t *testing.T) { + cs := session(t, &backend{}) + assert.NotContains(t, cs.InitializeResult().Instructions, "markPii") + tools, err := cs.ListTools(context.Background(), nil) + require.NoError(t, err) + assert.Equal(t, mcpserver.ToolDescription, tools.Tools[0].Description) + text, _ := call(t, cs, "help", nil) + assert.NotContains(t, text, "markPii") + + on := caller + on.AgentMarksPII = true + cs = sessionAs(t, &backend{}, on) + assert.Contains(t, cs.InitializeResult().Instructions, "mark it with markPii. Only the owner can unmark a column.") + tools, err = cs.ListTools(context.Background(), nil) + require.NoError(t, err) + assert.Equal(t, mcpserver.MarkPIIToolDescription, tools.Tools[0].Description) + assert.False(t, tools.Tools[0].Annotations.ReadOnlyHint, "markPii changes the PII set") + text, _ = call(t, cs, "help", nil) + assert.Contains(t, text, " markPii {") + assert.Less(t, strings.Index(text, "markPii"), strings.Index(text, " help ")) + text, _ = call(t, cs, "help", map[string]any{"topic": "markPii"}) + assert.Contains(t, text, "only the owner can unmark one") + text, _ = call(t, cs, "drop", nil) + assert.Equal(t, `error: "unknown action \"drop\": use listTables, describeTable, query, markPii or help"`, text) +} + +func TestMarkPII(t *testing.T) { + on := caller + on.AgentMarksPII = true + b := &backend{marked: &entity.PIIMarked{Table: "public.users", Marked: []string{"email", "phone"}, Already: []string{"name"}}} + cs := sessionAs(t, b, on) + text, isErr := call(t, cs, "markPii", map[string]any{"table": "users", "columns": []string{"email", "phone", "name"}}) + require.False(t, isErr) + assert.Equal(t, `table: public.users +marked: 2 +columns[3]{name,pii}: + email,marked + phone,marked + name,already`, text) + assert.Equal(t, []any{int64(7), int64(9), "", "users", []string{"email", "phone", "name"}}, b.markArgs) + + _, isErr = call(t, cs, "markPii", map[string]any{"table": "users", "columns": "email"}) + assert.True(t, isErr, "columns is a list") + + b.markErr = crud.ErrAgentMarksPIIOff + text, isErr = call(t, cs, "markPii", map[string]any{"table": "users", "schema": "public", "columns": []string{"email"}}) + assert.True(t, isErr) + assert.Contains(t, text, "the owner has not let agents mark PII columns") + assert.Equal(t, "public", b.markArgs[2]) +} + func TestHandlerNeedsACaller(t *testing.T) { h := mcpserver.NewHandler(&backend{}, "test") body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}` diff --git a/backend/internal/data/entity/entity.go b/backend/internal/data/entity/entity.go index ea2b22a..ca7f370 100644 --- a/backend/internal/data/entity/entity.go +++ b/backend/internal/data/entity/entity.go @@ -53,9 +53,11 @@ type Database struct { Driver string DSNMasked string PIICount int - MCPURL string - CreatedAt time.Time - UpdatedAt time.Time + // AgentMarksPII is whether the MCP agent may add PII columns. + AgentMarksPII bool + MCPURL string + CreatedAt time.Time + UpdatedAt time.Time } // PIIColumn is a column queries may not read. @@ -65,6 +67,14 @@ type PIIColumn struct { Column string } +// PIIMarked is what an agent's markPii did to one table: the columns it +// marked and those that were PII already. +type PIIMarked struct { + Table string + Marked []string + Already []string +} + // SchemaColumn is a column, and whether it is or looks like PII. type SchemaColumn struct { Name string diff --git a/backend/internal/data/model/model.go b/backend/internal/data/model/model.go index 2b28197..d7a06b7 100644 --- a/backend/internal/data/model/model.go +++ b/backend/internal/data/model/model.go @@ -25,8 +25,11 @@ type Database struct { DSNSealed string `gorm:"not null"` // DSNMasked is the DSN with its password replaced, safe to show. DSNMasked string - CreatedAt time.Time - UpdatedAt time.Time + // AgentMarksPII lets the database's MCP agent add PII columns. It can + // never remove one; only the owner can. + AgentMarksPII bool `gorm:"not null;default:false"` + CreatedAt time.Time + UpdatedAt time.Time } // PIIColumn is a column queries may not read. diff --git a/backend/internal/data/view/view.go b/backend/internal/data/view/view.go index 11ea9d8..e6dc27a 100644 --- a/backend/internal/data/view/view.go +++ b/backend/internal/data/view/view.go @@ -54,6 +54,7 @@ type Database struct { Driver string `json:"driver"` DSNMasked string `json:"dsnMasked"` PIIColumnCount int `json:"piiColumnCount"` + AgentMarksPII bool `json:"agentMarksPii"` MCPURL string `json:"mcpUrl"` CreatedAt time.Time `json:"createdAt"` UpdatedAt time.Time `json:"updatedAt"` @@ -92,6 +93,10 @@ type PIIColumnsRequest struct { PIIColumns []PIIColumn `json:"piiColumns"` } +type AgentPIIRequest struct { + AgentMarksPII bool `json:"agentMarksPii"` +} + type PIIColumnsResponse struct { PIIColumns []PIIColumn `json:"piiColumns"` } @@ -166,7 +171,7 @@ func FromDrivers(in []entity.Driver) DriversResponse { func FromDatabase(d *entity.Database) Database { return Database{ ID: d.ID, Name: d.Name, Driver: d.Driver, DSNMasked: d.DSNMasked, PIIColumnCount: d.PIICount, - MCPURL: d.MCPURL, CreatedAt: d.CreatedAt, UpdatedAt: d.UpdatedAt, + AgentMarksPII: d.AgentMarksPII, MCPURL: d.MCPURL, CreatedAt: d.CreatedAt, UpdatedAt: d.UpdatedAt, } } diff --git a/backend/internal/handler/api/api.go b/backend/internal/handler/api/api.go index 6b68484..74925d3 100644 --- a/backend/internal/handler/api/api.go +++ b/backend/internal/handler/api/api.go @@ -88,6 +88,7 @@ func Register(p Params) { r.Get("/databases/:id/schema", h.schema) r.Get("/databases/:id/pii", h.getPII) r.Put("/databases/:id/pii", h.setPII) + r.Put("/databases/:id/pii/agent", h.setAgentPII) r.Post("/databases/:id/query", h.query) r.Get("/databases/:id/queries", h.queries) @@ -381,6 +382,18 @@ func (h *handler) getPII(c fiber.Ctx) error { return c.JSON(view.FromPIIColumns(cols)) } +func (h *handler) setAgentPII(c fiber.Ctx) error { + var rq view.AgentPIIRequest + if err := c.Bind().JSON(&rq); err != nil { + return badPayload(c) + } + d, err := h.crud.SetAgentMarksPII(c.Context(), userID(c), dbID(c), rq.AgentMarksPII) + if err != nil { + return respondErr(c, err) + } + return c.JSON(view.DatabaseResponse{Database: view.FromDatabase(d)}) +} + func (h *handler) setPII(c fiber.Ctx) error { var rq view.PIIColumnsRequest if err := c.Bind().JSON(&rq); err != nil { diff --git a/backend/internal/handler/mcp/mcp.go b/backend/internal/handler/mcp/mcp.go index 34df064..054b754 100644 --- a/backend/internal/handler/mcp/mcp.go +++ b/backend/internal/handler/mcp/mcp.go @@ -356,6 +356,6 @@ func (h *handler) serveMCP(w http.ResponseWriter, r *http.Request) { h.challenge(w, idStr, "invalid or expired token") return } - ctx := mcpserver.WithCaller(r.Context(), mcpserver.Caller{UserID: owner.UserID, DatabaseID: id, DatabaseName: owner.Name, Driver: owner.Driver}) + ctx := mcpserver.WithCaller(r.Context(), mcpserver.Caller{UserID: owner.UserID, DatabaseID: id, DatabaseName: owner.Name, Driver: owner.Driver, AgentMarksPII: owner.AgentMarksPII}) h.mcp.ServeHTTP(w, r.WithContext(ctx)) } diff --git a/backend/internal/repository/repository.go b/backend/internal/repository/repository.go index e0ab417..5576a76 100644 --- a/backend/internal/repository/repository.go +++ b/backend/internal/repository/repository.go @@ -125,6 +125,20 @@ func (r *Repository) UpdateDatabase(ctx context.Context, d *model.Database) erro return nil } +// SetAgentMarksPII turns on or off whether a database's MCP agent may add +// PII columns. +func (r *Repository) SetAgentMarksPII(ctx context.Context, userID, id int64, on bool) error { + res := r.db.WithContext(ctx).Model(&model.Database{}).Where("id = ? AND user_id = ?", id, userID). + Updates(map[string]any{"agent_marks_pii": on, "updated_at": time.Now()}) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return ErrNotFound + } + return nil +} + // DeleteDatabase removes a database with its PII columns and query log. func (r *Repository) DeleteDatabase(ctx context.Context, userID, id int64) error { return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { @@ -182,10 +196,24 @@ func (r *Repository) ReplacePIIColumns(ctx context.Context, databaseID int64, co if len(columns) == 0 { return nil } - return tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&columns).Error + return tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(&columns, piiBatch).Error }) } +// piiBatch keeps one INSERT under SQLite's limit of 32766 bound variables +// (six per column), so the largest PII set still saves. +const piiBatch = 1000 + +// AddPIIColumns adds columns to the database's PII set, keeping every column +// already in it. +func (r *Repository) AddPIIColumns(ctx context.Context, databaseID int64, columns []model.PIIColumn) error { + for i := range columns { + columns[i].ID = r.newID() + columns[i].DatabaseID = databaseID + } + return r.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(&columns, piiBatch).Error +} + // GetTmphost returns a person's tmphost connection. func (r *Repository) GetTmphost(ctx context.Context, userID int64) (*model.TmphostConnection, error) { var c model.TmphostConnection diff --git a/backend/internal/testutil/testutil.go b/backend/internal/testutil/testutil.go index 615a9a3..b20b522 100644 --- a/backend/internal/testutil/testutil.go +++ b/backend/internal/testutil/testutil.go @@ -30,11 +30,13 @@ import ( const GoodDSN = "postgres://app:secret@db.internal:5432/shop" // Driver is a fake engine with a users table whose email column looks -// like PII. Query returns Rows rows, or QueryErr. +// like PII. Query returns Rows rows, or QueryErr. Extra tables are added to +// the schema. type Driver struct { Rows int QueryErr error SchemaErr error + Extra []dbdriver.Table } func (*Driver) ID() string { return "postgres" } @@ -58,7 +60,7 @@ func (d *Driver) Schema(context.Context, string) ([]dbdriver.Table, error) { if d.SchemaErr != nil { return nil, d.SchemaErr } - return []dbdriver.Table{ + return append([]dbdriver.Table{ {Schema: "public", Name: "users", Kind: "table", Columns: []dbdriver.Column{ {Name: "id", DataType: "integer"}, {Name: "email", DataType: "text", SuggestedPII: true}, @@ -66,7 +68,7 @@ func (d *Driver) Schema(context.Context, string) ([]dbdriver.Table, error) { {Schema: "public", Name: "orders", Kind: "table", Columns: []dbdriver.Column{ {Name: "id", DataType: "integer"}, {Name: "total", DataType: "numeric"}, }}, - }, nil + }, d.Extra...), nil } func (d *Driver) Query(_ context.Context, _, _ string, maxRows int, _ time.Duration) (*dbdriver.Result, error) { if d.QueryErr != nil { diff --git a/docs/API.md b/docs/API.md index 6662728..bc570f2 100644 --- a/docs/API.md +++ b/docs/API.md @@ -39,6 +39,7 @@ A database is one connection a person registered; it gets its own MCP server. "driver": "postgres", "dsnMasked": "postgres://app:****@db.internal:5432/analytics?sslmode=require", "piiColumnCount": 3, + "agentMarksPii": false, "mcpUrl": "https://sqlmcp.example.com/mcp/", "createdAt": "RFC3339", "updatedAt": "RFC3339" @@ -57,6 +58,7 @@ A database is one connection a person registered; it gets its own MCP server. | GET | `/databases/:id/schema` | | `{ "tables": [{ "schema": "public", "name": "users", "columns": [{ "name": "email", "dataType": "text", "pii": true, "suggestedPii": true }] }] }` | | GET | `/databases/:id/pii` | | `{ "piiColumns": [{ "schema": "public", "table": "users", "column": "email" }] }` | | PUT | `/databases/:id/pii` | `{ "piiColumns": [PiiColumn] }` (replaces the whole set; an empty `schema` is the engine's default: `public`, or the DSN's database on MySQL) | `{ "piiColumns": [PiiColumn] }` | +| PUT | `/databases/:id/pii/agent` | `{ "agentMarksPii": true }` (lets the MCP agent add PII columns with `markPii`; it can never remove one) | `{ "database": Database }` | | POST | `/databases/:id/query` | `{ "sql": "select ..." }` | `{ "columns": ["id"], "rows": [[1]], "rowCount": 1, "truncated": false }`; a query touching a PII column is 422 with the reason | | GET | `/databases/:id/queries?limit=50` | | `{ "queries": [{ "id", "source": "mcp" \| "console", "sql", "status": "ok" \| "rejected" \| "error", "reason", "rowCount", "resultUrl", "durationMs", "createdAt" }] }` | | GET | `/databases/:id/queries/stream` | | `text/event-stream`: one `event: query` per query as it is recorded, `data` a query log entry as above and `id` its id; `: ping` comments every 25 s. 401 without a session, 404 for someone else's database | @@ -96,7 +98,9 @@ registration, RFC 8414/9728 discovery and a server-rendered consent page: - `/mcp/{id}/oauth2/register`, `/mcp/{id}/oauth2/authorize`, `/mcp/{id}/oauth2/token` One tool, `sql`, taking `{ action, payload }`: `listTables`, `describeTable`, -`query`, `help`. +`query`, `help`, and `markPii` when the database's `agentMarksPii` is on (the +tool is then no longer annotated read-only, and the instructions ask the agent +to mark the personal data it finds). Answers are [TOON](https://toonformat.dev) text (spec v3, comma-delimited), which spends fewer tokens than JSON: a table names its columns once, then @@ -119,8 +123,17 @@ query rowCount: 120 rows[25]{id,name}: 1,alice 2,null +markPii table: public.users ({table, schema?, columns: [...]}) + marked: 1 + columns[2]{name,pii}: + phone,marked + email,already ``` +`markPii` only adds: every column must exist in the table, a column already +marked is reported as `already`, and nothing an agent sends unmarks one. The +setting is checked on every call, so turning it off stops an agent at once. + A query result shows at most the inline limit (`QUERY_INLINE_ROWS`, default 25) of its rows. diff --git a/frontend/src/api.js b/frontend/src/api.js index 6868c5d..4126d58 100644 --- a/frontend/src/api.js +++ b/frontend/src/api.js @@ -173,6 +173,12 @@ export async function savePiiColumns(id, piiColumns) { return data?.piiColumns || [] } +/** Lets, or stops, the database's MCP agent marking PII columns; the updated database. */ +export async function setAgentMarksPii(id, agentMarksPii) { + const data = await request(`/databases/${enc(id)}/pii/agent`, { method: 'PUT', body: { agentMarksPii } }) + return data?.database +} + /** `{ columns, rows, rowCount, truncated }`; a query touching PII rejects with status 422. */ export function runQuery(id, sql) { return request(`/databases/${enc(id)}/query`, { method: 'POST', body: { sql } }) diff --git a/frontend/src/components/PiiEditor.vue b/frontend/src/components/PiiEditor.vue index b90de3c..139d730 100644 --- a/frontend/src/components/PiiEditor.vue +++ b/frontend/src/components/PiiEditor.vue @@ -1,13 +1,16 @@ diff --git a/frontend/test/PiiEditor.test.js b/frontend/test/PiiEditor.test.js index 5e63ea4..809a3af 100644 --- a/frontend/test/PiiEditor.test.js +++ b/frontend/test/PiiEditor.test.js @@ -19,10 +19,10 @@ const SCHEMA = { describe('PiiEditor', () => { afterEach(() => { vi.unstubAllGlobals(); useToasts().toasts.value = [] }) - async function mountEditor(extra = {}) { + async function mountEditor(extra = {}, props = {}) { setActivePinia(createPinia()) const f = mockFetch({ 'GET /databases/d1/schema': SCHEMA, 'GET /databases/d1/pii': { piiColumns: [] }, ...extra }) - const w = mount(PiiEditor, { props: { databaseId: 'd1' } }) + const w = mount(PiiEditor, { props: { databaseId: 'd1', ...props } }) await flush(); await flush() return { w, f } } @@ -61,6 +61,56 @@ describe('PiiEditor', () => { expect(w.find('[data-test="selected-count"]').text()).toBe('3 columns selected') }) + it('keeps columns the agent marked after the editor loaded', async () => { + let calls = 0 + const { w, f } = await mountEditor({ + // By the save, the agent has marked orders.id. + 'GET /databases/d1/pii': () => ({ piiColumns: [ + { schema: 'public', table: 'users', column: 'email' }, + ...(calls++ ? [{ schema: 'public', table: 'orders', column: 'id' }] : []), + ] }), + 'PUT /databases/d1/pii': (init) => ({ piiColumns: JSON.parse(init.body).piiColumns }), + }, { agentMarksPii: true }) + await w.find('[data-test="select-suggested"]').trigger('click') + await w.find('[data-test="save-pii"]').trigger('click') + await flush(); await flush() + const put = f.mock.calls.find((c) => c[1]?.method === 'PUT') + expect(JSON.parse(put[1].body).piiColumns).toEqual([ + { schema: 'public', table: 'orders', column: 'id' }, + { schema: 'public', table: 'users', column: 'email' }, + { schema: 'public', table: 'users', column: 'phone' }, + ]) + expect(w.emitted('saved')[0]).toEqual([3]) + }) + + it('turns the agent marking PII on and off', async () => { + const { w, f } = await mountEditor({ + 'PUT /databases/d1/pii/agent': (init) => ({ database: { id: 'd1', ...JSON.parse(init.body) } }), + }) + const sw = w.find('[data-test="agent-marks-pii"]') + expect(sw.attributes('aria-checked')).toBe('false') + await sw.trigger('click') + await flush(); await flush() + expect(JSON.parse(f.mock.calls.at(-1)[1].body)).toEqual({ agentMarksPii: true }) + expect(w.emitted('agent-changed')[0]).toEqual([{ id: 'd1', agentMarksPii: true }]) + expect(useToasts().toasts.value.at(-1).message).toContain('can now mark') + + await w.setProps({ agentMarksPii: true }) + expect(sw.attributes('aria-checked')).toBe('true') + await sw.trigger('click') + await flush(); await flush() + expect(JSON.parse(f.mock.calls.at(-1)[1].body)).toEqual({ agentMarksPii: false }) + expect(useToasts().toasts.value.at(-1).message).toContain('no longer') + }) + + it('reports a failed switch', async () => { + const { w } = await mountEditor({ 'PUT /databases/d1/pii/agent': [500, { error: { message: 'boom' } }] }) + await w.find('[data-test="agent-marks-pii"]').trigger('click') + await flush(); await flush() + expect(w.emitted('agent-changed')).toBeUndefined() + expect(useToasts().toasts.value.at(-1).message).toBe('boom') + }) + it('shows a schema load error', async () => { const { w } = await mountEditor({ 'GET /databases/d1/schema': [500, { error: { message: 'permission denied' } }] }) expect(w.text()).toContain('permission denied') diff --git a/frontend/test/api.test.js b/frontend/test/api.test.js index c761525..155e700 100644 --- a/frontend/test/api.test.js +++ b/frontend/test/api.test.js @@ -90,6 +90,7 @@ describe('api', () => { 'GET /databases/x1/schema': { tables: [{ name: 't' }] }, 'GET /databases/x1/pii': { piiColumns: [{ schema: 'public', table: 't', column: 'c' }] }, 'PUT /databases/x1/pii': (init) => ({ piiColumns: JSON.parse(init.body).piiColumns }), + 'PUT /databases/x1/pii/agent': (init) => ({ database: { ...db, ...JSON.parse(init.body) } }), 'POST /databases/x1/query': { columns: ['id'], rows: [[1]], rowCount: 1, truncated: false }, 'GET /databases/x1/queries?limit=10': { queries: [{ id: 'q' }] }, }) @@ -111,6 +112,8 @@ describe('api', () => { expect((await api.runQuery('x1', 'select 1')).rowCount).toBe(1) expect(bodyOf(f, 10)).toEqual({ sql: 'select 1' }) expect(await api.fetchQueries('x1', 10)).toEqual([{ id: 'q' }]) + expect(await api.setAgentMarksPii('x1', true)).toEqual({ ...db, agentMarksPii: true }) + expect(bodyOf(f, 12)).toEqual({ agentMarksPii: true }) }) it('defaults missing lists to empty arrays', async () => { diff --git a/frontend/test/usePiiSelection.test.js b/frontend/test/usePiiSelection.test.js index 9353e00..32f0534 100644 --- a/frontend/test/usePiiSelection.test.js +++ b/frontend/test/usePiiSelection.test.js @@ -168,4 +168,19 @@ describe('usePiiSelection', () => { expect(s.dirty.value).toBe(false) expect(s.selectedCount.value).toBe(3) }) + + it('applies its changes onto a set that changed on the server', () => { + const s = setup() + s.setColumn(users, 'email', false) + s.setColumn(users, 'phone', true) + const latest = [ + { schema: 'public', table: 'users', column: 'email' }, + { schema: 'sales', table: 'orders', column: 'shipping_address' }, // the agent marked it meanwhile + ] + expect(s.toPayloadOnto(latest)).toEqual([ + { schema: 'public', table: 'users', column: 'phone' }, + { schema: 'sales', table: 'orders', column: 'shipping_address' }, + ]) + expect(s.toPayloadOnto()).toEqual([{ schema: 'public', table: 'users', column: 'phone' }]) + }) }) From ce3890bf8d1253c4fe3262a8f8cf3304a7c13f69 Mon Sep 17 00:00:00 2001 From: tmphost-dev Date: Sat, 24 Oct 2026 01:25:06 +0000 Subject: [PATCH 2/2] Cover a failed agentMarksPii update Co-Authored-By: Claude Opus 5.5 --- backend/internal/controller/crud/crud_test.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/backend/internal/controller/crud/crud_test.go b/backend/internal/controller/crud/crud_test.go index 1da7ee0..db1659c 100644 --- a/backend/internal/controller/crud/crud_test.go +++ b/backend/internal/controller/crud/crud_test.go @@ -277,7 +277,7 @@ func TestAgentMarksPII(t *testing.T) { assert.ErrorIs(t, err, crud.ErrAgentMarksPIIOff) } -func TestAgentMarkPIIWhenTheStoreFails(t *testing.T) { +func TestAgentPIIWhenTheStoreFails(t *testing.T) { env := testutil.New(t) alice := signIn(t, env, "alice@example.com") id := newDB(t, env, alice) @@ -293,6 +293,10 @@ func TestAgentMarkPIIWhenTheStoreFails(t *testing.T) { require.NoError(t, db.Exec("DROP TABLE pii_columns").Error) _, err = env.Ctrl.AgentMarkPII(ctx, alice, id, "", "users", []string{"email"}) assert.Error(t, err) + + require.NoError(t, db.Exec("DROP TABLE databases").Error) + _, err = env.Ctrl.SetAgentMarksPII(ctx, alice, id, false) + assert.ErrorContains(t, err, "no such table") } func TestTooManyAgentPIIColumns(t *testing.T) {