og-sf-soql

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { QueryResult, View } from '../types'
5import { GAP, cell, fitWidths, parseArgs, queriesIn, resultFromCsv, resultFromJson, toCsv, toMarkdown } from './soql'
6import type { QueryCall } from './soql'
7
8const COMMAND = 'og-soql'
9const PANE = 'og-sf-soql'
10// Rows drawn in the pane; the copies always carry them all.
11const MAX_ROWS = 500
12
13const view = atom({ plugin: 'og-sf-soql', key: 'view' } as const, null)
14
15// The transcript's one-line summary for each query call Claude made, by tool_use_id.
16const summaries = new Map<string, string>()
17
18const SHELLS = new Set(['Bash', 'PowerShell'])
19
20const plural = (n: number, word: string) => `${n} ${word}${n === 1 ? '' : 's'}`
21
22const summaryOf = (result: QueryResult) =>
23  `📋 ${plural(result.totalSize, 'record')}${result.object ? ` · ${result.object}` : ''} · /${COMMAND} to see`
24
25/** Runs the sf CLI; on Windows through `cmd /c`, as `sf` is a batch file there. */
26async function runSf($: EngineInterface, args: string[]) {
27  const isWindows = (await $.env.get('OS')) === 'Windows_NT'
28  return $.process.run(isWindows ? ['cmd.exe', '/d', '/s', '/c', 'sf', ...args] : ['sf', ...args], { timeoutMs: 120_000 })
29}
30
31const queryArgs = (source: { query?: string; file?: string; org?: string; isTooling: boolean }) => [
32  'data',
33  'query',
34  ...(source.query !== undefined ? ['--query', source.query] : ['--file', source.file ?? '']),
35  ...(source.org !== undefined ? ['--target-org', source.org] : []),
36  ...(source.isTooling ? ['--use-tooling-api'] : []),
37  '--json',
38]
39
40/** sf's own error message from its JSON, or the first line of what it printed. */
41const errorOf = (stdout: string, stderr: string) => {
42  try {
43    const message = (JSON.parse(stdout) as { message?: unknown }).message
44    if (typeof message === 'string') return message
45  } catch {
46    // Not JSON: what it printed.
47  }
48  return (stderr || stdout).trim().split('\n')[0] || 'the query failed'
49}
50
51/** The rows of a query Claude ran: from its own output, or, for the human table, a --json run of it. */
52async function capture($: EngineInterface, call: QueryCall, stdout: string): Promise<QueryResult | undefined> {
53  const text = call.query ?? (call.file !== undefined ? await $.fs.read(call.file).then(String, () => call.file) : undefined)
54  const source = { query: (text ?? '').trim(), org: call.org, isTooling: call.isTooling }
55  if (call.format === 'json') return resultFromJson(stdout, source)
56  if (call.format === 'csv') return resultFromCsv(stdout, source)
57  const ran = await runSf($, queryArgs(call)).catch(() => undefined)
58
59  return ran === undefined ? undefined : resultFromJson(ran.stdout, source)
60}
61
62async function show($: EngineInterface, result: QueryResult) {
63  await update($, view, () => ({ result, record: null }))
64}
65
66async function openPane($: EngineInterface) {
67  const { isPlaced } = await $.ui.open({ id: PANE, title: 'SOQL', focus: true, closeOnEscape: true })
68  return isPlaced
69}
70
71/** Puts text on the system clipboard: pbcopy, Set-Clipboard on Windows, wl-copy or xclip on Linux. */
72async function toClipboard($: EngineInterface, text: string) {
73  const isWindows = (await $.env.get('OS')) === 'Windows_NT'
74  const commands = isWindows
75    ? [['powershell', '-NoProfile', '-Command', '[Console]::In.ReadToEnd() | Set-Clipboard']]
76    : [['pbcopy'], ['wl-copy'], ['xclip', '-selection', 'clipboard']]
77  for (const argv of commands) {
78    const ran = await $.process.run(argv, { stdin: text }).catch(() => undefined)
79    if (ran?.exitCode === 0) return
80  }
81  throw new Error('no clipboard command worked (pbcopy, wl-copy, xclip, Set-Clipboard)')
82}
83
84async function copyAs($: EngineInterface, result: QueryResult, as: 'Markdown' | 'CSV') {
85  const text = as === 'Markdown' ? toMarkdown(result.columns, result.rows) : toCsv(result.columns, result.rows)
86  await toClipboard($, text).then(
87    () => $.ui.toast(`Copied ${plural(result.rows.length, 'row')} as ${as}.`),
88    error => $.ui.toast(`og-sf-soql: ${error instanceof Error ? error.message : String(error)}`),
89  )
90}
91
92async function openRecord($: EngineInterface, result: QueryResult, id: string) {
93  const ran = await runSf($, ['org', 'open', '--path', `/${id}`, ...(result.org !== undefined ? ['--target-org', result.org] : [])]).catch(
94    () => undefined,
95  )
96  $.ui.toast(ran?.exitCode === 0 ? `Opening ${id} in the browser.` : `og-sf-soql: could not open ${id}.`)
97}
98
99export const register: Register = on => {
100  summaries.clear()
101
102  on('session.start', async ($, e, next) => {
103    const started = await next(e)
104    await $.command.register({
105      name: COMMAND,
106      description: 'The last SOQL result as a table (or /og-soql [-o org] [-t] SELECT … to run one)',
107      argumentHint: '[-o org] [-t] [SELECT …]',
108    })
109
110    return started
111  })
112
113  // Claude's queries: the result goes to the pane, the transcript keeps one line.
114  on('tool.call', async ($, e, next) => {
115    const command = (e as { command?: unknown }).command
116    const calls = SHELLS.has(e.tool) && typeof command === 'string' ? queriesIn(command) : []
117    const call = calls.at(-1)
118    if (call === undefined) return next(e)
119    const ran = await next(e)
120    const stdout = (ran.result as { stdout?: unknown } | undefined)?.stdout
121    if (ran.deny !== undefined || ran.isError || typeof stdout !== 'string') return ran
122    const result = await capture($, call, stdout).catch(() => undefined)
123    if (result !== undefined) {
124      await show($, result)
125      summaries.set(e.tool_use_id, summaryOf(result))
126    }
127
128    return ran
129    // A failure here never stops the query: before it ran, it runs; after, its result is replayed.
130  }).catch(($, e, next) => next(e))
131
132  on('ui.render', { component: 'ToolResult' }, async ($, e, next) => {
133    const summary = summaries.get(e.props.tool_use_id)
134    if (summary === undefined || e.props.isErrored) return next(e)
135    const { Text } = $.ui.resolve(e)
136
137    return <Text dimColor>{summary}</Text>
138  })
139
140  on('command.run', { command: COMMAND }, async ($, e) => {
141    const source = parseArgs(e.args)
142    if (source.query === '') {
143      const shown = await read($, view)
144      if (shown === null) return { text: `No query yet. /${COMMAND} SELECT … runs one.` }
145      await update($, view, one => (one === null ? one : { ...one, record: null }))
146      return (await openPane($)) ? {} : { text: summaryOf(shown.result) }
147    }
148    const ran = await runSf($, queryArgs(source))
149    const result = resultFromJson(ran.stdout, source)
150    if (result === undefined) return { text: `Query failed: ${errorOf(ran.stdout, ran.stderr)}` }
151    await show($, result)
152
153    return (await openPane($)) ? {} : { text: summaryOf(result) }
154  }).catch(($, e) => ({ text: `og-sf-soql: ${String(e)}` }))
155
156  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
157    const { Box, Text, Button } = $.ui.resolve(e)
158    const shown: View | null = await read($, view)
159    if (shown === null) return <Text dimColor>No query yet.</Text>
160    const { result } = shown
161    const width = Math.max(30, (e.props.bodyColumns ?? 100) - 4)
162    const setRecord = (record: number | null) => update($, view, one => (one === null ? one : { ...one, record }))
163    const row = shown.record === null ? undefined : result.rows[shown.record]
164
165    if (row !== undefined) {
166      const label = Math.min(30, Math.max(...result.columns.map(column => column.length)))
167      const id = row.Id
168      return (
169        <Box flexDirection="column">
170          <Text bold>
171            {result.object ?? 'Record'} {shown.record === null ? '' : shown.record + 1} of {result.rows.length}
172          </Text>
173          <Box flexDirection="column" marginTop={1}>
174            {result.columns.map(column => (
175              <Text key={`field:${column}`}>
176                <Text dimColor>{column.padEnd(label)}</Text> {row[column] ?? ''}
177              </Text>
178            ))}
179          </Box>
180          <Box marginTop={1} gap={1}>
181            <Button key="back" label="Back" hotkey="b" variant="primary" onPress={() => setRecord(null)} />
182            {id !== undefined && id !== '' && <Button key="open" label="Open in browser" hotkey="o" onPress={() => openRecord($, result, id)} />}
183          </Box>
184        </Box>
185      )
186    }
187
188    const widths = fitWidths(result.columns, result.rows, width)
189    const gap = ' '.repeat(GAP)
190    const line = (values: string[]) => values.map((value, i) => cell(value, widths[i] ?? 6)).join(gap)
191    const drawn = result.rows.slice(0, MAX_ROWS)
192
193    return (
194      <Box flexDirection="column">
195        <Text>{cell(result.query, width).trimEnd()}</Text>
196        <Text dimColor>
197          {[
198            result.org,
199            result.isTooling ? 'Tooling API' : undefined,
200            plural(result.totalSize, 'record'),
201            drawn.length < result.rows.length ? `showing ${drawn.length}` : undefined,
202            'Tab moves · Enter opens a record',
203          ]
204            .filter(Boolean)
205            .join(' · ')}
206        </Text>
207        <Box flexDirection="column" marginTop={1}>
208          <Text bold>{line(result.columns)}</Text>
209          {drawn.length === 0 && <Text dimColor>No records.</Text>}
210          {drawn.map((one, i) => (
211            <Button key={`row:${i}`} plain onPress={() => setRecord(i)}>
212              {line(result.columns.map(column => one[column] ?? ''))}
213            </Button>
214          ))}
215        </Box>
216        <Box marginTop={1} gap={1}>
217          <Button key="markdown" label="Copy Markdown" hotkey="m" variant="primary" onPress={() => copyAs($, result, 'Markdown')} />
218          <Button key="csv" label="Copy CSV" hotkey="v" onPress={() => copyAs($, result, 'CSV')} />
219        </Box>
220      </Box>
221    )
222  })
223}
224

