query-guard

1// The hooks Claude Code loads for Query Guard.
2// The Bash hook holds a flagged command until you press Proceed or Cancel.
3// The pane, the band above the prompt, and the status line show what it holds.
4
5import { atom, read, update } from 'claude-code'
6import type { EngineInterface, Register } from 'claude-code'
7
8import type { HeldCall, RowCount } from '../types'
9import { analyze, refusal } from './rules'
10import type { Finding } from './rules'
11
12const PANE = 'query-guard'
13const POLL = ['sleep', '0.25']
14
15const pending = atom({ plugin: 'query-guard', key: 'pending' } as const, [])
16const held = atom({ plugin: 'query-guard', key: 'held' } as const, 0)
17const where = atom({ plugin: 'query-guard', key: 'where' } as const, 'pane')
18
19type Decision = 'proceed' | 'cancel'
20
21// The person's answers, by tool_use_id. A module variable on purpose: the
22// waiting hook reads it between `$` calls, and every `$.state.get` of one
23// dispatch reads the moment the dispatch began.
24const decisions = new Map<string, Decision>()
25
26const statusText = (count: number) => `🛡 Query Guard · ${count} held this session`
27
28// A table name SQLite can be asked about: plain identifiers and one dot.
29const COUNTABLE = /^[A-Za-z_][\w$]*(?:\.[A-Za-z_][\w$]*)?$/
30
31async function rowCounts($: EngineInterface, finding: Finding): Promise<RowCount[]> {
32  const file = finding.sqliteFile
33  if (file === undefined) return []
34  const tables = [...new Set(finding.hits.filter(h => h.isTable).map(h => h.target))]
35  if (tables.length === 0) return []
36
37  const exists = await $.process.run(['test', '-f', file]).catch(() => undefined)
38  if (exists?.exitCode !== 0) return []
39
40  const counts: RowCount[] = []
41  for (const table of tables) {
42    if (!COUNTABLE.test(table)) continue
43    const quoted = table
44      .split('.')
45      .map(part => `"${part}"`)
46      .join('.')
47    const ran = await $.process
48      .run(['sqlite3', '-readonly', '-batch', '-noheader', file, `SELECT COUNT(*) FROM ${quoted};`], {
49        timeoutMs: 5_000,
50      })
51      .catch(() => undefined)
52    const rows = Number(ran?.stdout.trim())
53    if (ran?.exitCode === 0 && Number.isInteger(rows)) counts.push({ table, rows })
54  }
55  return counts
56}
57
58async function decide($: EngineInterface, id: string, decision: Decision) {
59  decisions.set(id, decision)
60  await update($, pending, list => list.filter(call => call.id !== id))
61}
62
63export const register: Register = on => {
64  on('session.start', async ($, e, next) => {
65    $.ui.status(statusText(await read($, held)))
66    return next(e)
67  })
68
69  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
70    const finding = analyze(e.command)
71    if (finding === null) return next(e)
72
73    const call: HeldCall = {
74      id: e.tool_use_id,
75      command: e.command,
76      client: finding.client,
77      reasons: finding.hits.map(h => h.reason),
78      effects: finding.hits.map(h => h.effect),
79      targets: [...new Set(finding.hits.map(h => h.target))],
80      database: finding.sqliteFile,
81      rowCounts: await rowCounts($, finding),
82    }
83
84    decisions.delete(call.id)
85    await update($, pending, list => [...list, call])
86    const count = await update($, held, n => n + 1)
87    $.ui.status(statusText(count))
88
89    const opened = await $.ui
90      .open({ id: PANE, title: 'Query Guard', focus: true })
91      .catch(() => ({ isPlaced: false }))
92    await update($, where, () => (opened.isPlaced ? 'pane' : 'band'))
93    if (!opened.isPlaced) {
94      // Too narrow for a pane: the report is drawn above the prompt instead;
95      // close the waiting pane so it does not also appear on a resize.
96      await $.ui.close({ id: PANE }).catch(() => undefined)
97    }
98
99    let decision = decisions.get(call.id)
100    while (decision === undefined) {
101      if (next.signal.aborted) {
102        await decide($, call.id, 'cancel')
103        break
104      }
105      await $.process.run(POLL).catch(() => undefined)
106      decision = decisions.get(call.id)
107    }
108    decisions.delete(call.id)
109
110    if ((await read($, pending)).length === 0) {
111      await $.ui.close({ id: PANE }).catch(() => undefined)
112    }
113
114    if (decision === 'proceed') return next(e)
115    return { deny: refusal(call.command, call.effects, call.rowCounts) }
116  }).catch(($, e, next) =>
117    next.called
118      ? next(e)
119      : { deny: 'Query Guard could not finish checking this database command, so it was not run. Ask the user to run it themselves or retry.' },
120  )
121
122  on('ui.close', { id: PANE }, async ($, e, next) => {
123    // Closing the pane by hand while a call is held cancels that call.
124    if (e.origin.kind === 'person') {
125      for (const call of await read($, pending)) await decide($, call.id, 'cancel')
126    }
127    return next(e)
128  })
129
130  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
131    const list = await read($, pending)
132    return report($, e, list, e.props.bodyColumns)
133  })
134
135  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
136    const list = await read($, pending)
137    if (list.length === 0 || (await read($, where)) !== 'band') return next(e)
138    return report($, e, list, e.props.bodyColumns)
139  })
140}
141
142function report(
143  $: EngineInterface,
144  e: Parameters<EngineInterface['ui']['resolve']>[0],
145  list: readonly HeldCall[],
146  columns: number,
147) {
148  const { Box, Text, Button } = $.ui.resolve(e)
149  const call = list[0]
150  if (call === undefined) {
151    return (
152      <Box>
153        <Text dimColor>Nothing held. Query Guard is watching SQL and dbt commands.</Text>
154      </Box>
155    )
156  }
157  const counts = new Map(call.rowCounts.map(c => [c.table, c.rows]))
158  const width = Math.max(20, columns)
159
160  return (
161    <Box flexDirection="column" width={width}>
162      <Text bold color="red">
163        🛡 Query Guard held a {call.client} command
164      </Text>
165      <Box flexDirection="column" marginTop={1}>
166        <Text dimColor>Command</Text>
167        <Text wrap="wrap">{call.command}</Text>
168      </Box>
169      <Box flexDirection="column" marginTop={1}>
170        <Text dimColor>Why it was flagged</Text>
171        {call.reasons.map(reason => (
172          <Text wrap="wrap">• {reason}</Text>
173        ))}
174      </Box>
175      <Box flexDirection="column" marginTop={1}>
176        <Text dimColor>Target{call.targets.length > 1 ? 's' : ''}</Text>
177        {call.targets.map(target => (
178          <Text wrap="wrap">
179            {target}
180            {counts.has(target) ? (
181              <Text color="yellow"> · {counts.get(target)?.toLocaleString('en-US')} rows now</Text>
182            ) : (
183              ''
184            )}
185          </Text>
186        ))}
187        {call.database !== undefined && <Text dimColor>in {call.database}</Text>}
188      </Box>
189      <Box marginTop={1} gap={2}>
190        <Button key="proceed" hotkey="1" label="Proceed" onPress={() => decide($, call.id, 'proceed')} />
191        <Button key="cancel" hotkey="2" label="Cancel" variant="primary" onPress={() => decide($, call.id, 'cancel')} />
192      </Box>
193      {list.length > 1 && <Text dimColor>{list.length - 1} more held after this one</Text>}
194    </Box>
195  )
196}
197

