Let the owner delegate marking PII columns to the agent

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.

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 (`<APP_BASE_URL>/mcp/<id>`) and
    add it to your client, for example
    `claude mcp add --transport http shop https://sql.example.com/mcp/<id>`.
@@ -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

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)

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
+}

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,121 @@ 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 TestAgentPIIWhenTheStoreFails(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)
+
+	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) {
+	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")

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.