1// Pure logic: finding sf queries in a shell command, reading their records,
2// and laying them out as a table. No `$` here, so tests drive it alone.
3
4import type { QueryResult } from '../types'
5
6/** A query found in a command: its text or file, the org and how sf prints it. */
7export type QueryCall = {
8  query?: string
9  file?: string
10  org?: string
11  isTooling: boolean
12  format: 'human' | 'json' | 'csv'
13}
14
15type Source = { query: string; org?: string; isTooling: boolean }
16
17/** Splits a command line into simple commands of words, quotes removed. */
18const commandsOf = (line: string): string[][] => {
19  const commands: string[][] = []
20  let words: string[] = []
21  let word = ''
22  let hasWord = false
23  const endWord = () => {
24    if (hasWord) words.push(word)
25    word = ''
26    hasWord = false
27  }
28  for (let i = 0; i < line.length; i++) {
29    const c = line[i] as string
30    if (c === "'" || c === '"') {
31      const end = line.indexOf(c, i + 1)
32      word += line.slice(i + 1, end < 0 ? line.length : end)
33      hasWord = true
34      i = end < 0 ? line.length : end
35    } else if (c === '\\' && i + 1 < line.length) {
36      word += line[++i]
37      hasWord = true
38    } else if (c === ' ' || c === '\t') {
39      endWord()
40    } else if (';&|\n()'.includes(c)) {
41      endWord()
42      if (words.length > 0) commands.push(words)
43      words = []
44    } else {
45      word += c
46      hasWord = true
47    }
48  }
49  endWord()
50  if (words.length > 0) commands.push(words)
51
52  return commands
53}
54
55const flag = (args: string[], names: string[]): string | undefined => {
56  for (let i = 0; i < args.length; i++) {
57    const arg = args[i] as string
58    for (const name of names) {
59      if (arg === name) return args[i + 1]
60      if (name.startsWith('--') && arg.startsWith(`${name}=`)) return arg.slice(name.length + 1)
61    }
62  }
63
64  return undefined
65}
66
67const program = (word: string) => (word.split(/[\\/]/).pop() ?? word).toLowerCase().replace(/\.(exe|cmd)$/, '')
68
69/** Every `sf data query` (or `sfdx force:data:soql:query`) a command line runs. */
70export const queriesIn = (line: string): QueryCall[] =>
71  commandsOf(line).flatMap(words => {
72    const start = words.findIndex(word => program(word) === 'sf' || program(word) === 'sfdx')
73    if (start < 0) return []
74    const args = words.slice(start + 1)
75    const isQuery = (args[0] === 'data' && args[1] === 'query') || /^(force:)?data:soql:query$/.test(args[0] ?? '')
76    if (!isQuery) return []
77    const query = flag(args, ['--query', '-q'])
78    const file = flag(args, ['--file', '-f'])
79    const org = flag(args, ['--target-org', '-o', '--targetusername', '-u'])
80    const resultFormat = flag(args, ['--result-format', '-r', '--resultformat'])
81    const format = args.includes('--json') || resultFormat === 'json' ? 'json' : resultFormat === 'csv' ? 'csv' : 'human'
82    if (query === undefined && file === undefined) return []
83    const isTooling = args.includes('--use-tooling-api') || args.includes('-t') || args.includes('--usetoolingapi')
84    const call: QueryCall = query !== undefined ? { query, isTooling, format } : { file, isTooling, format }
85    if (org !== undefined) call.org = org
86
87    return [call]
88  })
89
90/** `/og-soql` arguments: `-o <org>`, `-t`, then the query. */
91export const parseArgs = (text: string): Source => {
92  const words = text.trim().split(/\s+/).filter(Boolean)
93  let org: string | undefined
94  let isTooling = false
95  while (words[0] === '-o' || words[0] === '-t') {
96    if (words.shift() === '-o') org = words.shift()
97    else isTooling = true
98  }
99  const source: Source = { query: words.join(' '), isTooling }
100  if (org !== undefined) source.org = org
101
102  return source
103}
104
105/** The object of the outer FROM, past any subquery. */
106export const objectOf = (query: string) => {
107  let depth = 0
108  let outer = ''
109  for (const c of query) {
110    if (c === '(') depth++
111    else if (c === ')') depth--
112    else if (depth === 0) outer += c
113  }
114
115  return /\bfrom\s+(\w+)/i.exec(outer)?.[1]
116}
117
118/** The SELECT list's field paths, a subquery named by its relationship. */
119const selectFields = (query: string) => {
120  const list = /^\s*select\s+([\s\S]*?)\s+from\s+\w+[^)]*$/i.exec(query.replace(/\([^()]*\)/g, sub => sub.replace(/,/g, ' ')))?.[1]
121  if (list === undefined) return []
122
123  return list
124    .split(',')
125    .map(field => field.trim())
126    .map(field => (field.startsWith('(') ? (/\bfrom\s+(\w+)/i.exec(field)?.[1] ?? field) : field))
127    .filter(Boolean)
128}
129
130type SfRecord = Record<string, unknown>
131
132/** Records as flat rows: `Owner.Name` for a relationship, `[n]` for a subquery, '' for null. */
133export const flatten = (records: SfRecord[]) => {
134  const columns: string[] = []
135  const add = (column: string) => {
136    if (!columns.includes(column)) columns.push(column)
137  }
138  const rowOf = (record: SfRecord, prefix: string, row: Record<string, string>) => {
139    for (const [key, value] of Object.entries(record)) {
140      if (key === 'attributes') continue
141      const path = `${prefix}${key}`
142      if (value !== null && typeof value === 'object' && Array.isArray((value as { records?: unknown }).records)) {
143        add(path)
144        row[path] = `[${(value as { totalSize?: number; records: unknown[] }).totalSize ?? (value as { records: unknown[] }).records.length}]`
145      } else if (value !== null && typeof value === 'object') {
146        rowOf(value as SfRecord, `${path}.`, row)
147      } else {
148        add(path)
149        row[path] = value === null || value === undefined ? '' : String(value)
150      }
151    }
152    return row
153  }
154  const rows = records.map(record => rowOf(record, '', {}))
155
156  // A relationship null in one record and filled in another is its fields' columns alone.
157  return { columns: columns.filter(column => !columns.some(other => other.startsWith(`${column}.`))), rows }
158}
159
160const resultOf = (source: Source, columns: string[], rows: Record<string, string>[], totalSize: number): QueryResult => {
161  const result: QueryResult = {
162    query: source.query,
163    isTooling: source.isTooling,
164    columns: columns.length > 0 ? columns : selectFields(source.query),
165    rows,
166    totalSize,
167  }
168  const object = objectOf(source.query)
169  if (object !== undefined) result.object = object
170  if (source.org !== undefined) result.org = source.org
171
172  return result
173}
174
175/** The result of `sf data query --json`; undefined for an error or anything else. */
176export const resultFromJson = (stdout: string, source: Source): QueryResult | undefined => {
177  let parsed: { status?: number; result?: { records?: unknown; totalSize?: number } }
178  try {
179    parsed = JSON.parse(stdout)
180  } catch {
181    return undefined
182  }
183  const records = parsed.result?.records
184  if (parsed.status !== 0 || !Array.isArray(records)) return undefined
185  const { columns, rows } = flatten(records as SfRecord[])
186
187  return resultOf(source, columns, rows, parsed.result?.totalSize ?? rows.length)
188}
189
190/** CSV (RFC 4180) as rows of fields. */
191export const parseCsv = (text: string): string[][] => {
192  const rows: string[][] = []
193  let row: string[] = []
194  let field = ''
195  let quoted = false
196  for (let i = 0; i < text.length; i++) {
197    const c = text[i] as string
198    if (quoted) {
199      if (c === '"' && text[i + 1] === '"') {
200        field += '"'
201        i++
202      } else if (c === '"') quoted = false
203      else field += c
204    } else if (c === '"') quoted = true
205    else if (c === ',') {
206      row.push(field)
207      field = ''
208    } else if (c === '\n' || c === '\r') {
209      if (c === '\r' && text[i + 1] === '\n') i++
210      row.push(field)
211      rows.push(row)
212      row = []
213      field = ''
214    } else field += c
215  }
216  if (field !== '' || row.length > 0) rows.push([...row, field])
217
218  return rows
219}
220
221/** The result of `sf data query -r csv`. */
222export const resultFromCsv = (stdout: string, source: Source): QueryResult | undefined => {
223  const [header, ...lines] = parseCsv(stdout.trim())
224  if (header === undefined || header.length === 0) return undefined
225  const rows = lines.map(line => Object.fromEntries(header.map((column, i) => [column, line[i] ?? ''])))
226
227  return resultOf(source, header, rows, rows.length)
228}
229
230const csvField = (value: string) => (/[",\n\r]/.test(value) ? `"${value.replace(/"/g, '""')}"` : value)
231
232export const toCsv = (columns: string[], rows: Record<string, string>[]) =>
233  [columns, ...rows.map(row => columns.map(column => row[column] ?? ''))].map(fields => fields.map(csvField).join(',')).join('\n')
234
235const mdCell = (value: string) => value.replace(/\|/g, '\\|').replace(/\s*\n\s*/g, ' ')
236
237export const toMarkdown = (columns: string[], rows: Record<string, string>[]) =>
238  [
239    `| ${columns.map(mdCell).join(' | ')} |`,
240    `| ${columns.map(() => '---').join(' | ')} |`,
241    ...rows.map(row => `| ${columns.map(column => mdCell(row[column] ?? '')).join(' | ')} |`),
242  ].join('\n')
243
244const MIN = 6
245const MAX = 40
246export const GAP = 2
247
248/**
249 * Column widths that fit `width` with GAP spaces between columns: each its
250 * longest value (40 at most); when too wide, the widest shrink first, to 6.
251 */
252export const fitWidths = (columns: string[], rows: Record<string, string>[], width: number) => {
253  const widths = columns.map(column =>
254    Math.min(MAX, Math.max(column.length, ...rows.map(row => [...(row[column] ?? '')].length))),
255  )
256  const room = width - GAP * (columns.length - 1)
257  while (widths.reduce((a, b) => a + b, 0) > room) {
258    const widest = widths.indexOf(Math.max(...widths))
259    if ((widths[widest] ?? 0) <= MIN) break
260    widths[widest] = (widths[widest] ?? MIN) - 1
261  }
262
263  return widths
264}
265
266/** A value padded or cut to a width. */
267export const cell = (value: string, width: number) => {
268  const chars = [...value.replace(/\s+/g, ' ')]
269  return chars.length <= width ? chars.join('').padEnd(width) : `${chars.slice(0, width - 1).join('')}…`
270}
271

1/** A query result as the pane shows it: flat rows of text, one column per field path. */
2export type QueryResult = {
3  query: string
4  /** The org named with -o, when one was. */
5  org?: string
6  isTooling: boolean
7  /** The object queried (`Account`), from the FROM clause. */
8  object?: string
9  columns: string[]
10  rows: Record<string, string>[]
11  /** What the org reported, which may exceed the rows returned. */
12  totalSize: number
13}
14
15/** What the pane shows: the last result, and the record open in it. */
16export type View = {
17  result: QueryResult
18  /** The row shown vertically; null for the table. */
19  record: number | null
20}
21
22declare module 'claude-code' {
23  interface PluginState {
24    'og-sf-soql': {
25      view: View | null
26    }
27  }
28}
29