1// The detection rules. Each rule reads the text of one Bash command.
2// Nothing here calls Claude Code, so the tests can check each rule directly.
3
4export type RuleId =
5  | 'drop-table'
6  | 'drop-database'
7  | 'drop-schema'
8  | 'truncate'
9  | 'delete-without-where'
10  | 'update-without-where'
11  | 'drop-column'
12  | 'dbt-full-refresh'
13  | 'production-target'
14
15export type Hit = {
16  rule: RuleId
17  /** Why it was held, for the person. */
18  reason: string
19  /** What it would destroy, for the refusal Claude reads. */
20  effect: string
21  /** The table, schema, database or dbt selection. */
22  target: string
23  /** True when the target is a table whose rows can be counted. */
24  isTable: boolean
25}
26
27export type Finding = {
28  client: string
29  hits: Hit[]
30  /** The SQLite database file a sqlite3 command names, when it names one. */
31  sqliteFile?: string
32}
33
34const CLIENTS = ['sqlite3', 'psql', 'mysql', 'snowsql', 'databricks', 'dbt'] as const
35type Client = (typeof CLIENTS)[number]
36
37// A client invoked as a command word: at the start, after a shell operator,
38// after `env VAR=x`, or as a path ending in the name (/usr/bin/psql).
39const CLIENT_RE = new RegExp(
40  String.raw`(?:^|[\s;&|(\x60$]|/)(${CLIENTS.join('|')})(?=$|[\s;&|)])`,
41  'g',
42)
43
44// An identifier, optionally schema-qualified and quoted in any dialect.
45const NAME = String.raw`(?:[\w$]+|"[^"]+"|\x60[^\x60]+\x60|\[[^\]]+\])(?:\.(?:[\w$]+|"[^"]+"|\x60[^\x60]+\x60|\[[^\]]+\]))*`
46
47/** Which database clients the command runs, in order of appearance. */
48export function findClients(command: string): Client[] {
49  const found: Client[] = []
50  for (const match of command.matchAll(CLIENT_RE)) {
51    const client = (match[1] ?? '') as Client
52    if (!found.includes(client)) found.push(client)
53  }
54  return found
55}
56
57/** Strips one layer of quoting from an identifier: "a"."b" -> a.b */
58export function cleanName(name: string): string {
59  return name.replace(/["\x60[\]]/g, '')
60}
61
62// Where one SQL statement ends inside a shell command: a semicolon, or a
63// closing quote followed by a shell operator or the end of the command.
64function statementEnd(text: string, from: number): number {
65  const rest = text.slice(from)
66  const match = rest.match(/;|['"]\s*(?:$|\|\||&&|[|;&)>])/)
67  return match?.index === undefined ? text.length : from + match.index
68}
69
70function sqlHits(command: string): Hit[] {
71  const hits: Hit[] = []
72  // `--` and `/* */` comments could hide a WHERE or fake one; drop them.
73  const sql = command.replace(/\/\*[\s\S]*?\*\//g, ' ').replace(/--[^\n]*/g, ' ')
74
75  for (const m of sql.matchAll(
76    new RegExp(String.raw`\bDROP\s+(TABLE|DATABASE|SCHEMA)\s+(?:IF\s+EXISTS\s+)?(${NAME})`, 'gi'),
77  )) {
78    const kind = (m[1] ?? '').toUpperCase()
79    const target = cleanName(m[2] ?? '')
80    if (kind === 'TABLE') {
81      hits.push({
82        rule: 'drop-table',
83        reason: `DROP TABLE ${target}`,
84        effect: `dropped table ${target} with all of its rows`,
85        target,
86        isTable: true,
87      })
88    } else if (kind === 'DATABASE') {
89      hits.push({
90        rule: 'drop-database',
91        reason: `DROP DATABASE ${target}`,
92        effect: `dropped database ${target} and every table in it`,
93        target,
94        isTable: false,
95      })
96    } else {
97      hits.push({
98        rule: 'drop-schema',
99        reason: `DROP SCHEMA ${target}`,
100        effect: `dropped schema ${target} and every object in it`,
101        target,
102        isTable: false,
103      })
104    }
105  }
106
107  for (const m of sql.matchAll(
108    new RegExp(String.raw`\bTRUNCATE\s+(?:TABLE\s+)?(?:ONLY\s+)?(${NAME})`, 'gi'),
109  )) {
110    const target = cleanName(m[1] ?? '')
111    hits.push({
112      rule: 'truncate',
113      reason: `TRUNCATE ${target}`,
114      effect: `removed every row from ${target}`,
115      target,
116      isTable: true,
117    })
118  }
119
120  for (const m of sql.matchAll(new RegExp(String.raw`\bDELETE\s+FROM\s+(?:ONLY\s+)?(${NAME})`, 'gi'))) {
121    const after = m.index + m[0].length
122    const statement = sql.slice(after, statementEnd(sql, after))
123    if (/\bWHERE\b/i.test(statement)) continue
124    const target = cleanName(m[1] ?? '')
125    hits.push({
126      rule: 'delete-without-where',
127      reason: `DELETE FROM ${target} has no WHERE clause`,
128      effect: `deleted every row in ${target}`,
129      target,
130      isTable: true,
131    })
132  }
133
134  for (const m of sql.matchAll(new RegExp(String.raw`\bUPDATE\s+(?:ONLY\s+)?(${NAME})\s+SET\b`, 'gi'))) {
135    const after = m.index + m[0].length
136    const statement = sql.slice(after, statementEnd(sql, after))
137    if (/\bWHERE\b/i.test(statement)) continue
138    const target = cleanName(m[1] ?? '')
139    hits.push({
140      rule: 'update-without-where',
141      reason: `UPDATE ${target} has no WHERE clause`,
142      effect: `overwritten the updated columns in every row of ${target}`,
143      target,
144      isTable: true,
145    })
146  }
147
148  for (const m of sql.matchAll(
149    new RegExp(
150      String.raw`\bALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?(${NAME})\s+[^;]*?\bDROP\s+COLUMN\s+(?:IF\s+EXISTS\s+)?(${NAME})`,
151      'gi',
152    ),
153  )) {
154    const table = cleanName(m[1] ?? '')
155    const column = cleanName(m[2] ?? '')
156    hits.push({
157      rule: 'drop-column',
158      reason: `ALTER TABLE ${table} DROP COLUMN ${column}`,
159      effect: `dropped column ${column} from ${table}, with its value in every row`,
160      target: table,
161      isTable: true,
162    })
163  }
164
165  return hits
166}
167
168function dbtHits(command: string): Hit[] {
169  const isRebuild = /\bdbt\s+(?:[^\n;&|]*?\s)?(run|build|seed)\b[^\n;&|]*?--full-refresh\b/.test(command)
170  if (!isRebuild) return []
171  const selection = command.match(/(?:--select|--models|-s|-m)[\s=]+((?:[^\s-][^\s]*\s*)+)/)
172  const target = selection ? (selection[1] ?? '').trim() : 'every model in the project'
173  return [
174    {
175      rule: 'dbt-full-refresh',
176      reason: 'dbt --full-refresh drops and rebuilds incremental models',
177      effect: `dropped and rebuilt ${target} from scratch, discarding incremental history`,
178      target,
179      isTable: false,
180    },
181  ]
182}
183
184/** True when a host, URL or connection name says prod or production as a word. */
185export function isProdName(value: string): boolean {
186  return /(?:^|[^a-z])prod(?:uction)?(?:[^a-z]|$)/i.test(value)
187}
188
189/**
190 * The parts of a command that name where it connects: URLs, host and
191 * account flags, `host=` style connection strings, dbt targets and profiles,
192 * and environment variables such as PGHOST or DATABASE_URL.
193 */
194export function connectionValues(command: string, clients: readonly string[]): string[] {
195  const values: string[] = []
196  const add = (value: string | undefined) => {
197    if (value) values.push(value.replace(/^['"]|['"]$/g, ''))
198  }
199
200  for (const m of command.matchAll(/\b[a-z][a-z0-9+.-]*:\/\/[^\s'"]+/gi)) add(m[0])
201  for (const m of command.matchAll(/\bjdbc:[^\s'"]+/gi)) add(m[0])
202
203  const flags = ['--host', '--hostname', '--server-hostname', '--account', '--accountname', '--target', '--profile', '--connection', '--dbname', '--database', '--url']
204  const shortFlags = ['-h', '-a', '-t', '-d', '-D']
205  if (clients.includes('snowsql')) shortFlags.push('-c')
206  if (clients.includes('databricks')) shortFlags.push('-p')
207  for (const m of command.matchAll(/(?:^|\s)(--?[A-Za-z-]+)(?:=|\s+)("[^"]*"|'[^']*'|[^\s]+)/g)) {
208    const flag = m[1] ?? ''
209    if (flags.includes(flag) || shortFlags.includes(flag)) add(m[2])
210  }
211  // -hprod-db, as mysql takes it.
212  for (const m of command.matchAll(/(?:^|\s)-h([^\s=-][^\s]*)/g)) add(m[1])
213
214  for (const m of command.matchAll(
215    /\b(host|hostaddr|dbname|service|server|account|[A-Z_]*(?:HOST|URL|URI|ACCOUNT|SERVER|DSN|TARGET|PROFILE|CONNECTION)[A-Z_]*)=("[^"]*"|'[^']*'|[^\s'";]+)/g,
216  )) {
217    add(m[2])
218  }
219
220  return values
221}
222
223/** The database file a sqlite3 command opens, when it names one on disk. */
224export function sqliteFile(command: string): string | undefined {
225  const at = command.search(/(?:^|[\s;&|(/])sqlite3(?=\s)/)
226  if (at < 0) return undefined
227  const rest = command.slice(at).replace(/^.*?sqlite3/, '')
228  const tokens = [...rest.matchAll(/"([^"]*)"|'([^']*)'|(\S+)/g)].map(m => ({
229    text: m[1] ?? m[2] ?? m[3] ?? '',
230    isQuoted: m[3] === undefined,
231  }))
232  const takesValue = ['-cmd', '-init', '-separator', '-newline', '-nullvalue', '-vfs', '-maxsize', '-mmap', '-pagecache', '-lookaside', '-heap', '-zip']
233  for (let i = 0; i < tokens.length; i += 1) {
234    const { text, isQuoted } = tokens[i] ?? { text: '', isQuoted: false }
235    if (!isQuoted && /^[;&|<>]/.test(text)) return undefined
236    if (!isQuoted && text.startsWith('-')) {
237      if (takesValue.includes(text.replace(/^--/, '-'))) i += 1
238      continue
239    }
240    if (text === ':memory:' || text === '') return undefined
241    return text
242  }
243  return undefined
244}
245
246/** Reads a Bash command and answers what Query Guard holds it for, or null. */
247export function analyze(command: string): Finding | null {
248  const clients = findClients(command)
249  if (clients.length === 0) return null
250
251  const hits: Hit[] = [...sqlHits(command)]
252  if (clients.includes('dbt')) hits.push(...dbtHits(command))
253
254  const prod = connectionValues(command, clients).filter(isProdName)
255  if (prod.length > 0) {
256    const where = [...new Set(prod)].join(', ')
257    hits.push({
258      rule: 'production-target',
259      reason: `connects to a production target: ${where}`,
260      effect: `run against production (${where})`,
261      target: where,
262      isTable: false,
263    })
264  }
265
266  if (hits.length === 0) return null
267  const client = clients.find(c => c !== 'dbt' || hits.some(h => h.rule === 'dbt-full-refresh')) ?? clients[0] ?? 'sql'
268  return {
269    client,
270    hits,
271    sqliteFile: clients.includes('sqlite3') ? sqliteFile(command) : undefined,
272  }
273}
274
275/** The text Claude reads when the person cancels a held call. */
276export function refusal(command: string, effects: readonly string[], counts: readonly { table: string; rows: number }[]): string {
277  const lines = effects.map(effect => {
278    const count = counts.find(c => effect.includes(` ${c.table}`))
279    return count ? `- ${effect} (${count.rows.toLocaleString('en-US')} rows right now)` : `- ${effect}`
280  })
281  return [
282    'Query Guard: the user cancelled this command. Nothing was run.',
283    `Command: ${command}`,
284    'Had it run, it would have:',
285    ...lines,
286    'Do not retry it or work around the guard with a script or alias; ask the user how they want to proceed.',
287  ].join('\n')
288}
289

1// The state Query Guard keeps for one session.
2// The pane and the band read it to draw the report.
3
4/** One table a held command touches, with its row count when it could be read. */
5export type RowCount = { table: string; rows: number }
6
7/** A Bash call Query Guard is holding until the person answers it. */
8export type HeldCall = {
9  /** The tool_use_id of the held Bash call. */
10  id: string
11  command: string
12  /** The database client or tool the command runs: sqlite3, psql, dbt, ... */
13  client: string
14  /** Why it was held, one line per rule that matched. */
15  reasons: string[]
16  /** What it would destroy or touch, one line per effect. */
17  effects: string[]
18  /** The tables, schemas, databases or dbt selections it targets. */
19  targets: string[]
20  /** The SQLite database file, for sqlite3 commands that name one. */
21  database?: string
22  /** Current row counts, read with SELECT COUNT(*) from local SQLite files only. */
23  rowCounts: RowCount[]
24}
25
26declare module 'claude-code' {
27  interface PluginState {
28    'query-guard': {
29      /** Held calls, oldest first; the first one is the one on screen. */
30      pending: HeldCall[]
31      /** How many calls were held this session. */
32      held: number
33      /** Where the report is drawn: the pane, or above the prompt when no pane fits. */
34      where: 'pane' | 'band'
35    }
36  }
37}
38