query-table

1import type { ElementTable, EngineInterface, Register, RenderChildren } from 'claude-code'
2
3import type { BashQuery } from '../types'
4import { bashQueryOf, displaySql, isGroupedColumn, isNumeric, parseBqOutput, parseDbtShow, parseJsonRows, type Table, withThousands } from './parse'
5
6const BIGQUERY_TOOL = 'mcp__bigquery__execute_sql'
7const MAX_CELL = 32
8const MAX_ROWS = 50
9const MAX_SQL = 10000
10const SEPARATOR = ' │ '
11const BASH_QUERY = { plugin: 'query-table', key: 'bashQuery' } as const
12
13const length = (text: string) => [...text].length
14
15const fit = (text: string, width: number, alignRight: boolean): string => {
16  const clipped = length(text) > width ? `${[...text].slice(0, width - 1).join('')}…` : text
17  const padding = ' '.repeat(width - length(clipped))
18  return alignRight ? padding + clipped : clipped + padding
19}
20
21const layout = (table: Table, available: number) => {
22  const isNumberColumn = table.columns.map((_, i) =>
23    table.rows.some(row => isNumeric(row[i] ?? '')) && table.rows.every(row => row[i] === '∅' || isNumeric(row[i] ?? '')),
24  )
25  const rows = table.rows.map(row =>
26    row.map((cell, i) => (isNumberColumn[i] && isGroupedColumn(table.columns[i] ?? '') ? withThousands(cell) : cell)),
27  )
28  const widths = table.columns.map((column, i) =>
29    Math.min(MAX_CELL, Math.max(length(column), ...rows.map(row => length(row[i] ?? '')))),
30  )
31  let used = 0
32  const visible = widths.findIndex(width => (used += width + length(SEPARATOR)) > available + length(SEPARATOR))
33  const shown = visible < 0 ? widths.length : Math.max(1, visible)
34  return { rows, widths: widths.slice(0, shown), isNumberColumn, hidden: widths.length - shown }
35}
36
37type Query = { source: string; sql?: string }
38
39const toTable = (tool: string, output: unknown, query: Query): Table | undefined => {
40  if (tool === BIGQUERY_TOOL) return parseJsonRows(output)
41  const stdout = (output as { stdout?: unknown })?.stdout
42  if (typeof stdout !== 'string') return undefined
43  return query.source === 'bq query' ? parseBqOutput(stdout) : parseDbtShow(stdout)
44}
45
46const sqlOf = (input: unknown): string | undefined => {
47  const sql = (input as { sql?: unknown })?.sql
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  if (tool === BIGQUERY_TOOL) return { source: 'BigQuery', sql: sqlOf(input) }
101  if (tool !== 'Bash') return undefined
102  const { value }: { value?: BashQuery } = await $.state.get({ ...BASH_QUERY, id: toolUseId })
103  return value ? { ...value, sql: value.sql && displaySql(value.sql).slice(0, MAX_SQL) } : { source: 'dbt show' }
104}
105
106export const register: Register = on => {
107  on('ui.render', { component: 'ToolGroup' }, ($, e, next) =>
108    !e.props.isExpanded && e.props.calls.some(call => call.tool === BIGQUERY_TOOL)
109      ? next({ ...e, props: { ...e.props, isExpanded: true } })
110      : next(e),
111  )
112
113  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
114    const query = bashQueryOf(e.command)
115    if (query) await $.state.set({ ...BASH_QUERY, id: e.tool_use_id }, query)
116    return next(e)
117  })
118
119  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
120    if (e.props.isRunning || e.props.isErrored) return next(e)
121    const query = await queryOf($, e.props.tool, e.props.tool_use_id, e.props.input)
122    const table = query && toTable(e.props.tool, e.props.output, query)
123    if (!query || !table?.columns.length) return next(e)
124    return query.sql || e.props.tool === BIGQUERY_TOOL
125      ? drawTable($.ui.resolve(e), query.source, table, (e.viewport?.columns ?? 120) - 6, query.sql)
126      : next(e)
127  })
128
129  on('ui.render', { component: 'ToolResult' }, async ($, e, next) => {
130    if (e.props.isErrored) return next(e)
131    const query = await queryOf($, e.props.tool, e.props.tool_use_id)
132    const table = query && toTable(e.props.tool, e.props.output, query)
133    if (!query || !table?.columns.length) return next(e)
134    const elements = $.ui.resolve(e)
135    const drawnAbove = e.props.tool === BIGQUERY_TOOL || query.sql
136    return drawnAbove ? <elements.Box /> : drawTable(elements, query.source, table, (e.viewport?.columns ?? 120) - 6)
137  })
138}
139

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 /
9const DBT_TRAILING_DOT = /^-?\d[\d,]*\.$/
10
11export const isNumeric = (value: string): boolean => NUMERIC.test(value.trim())
12
13const 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 recordsOf = (text: string): Record<string, unknown>[] | undefined => {
58  const whole = tryJson(text.trim())
59  if (isRecordArray(whole)) return whole
60  if (isRecord(whole)) {
61    const values = Object.values(whole)
62    return values.length === 1 && isRecordArray(values[0]) ? values[0] : [whole]
63  }
64  const objects = splitConcatenatedObjects(text).map(tryJson)
65  return objects.length && objects.every(isRecord) ? (objects as Record<string, unknown>[]) : undefined
66}
67
68const cell = (value: unknown): string =>
69  value === null || value === undefined ? '∅' : typeof value === 'object' ? JSON.stringify(value) : String(value)
70
71export const parseJsonRows = (output: unknown): Table | undefined => {
72  const text = textOf(output)
73  const records = text ? recordsOf(text) : undefined
74  if (!records) return undefined
75  const columns = [...new Set(records.flatMap(Object.keys))]
76  return { columns, rows: records.map(r => columns.map(c => cell(r[c]))), notes: [] }
77}
78
79const splitDbtRow = (line: string): string[] =>
80  (line.match(DBT_ROW)?.[1] ?? '').split('┆').map(c => c.trim())
81
82const dbtCell = (value: string): string =>
83  value === 'null' ? '∅' : DBT_TRAILING_DOT.test(value) ? value.slice(0, -1) : value
84
85export const parseDbtShow = (stdout: string): Table | undefined => {
86  if (!DBT_BANNER.test(stdout)) return undefined
87  const lines = stdout.split('\n')
88  const start = lines.findIndex(l => l.startsWith('┌'))
89  const end = lines.findIndex(l => l.startsWith('└'))
90  if (start < 0 || end < start) return undefined
91  const body = lines.slice(start, end + 1).filter(l => !DBT_BORDER.test(l) && DBT_ROW.test(l))
92  if (!body.length) return undefined
93  const [header = [], ...rows] = body.map(splitDbtRow)
94  const notes = [...lines.slice(0, start), ...lines.slice(end + 1)]
95    .map(l => l.trim())
96    .filter(l => l && !/^(dbt-fusion|Loading |Query show_sql|\d+ rows?\.|=+ .* =+$)/.test(l))
97  return { columns: header, rows: rows.map(r => r.map(dbtCell)), notes }
98}
99
100export const displaySql = (sql: string): string => {
101  const lines = sql.replace(/\r\n?/g, '\n').split('\n').map(line => line.trimEnd())
102  const body = lines.slice(lines.findIndex(Boolean), lines.findLastIndex(Boolean) + 1)
103  const indent = Math.min(...body.filter(Boolean).map(line => line.length - line.trimStart().length))
104  return body.map(line => line.slice(indent)).join('\n')
105}
106
107const DBT_SHOW_COMMAND = /(^|[^\w-])dbt\s+show(\s|$)/
108const INLINE_FLAG = /\s--inline(?:=|\s+)/
109
110const unquoteDouble = (text: string): string | undefined => {
111  let value = ''
112  for (let i = 0; i < text.length; i++) {
113    const char = text[i]
114    if (char === '"') return value
115    if (char === '`' || (char === '$' && text[i + 1] === '(')) return undefined
116    if (char === '\\' && '"\\$`\n'.includes(text[i + 1] ?? '')) {
117      if (text[++i] !== '\n') value += text[i]
118      continue
119    }
120    value += char
121  }
122  return undefined
123}
124
125export const inlineDbtSql = (command: string): string | undefined => {
126  if (!DBT_SHOW_COMMAND.test(command)) return undefined
127  const flag = INLINE_FLAG.exec(command)
128  if (!flag) return undefined
129  const rest = command.slice(flag.index + flag[0].length)
130  if (rest.startsWith('"')) return unquoteDouble(rest.slice(1))
131  if (rest.startsWith("'")) {
132    const end = rest.indexOf("'", 1)
133    return end < 0 ? undefined : rest.slice(1, end)
134  }
135  return undefined
136}
137
138const PLAIN_NUMBER = /^(-?)(\d+)(\.\d+)?$/
139const UNGROUPED_COLUMN = /(^|_)(id|key|year)$/i
140
141export const isGroupedColumn = (column: string): boolean => !UNGROUPED_COLUMN.test(column)
142
143export const withThousands = (value: string): string => {
144  const match = PLAIN_NUMBER.exec(value)
145  if (!match) return value
146  const [, sign = '', digits = '', fraction = ''] = match
147  return sign + digits.replace(/\B(?=(\d{3})+$)/g, ',') + fraction
148}
149
150const BQ_QUERY_COMMAND = /(^|[^\w-])bq\s+(?:-\S+\s+)*query(\s|$)/
151const BQ_PROGRESS = /^\s*Waiting on \S+ \.\.\./
152const BQ_BORDER = /^\+(-+\+)+$/
153const SQL_START = /^\s*(\(|select|with|insert|update|delete|merge|create|declare)\b/i
154
155const bqCellBounds = (border: string): [number, number][] => {
156  const corners = [...border].flatMap((char, i) => (char === '+' ? [i] : []))
157  return corners.slice(1).map((end, i) => [(corners[i] ?? 0) + 1, end])
158}
159
160const splitBqRow = (line: string, bounds: [number, number][]): string[] => {
161  const chars = [...line]
162  return bounds.map(([start, end]) => chars.slice(start, end).join('').trim())
163}
164
165const parseBqPretty = (lines: string[]): Table | undefined => {
166  const start = lines.findIndex((l, i) => BQ_BORDER.test(l) && lines[i + 1]?.startsWith('|') && lines[i + 2] === l)
167  const border = lines[start]
168  if (!border) return undefined
169  const bounds = bqCellBounds(border)
170  const isRow = (l: string) => l.startsWith('|') && [...l].length === [...border].length
171  const end = lines.findLastIndex(l => BQ_BORDER.test(l) || isRow(l))
172  const [header, ...rows] = lines.slice(start, end + 1).filter(isRow).map(l => splitBqRow(l, bounds))
173  if (!header) return undefined
174  const notes = [...lines.slice(0, start), ...lines.slice(end + 1)].map(l => l.trim()).filter(Boolean)
175  return { columns: header, rows: rows.map(r => r.map(c => (c === 'NULL' ? '∅' : c))), notes }
176}
177
178export const parseBqOutput = (stdout: string): Table | undefined => {
179  const lines = stdout.split(/\r?\n|\r/).filter(l => !BQ_PROGRESS.test(l))
180  return parseBqPretty(lines) ?? parseJsonRows(lines.join('\n'))
181}
182
183const QUOTED_ARGUMENT = /\s(["'])/g
184
185const bqSql = (command: string, from: number): string | undefined => {
186  QUOTED_ARGUMENT.lastIndex = from
187  for (let match; (match = QUOTED_ARGUMENT.exec(command)); ) {
188    const rest = command.slice(match.index + match[0].length)
189    const end = rest.indexOf("'")
190    const sql = match[1] === '"' ? unquoteDouble(rest) : end < 0 ? undefined : rest.slice(0, end)
191    if (sql && SQL_START.test(sql)) return sql
192  }
193  return undefined
194}
195
196const withSql = (source: BashQuery['source'], sql: string | undefined): BashQuery => (sql ? { source, sql } : { source })
197
198export const bashQueryOf = (command: string): BashQuery | undefined => {
199  if (DBT_SHOW_COMMAND.test(command)) return withSql('dbt show', inlineDbtSql(command))
200  const bq = BQ_QUERY_COMMAND.exec(command)
201  return bq ? withSql('bq query', bqSql(command, bq.index + bq[0].length - 1)) : undefined
202}
203

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