query-table

1import type { ElementTable, EngineInterface, Register, RenderChildren } from 'claude-code'
2
3import type { BashQuery } from '../types'
4import { bashQueryOf, displaySql, isGroupedColumn, isNumeric, mcpQuerySource, parseBqOutput, parseDbtShow, parseJsonRows, type Table, textOf, withThousands } from './parse'
5
6const MAX_CELL = 32
7const MAX_ROWS = 50
8const MAX_SQL = 10000
9const SEPARATOR = ' │ '
10const BASH_QUERY = { plugin: 'query-table', key: 'bashQuery' } as const
11
12const length = (text: string) => [...text].length
13
14const fit = (text: string, width: number, alignRight: boolean): string => {
15  const clipped = length(text) > width ? `${[...text].slice(0, width - 1).join('')}…` : text
16  const padding = ' '.repeat(width - length(clipped))
17  return alignRight ? padding + clipped : clipped + padding
18}
19
20const layout = (table: Table, available: number) => {
21  const isNumberColumn = table.columns.map((_, i) =>
22    table.rows.some(row => isNumeric(row[i] ?? '')) && table.rows.every(row => row[i] === '∅' || isNumeric(row[i] ?? '')),
23  )
24  const rows = table.rows.map(row =>
25    row.map((cell, i) => (isNumberColumn[i] && isGroupedColumn(table.columns[i] ?? '') ? withThousands(cell) : cell)),
26  )
27  const widths = table.columns.map((column, i) =>
28    Math.min(MAX_CELL, Math.max(length(column), ...rows.map(row => length(row[i] ?? '')))),
29  )
30  let used = 0
31  const visible = widths.findIndex(width => (used += width + length(SEPARATOR)) > available + length(SEPARATOR))
32  const shown = visible < 0 ? widths.length : Math.max(1, visible)
33  return { rows, widths: widths.slice(0, shown), isNumberColumn, hidden: widths.length - shown }
34}
35
36type Query = { source: string; sql?: string }
37
38const toTable = (tool: string, output: unknown, query: Query): Table | undefined => {
39  if (mcpQuerySource(tool)) return parseJsonRows(output) ?? parseDbtShow(textOf(output) ?? '')
40  const stdout = (output as { stdout?: unknown })?.stdout
41  if (typeof stdout !== 'string') return undefined
42  return query.source === 'bq query' ? parseBqOutput(stdout) : parseDbtShow(stdout)
43}
44
45const sqlOf = (input: unknown): string | undefined => {
46  const fields = input as { sql?: unknown; sql_query?: unknown; query?: unknown } | undefined
47  const sql = [fields?.sql, fields?.sql_query, fields?.query].find(value => typeof value === 'string')
48  return typeof sql === 'string' ? displaySql(sql).slice(0, MAX_SQL) : undefined
49}
50
51const drawTable = ({ Box, Text, Code }: ElementTable, source: string, table: Table, available: number, sql?: string) => {
52  const { rows: formattedRows, widths, isNumberColumn, hidden } = layout(table, available)
53  const rows = formattedRows.slice(0, MAX_ROWS)
54  const tableWidth = widths.reduce((sum, width) => sum + width, 0) + length(SEPARATOR) * (widths.length - 1)
55  const ruleWidth = Math.min(available, Math.max(tableWidth, ...(sql ?? '').split('\n').map(length)))
56  const summary = [
57    `${table.rows.length} row${table.rows.length === 1 ? '' : 's'}`,
58    `${table.columns.length} column${table.columns.length === 1 ? '' : 's'}`,
59    hidden ? `${hidden} hidden (too wide)` : '',
60    table.rows.length > MAX_ROWS ? `showing first ${MAX_ROWS}` : '',
61  ].filter(Boolean)
62
63  const line = (cells: string[], render: (text: string, i: number) => RenderChildren) => (
64    <Box flexDirection="row">
65      {widths.map((width, i) => (
66        <Text wrap="truncate">
67          {i ? <Text dimColor>{SEPARATOR}</Text> : ''}
68          {render(fit(cells[i] ?? '', width, isNumberColumn[i] ?? false), i)}
69        </Text>
70      ))}
71    </Box>
72  )
73
74  return (
75    <Box flexDirection="column" alignSelf="flex-start" borderStyle="round" borderColor="gray" paddingX={1}>
76      <Text>
77        <Text bold color="magenta">{source}</Text>
78        <Text dimColor> · {summary.join(' · ')}</Text>
79      </Text>
80      {sql ? <Code source={sql} language="sql" wrap="wrap" /> : ''}
81      {sql ? <Text dimColor wrap="truncate">{'─'.repeat(ruleWidth)}</Text> : ''}
82      {line(table.columns, text => <Text bold underline color="cyan">{text}</Text>)}
83      {rows.map(row =>
84        line(row, (text, i) =>
85          row[i] === '∅' ? (
86            <Text dimColor italic>{text}</Text>
87          ) : isNumberColumn[i] && row[i]?.startsWith('-') ? (
88            <Text color="red">{text}</Text>
89          ) : (
90            <Text>{text}</Text>
91          ),
92        ),
93      )}
94      {table.notes.map(note => <Text dimColor wrap="truncate">{note}</Text>)}
95    </Box>
96  )
97}
98
99const queryOf = async ($: EngineInterface, tool: string, toolUseId: string, input?: unknown): Promise<Query | undefined> => {
100  const source = mcpQuerySource(tool)
101  if (source) return { source, sql: sqlOf(input) }
102  if (tool !== 'Bash') return undefined
103  const { value }: { value?: BashQuery } = await $.state.get({ ...BASH_QUERY, id: toolUseId })
104  return value ? { ...value, sql: value.sql && displaySql(value.sql).slice(0, MAX_SQL) } : { source: 'dbt show' }
105}
106
107export const register: Register = on => {
108  on('ui.render', { component: 'ToolGroup' }, ($, e, next) =>
109    !e.props.isExpanded && e.props.calls.some(call => mcpQuerySource(call.tool))
110      ? next({ ...e, props: { ...e.props, isExpanded: true } })
111      : next(e),
112  )
113
114  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
115    const query = bashQueryOf(e.command)
116    if (query) await $.state.set({ ...BASH_QUERY, id: e.tool_use_id }, query)
117    return next(e)
118  })
119
120  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
121    if (e.props.isRunning || e.props.isErrored) return next(e)
122    const query = await queryOf($, e.props.tool, e.props.tool_use_id, e.props.input)
123    const table = query && toTable(e.props.tool, e.props.output, query)
124    if (!query || !table?.columns.length) return next(e)
125    return query.sql || mcpQuerySource(e.props.tool)
126      ? drawTable($.ui.resolve(e), query.source, table, (e.viewport?.columns ?? 120) - 6, query.sql)
127      : next(e)
128  })
129
130  on('ui.render', { component: 'ToolResult' }, async ($, e, next) => {
131    if (e.props.isErrored) return next(e)
132    const query = await queryOf($, e.props.tool, e.props.tool_use_id)
133    const table = query && toTable(e.props.tool, e.props.output, query)
134    if (!query || !table?.columns.length) return next(e)
135    const elements = $.ui.resolve(e)
136    const drawnAbove = mcpQuerySource(e.props.tool) || query.sql
137    return drawnAbove ? <elements.Box /> : drawTable(elements, query.source, table, (e.viewport?.columns ?? 120) - 6)
138  })
139}
140

1import type { BashQuery } from '../types'
2
3export type Table = { columns: string[]; rows: string[][]; notes: string[] }
4
5const NUMERIC = /^-?\d[\d,]*\.?\d*(e[+-]?\d+)?$/i
6const DBT_ROW = /^│(.*)│$/
7const DBT_BORDER = /^[┌╞├└]/
8const DBT_BANNER = /^\s*dbt(?:-fusion)? v?\d/
9const DBT_TRAILING_DOT = /^-?\d[\d,]*\.$/
10
11export const isNumeric = (value: string): boolean => NUMERIC.test(value.trim())
12
13export const textOf = (output: unknown): string | undefined => {
14  if (typeof output === 'string') return output
15  if (Array.isArray(output)) {
16    const texts = output.map(block => (block as { text?: unknown })?.text).filter(t => typeof t === 'string')
17    return texts.length ? texts.join('\n') : undefined
18  }
19  if (output && typeof output === 'object' && 'content' in output) return textOf((output as { content: unknown }).content)
20  return undefined
21}
22
23const tryJson = (text: string): unknown => {
24  try {
25    return JSON.parse(text)
26  } catch {
27    return undefined
28  }
29}
30
31const isRecord = (value: unknown): value is Record<string, unknown> =>
32  typeof value === 'object' && value !== null && !Array.isArray(value)
33
34const splitConcatenatedObjects = (text: string): string[] => {
35  const objects: string[] = []
36  let depth = 0
37  let start = -1
38  let inString = false
39  let escaped = false
40  for (let i = 0; i < text.length; i++) {
41    const char = text[i]
42    if (inString) {
43      if (char === '"' && !escaped) inString = false
44      escaped = !escaped && char === '\\'
45      continue
46    }
47    if (char === '"') inString = true
48    else if (char === '{' && depth++ === 0) start = i
49    else if (char === '}' && --depth === 0) objects.push(text.slice(start, i + 1))
50  }
51  return objects
52}
53
54const isRecordArray = (value: unknown): value is Record<string, unknown>[] =>
55  Array.isArray(value) && value.length > 0 && value.every(isRecord)
56
57const ROW_KEYS = ['rows', 'results', 'show', 'data']
58
59const wrappedRows = (value: Record<string, unknown>): Record<string, unknown>[] | undefined => {
60  const entries = Object.entries(value)
61  const [, rows] = entries.find(([key, rows]) => ROW_KEYS.includes(key) && isRecordArray(rows)) ?? (entries.length === 1 ? entries[0]! : [])
62  return isRecordArray(rows) ? rows : undefined
63}
64
65const recordsOf = (text: string): Record<string, unknown>[] | undefined => {
66  const whole = tryJson(text.trim())
67  if (isRecordArray(whole)) return whole
68  if (isRecord(whole)) return wrappedRows(whole) ?? [whole]
69  const objects = splitConcatenatedObjects(text).map(tryJson)
70  if (!objects.length || !objects.every(isRecord)) return undefined
71  const records = objects as Record<string, unknown>[]
72  return (records.length === 1 && wrappedRows(records[0]!)) || records
73}
74
75const cell = (value: unknown): string =>
76  value === null || value === undefined ? '∅' : typeof value === 'object' ? JSON.stringify(value) : String(value)
77
78export const parseJsonRows = (output: unknown): Table | undefined => {
79  const text = textOf(output)
80  const records = text ? recordsOf(text) : undefined
81  if (!records) return undefined
82  const columns = [...new Set(records.flatMap(Object.keys))]
83  return { columns, rows: records.map(r => columns.map(c => cell(r[c]))), notes: [] }
84}
85
86const splitDbtRow = (line: string): string[] =>
87  (line.match(DBT_ROW)?.[1] ?? '').split('┆').map(c => c.trim())
88
89const dbtCell = (value: string): string =>
90  value === 'null' ? '∅' : DBT_TRAILING_DOT.test(value) ? value.slice(0, -1) : value
91
92export const parseDbtShow = (stdout: string): Table | undefined => {
93  if (!DBT_BANNER.test(stdout)) return undefined
94  const lines = stdout.split('\n')
95  const start = lines.findIndex(l => l.startsWith('┌'))
96  const end = lines.findIndex(l => l.startsWith('└'))
97  if (start < 0 || end < start) return undefined
98  const body = lines.slice(start, end + 1).filter(l => !DBT_BORDER.test(l) && DBT_ROW.test(l))
99  if (!body.length) return undefined
100  const [header = [], ...rows] = body.map(splitDbtRow)
101  const notes = [...lines.slice(0, start), ...lines.slice(end + 1)]
102    .map(l => l.trim())
103    .filter(l => l && !/^(dbt(-fusion)? v?\d|Loading |Query show_sql|\d+ rows?\.|=+ .* =+$)/.test(l))
104  return { columns: header, rows: rows.map(r => r.map(dbtCell)), notes }
105}
106
107export const displaySql = (sql: string): string => {
108  const lines = sql.replace(/\r\n?/g, '\n').split('\n').map(line => line.trimEnd())
109  const body = lines.slice(lines.findIndex(Boolean), lines.findLastIndex(Boolean) + 1)
110  const indent = Math.min(...body.filter(Boolean).map(line => line.length - line.trimStart().length))
111  return body.map(line => line.slice(indent)).join('\n')
112}
113
114const DBT_SHOW_COMMAND = /(^|[^\w-])dbt\s+show(\s|$)/
115const INLINE_FLAG = /\s--inline(?:=|\s+)/
116
117const unquoteDouble = (text: string): string | undefined => {
118  let value = ''
119  for (let i = 0; i < text.length; i++) {
120    const char = text[i]
121    if (char === '"') return value
122    if (char === '`' || (char === '$' && text[i + 1] === '(')) return undefined
123    if (char === '\\' && '"\\$`\n'.includes(text[i + 1] ?? '')) {
124      if (text[++i] !== '\n') value += text[i]
125      continue
126    }
127    value += char
128  }
129  return undefined
130}
131
132export const inlineDbtSql = (command: string): string | undefined => {
133  if (!DBT_SHOW_COMMAND.test(command)) return undefined
134  const flag = INLINE_FLAG.exec(command)
135  if (!flag) return undefined
136  const rest = command.slice(flag.index + flag[0].length)
137  if (rest.startsWith('"')) return unquoteDouble(rest.slice(1))
138  if (rest.startsWith("'")) {
139    const end = rest.indexOf("'", 1)
140    return end < 0 ? undefined : rest.slice(1, end)
141  }
142  return undefined
143}
144
145const PLAIN_NUMBER = /^(-?)(\d+)(\.\d+)?$/
146const UNGROUPED_COLUMN = /(^|_)(id|key|year)$/i
147
148export const isGroupedColumn = (column: string): boolean => !UNGROUPED_COLUMN.test(column)
149
150export const withThousands = (value: string): string => {
151  const match = PLAIN_NUMBER.exec(value)
152  if (!match) return value
153  const [, sign = '', digits = '', fraction = ''] = match
154  return sign + digits.replace(/\B(?=(\d{3})+$)/g, ',') + fraction
155}
156
157const BQ_QUERY_COMMAND = /(^|[^\w-])bq\s+(?:-\S+\s+)*query(\s|$)/
158const BQ_PROGRESS = /^\s*Waiting on \S+ \.\.\./
159const BQ_BORDER = /^\+(-+\+)+$/
160const SQL_START = /^\s*(\(|select|with|insert|update|delete|merge|create|declare)\b/i
161
162const bqCellBounds = (border: string): [number, number][] => {
163  const corners = [...border].flatMap((char, i) => (char === '+' ? [i] : []))
164  return corners.slice(1).map((end, i) => [(corners[i] ?? 0) + 1, end])
165}
166
167const splitBqRow = (line: string, bounds: [number, number][]): string[] => {
168  const chars = [...line]
169  return bounds.map(([start, end]) => chars.slice(start, end).join('').trim())
170}
171
172const parseBqPretty = (lines: string[]): Table | undefined => {
173  const start = lines.findIndex((l, i) => BQ_BORDER.test(l) && lines[i + 1]?.startsWith('|') && lines[i + 2] === l)
174  const border = lines[start]
175  if (!border) return undefined
176  const bounds = bqCellBounds(border)
177  const isRow = (l: string) => l.startsWith('|') && [...l].length === [...border].length
178  const end = lines.findLastIndex(l => BQ_BORDER.test(l) || isRow(l))
179  const [header, ...rows] = lines.slice(start, end + 1).filter(isRow).map(l => splitBqRow(l, bounds))
180  if (!header) return undefined
181  const notes = [...lines.slice(0, start), ...lines.slice(end + 1)].map(l => l.trim()).filter(Boolean)
182  return { columns: header, rows: rows.map(r => r.map(c => (c === 'NULL' ? '∅' : c))), notes }
183}
184
185export const parseBqOutput = (stdout: string): Table | undefined => {
186  const lines = stdout.split(/\r?\n|\r/).filter(l => !BQ_PROGRESS.test(l))
187  return parseBqPretty(lines) ?? parseJsonRows(lines.join('\n'))
188}
189
190const QUOTED_ARGUMENT = /\s(["'])/g
191
192const bqSql = (command: string, from: number): string | undefined => {
193  QUOTED_ARGUMENT.lastIndex = from
194  for (let match; (match = QUOTED_ARGUMENT.exec(command)); ) {
195    const rest = command.slice(match.index + match[0].length)
196    const end = rest.indexOf("'")
197    const sql = match[1] === '"' ? unquoteDouble(rest) : end < 0 ? undefined : rest.slice(0, end)
198    if (sql && SQL_START.test(sql)) return sql
199  }
200  return undefined
201}
202
203const withSql = (source: BashQuery['source'], sql: string | undefined): BashQuery => (sql ? { source, sql } : { source })
204
205export const bashQueryOf = (command: string): BashQuery | undefined => {
206  if (DBT_SHOW_COMMAND.test(command)) return withSql('dbt show', inlineDbtSql(command))
207  const bq = BQ_QUERY_COMMAND.exec(command)
208  return bq ? withSql('bq query', bqSql(command, bq.index + bq[0].length - 1)) : undefined
209}
210
211const MCP_QUERY_TOOLS: [server: RegExp, tool: RegExp, source: string][] = [
212  [/bigquery/i, /^(execute_sql|execute_query|run_query|query)$/i, 'BigQuery'],
213  [/dbt/i, /^show$/i, 'dbt show'],
214  [/dbt/i, /^execute_sql$/i, 'dbt SQL'],
215]
216
217export const mcpQuerySource = (tool: string): string | undefined => {
218  const split = tool.lastIndexOf('__')
219  if (!tool.startsWith('mcp__') || split < 5) return undefined
220  const server = tool.slice(5, split)
221  const name = tool.slice(split + 2)
222  return MCP_QUERY_TOOLS.find(([serverPattern, toolPattern]) => serverPattern.test(server) && toolPattern.test(name))?.[2]
223}
224

1export type BashQuery = { source: 'dbt show' | 'bq query'; sql?: string }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'query-table': { bashQuery: StateFamily<BashQuery> }
6  }
7}
8