From b9cf959808859c96c904722d970ac1190b733821 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Mon, 2 Nov 2026 15:36:28 -0600 Subject: [PATCH 1/7] feat(mcp): add clickstack_describe_table tool Agents use clickstack_sql to query system.columns / system.tables to inspect tables that are not registered as sources. Add an off-the-shelf tool for it: list the tables in a database, or describe one table's columns, keys, engine, skipping indices, approximate size, and an optional small row sample. --- .changeset/mcp-describe-table.md | 5 + MCP.md | 1 + .../mcp/__tests__/describeTable.int.test.ts | 217 +++++++++++ .../src/mcp/__tests__/describeTable.test.ts | 360 ++++++++++++++++++ packages/api/src/mcp/tools/query/sql.ts | 4 +- .../src/mcp/tools/sources/describeTable.ts | 277 ++++++++++++++ .../mcp/tools/sources/describeTableQueries.ts | 198 ++++++++++ .../mcp/tools/sources/describeTableSchema.ts | 83 ++++ packages/api/src/mcp/tools/sources/index.ts | 2 + 9 files changed, 1146 insertions(+), 1 deletion(-) create mode 100644 .changeset/mcp-describe-table.md create mode 100644 packages/api/src/mcp/__tests__/describeTable.int.test.ts create mode 100644 packages/api/src/mcp/__tests__/describeTable.test.ts create mode 100644 packages/api/src/mcp/tools/sources/describeTable.ts create mode 100644 packages/api/src/mcp/tools/sources/describeTableQueries.ts create mode 100644 packages/api/src/mcp/tools/sources/describeTableSchema.ts diff --git a/.changeset/mcp-describe-table.md b/.changeset/mcp-describe-table.md new file mode 100644 index 0000000000..bd4e022999 --- /dev/null +++ b/.changeset/mcp-describe-table.md @@ -0,0 +1,5 @@ +--- +'@hyperdx/api': minor +--- + +feat(mcp): add `clickstack_describe_table` to list the tables in a database and describe a table that is not a registered source (columns, keys, engine, skipping indices, approximate size, optional row sample) diff --git a/MCP.md b/MCP.md index a380fb8909..d88f53b99d 100644 --- a/MCP.md +++ b/MCP.md @@ -122,6 +122,7 @@ with: | ----------------------------- | -------------------------------------------------------------------------------------------- | | `clickstack_list_sources` | List all data sources and connections as a lightweight catalog (IDs, names, kinds) | | `clickstack_describe_source` | Full column schema, attribute keys, sampled low-cardinality values, and a round-trippable `config` block (the exact shape `clickstack_save_source` accepts, incl. correlation IDs) for a single source; for metric sources also returns a per-kind metric-name sample | +| `clickstack_describe_table` | Inspect a ClickHouse table that is not a registered source: columns, keys, engine, skipping indices, approximate size, and an optional row sample. Without a table, lists the tables in a database (paginated) | | `clickstack_save_source` | Create (omit id) or update (provide id) a data source (log, trace, session, metric, or promql) so shipped telemetry becomes queryable. To clone a source, read its `config` from `clickstack_describe_source`, drop `id`, and pass it here | | `clickstack_delete_source` | Permanently delete a data source by ID | | `clickstack_list_metrics` | Paginated catalog of metric names on a metric source with optional kind / namePattern / time-window filters | diff --git a/packages/api/src/mcp/__tests__/describeTable.int.test.ts b/packages/api/src/mcp/__tests__/describeTable.int.test.ts new file mode 100644 index 0000000000..864a97986d --- /dev/null +++ b/packages/api/src/mcp/__tests__/describeTable.int.test.ts @@ -0,0 +1,217 @@ +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import mongoose from 'mongoose'; + +import * as config from '@/config'; +import { + getLoggedInAgent, + getServer, + getTestFixtureClickHouseClient, +} from '@/fixtures'; +import { McpContext } from '@/mcp/tools/types'; +import Connection from '@/models/connection'; + +import { callTool, createTestClient, getFirstText } from './mcpTestUtils'; + +const DATABASE = 'mcp_describe_table_test'; + +describe('clickstack_describe_table', () => { + const server = getServer(); + let connection: any; + let client: Client; + + beforeAll(async () => { + await server.start(); + const ch = await getTestFixtureClickHouseClient(); + await ch.command({ query: `DROP DATABASE IF EXISTS ${DATABASE}` }); + await ch.command({ query: `CREATE DATABASE ${DATABASE}` }); + await ch.command({ + query: ` + CREATE TABLE ${DATABASE}.events ( + Timestamp DateTime64(9), + ServiceName LowCardinality(String) COMMENT 'emitting service', + Body String, + Day Date MATERIALIZED toDate(Timestamp), + INDEX idx_body Body TYPE tokenbf_v1(8192, 3, 0) GRANULARITY 1 + ) + ENGINE = MergeTree + PARTITION BY toDate(Timestamp) + PRIMARY KEY ServiceName + ORDER BY (ServiceName, Timestamp) + COMMENT 'test events' + `, + }); + await ch.command({ + query: `CREATE VIEW ${DATABASE}.events_view AS SELECT ServiceName FROM ${DATABASE}.events`, + }); + await ch.insert({ + table: `${DATABASE}.events`, + values: [ + { Timestamp: '2024-01-01 00:00:00', ServiceName: 'api', Body: 'a' }, + { Timestamp: '2024-01-01 00:00:01', ServiceName: 'web', Body: 'b' }, + { Timestamp: '2024-01-01 00:00:02', ServiceName: 'api', Body: 'c' }, + ], + format: 'JSONEachRow', + }); + }); + + beforeEach(async () => { + const { team, user } = await getLoggedInAgent(server); + connection = await Connection.create({ + team: team._id, + name: 'Default', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + const context: McpContext = { + teamId: team._id.toString(), + userId: user._id.toString(), + }; + client = await createTestClient(context); + }); + + afterEach(async () => { + await client?.close(); + await server.clearDBs(); + }); + + afterAll(async () => { + const ch = await getTestFixtureClickHouseClient(); + await ch.command({ query: `DROP DATABASE IF EXISTS ${DATABASE}` }); + await server.stop(); + }); + + const call = (args: Record) => + callTool(client, 'clickstack_describe_table', { + connectionId: connection._id.toString(), + ...args, + }); + + it('describes columns, keys, indices, and size of a table', async () => { + const result = await call({ database: DATABASE, table: 'events' }); + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + + expect(output.table).toMatchObject({ + database: DATABASE, + name: 'events', + engine: 'MergeTree', + sortingKey: 'ServiceName, Timestamp', + primaryKey: 'ServiceName', + partitionKey: 'toDate(Timestamp)', + comment: 'test events', + approxTotalRows: 3, + }); + expect(output.table.approxTotalBytes).toBeGreaterThan(0); + + expect(output.columns).toEqual([ + { + name: 'Timestamp', + type: 'DateTime64(9)', + inSortingKey: true, + inPrimaryKey: false, + }, + { + name: 'ServiceName', + type: 'LowCardinality(String)', + comment: 'emitting service', + inSortingKey: true, + inPrimaryKey: true, + }, + { + name: 'Body', + type: 'String', + inSortingKey: false, + inPrimaryKey: false, + }, + { + name: 'Day', + type: 'Date', + defaultKind: 'MATERIALIZED', + defaultExpression: 'toDate(Timestamp)', + inSortingKey: false, + inPrimaryKey: false, + }, + ]); + expect(output.dataSkippingIndices).toEqual([ + { name: 'idx_body', type: 'tokenbf_v1', expression: 'Body' }, + ]); + expect(output.sampleRows).toBeUndefined(); + }); + + it('returns a capped sample of rows when asked', async () => { + const result = await call({ + database: DATABASE, + table: 'events', + sampleRows: 2, + }); + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.sampleRows).toHaveLength(2); + expect(output.sampleRows[0]).toHaveProperty('ServiceName'); + }); + + it('rejects more than 5 sample rows', async () => { + const result = await call({ + database: DATABASE, + table: 'events', + sampleRows: 50, + }); + expect(result.isError).toBe(true); + }); + + it('lists tables in a database across pages', async () => { + const first = JSON.parse( + getFirstText(await call({ database: DATABASE, limit: 1 })), + ); + expect(first.tables).toEqual([ + { name: 'events', engine: 'MergeTree', approxRows: 3 }, + ]); + expect(first.nextCursor).toEqual(expect.any(String)); + + const second = JSON.parse( + getFirstText( + await call({ database: DATABASE, limit: 1, cursor: first.nextCursor }), + ), + ); + expect(second.tables).toEqual([ + { name: 'events_view', engine: 'View', approxRows: null }, + ]); + expect(second.nextCursor).toBeUndefined(); + }); + + it('returns user errors for a missing table or database', async () => { + const missingTable = await call({ database: DATABASE, table: 'nope' }); + expect(missingTable.isError).toBe(true); + expect(getFirstText(missingTable)).toMatch(/not found/); + + const missingDb = await call({ database: 'mcp_no_such_db' }); + expect(missingDb.isError).toBe(true); + expect(getFirstText(missingDb)).toMatch(/not found/); + }); + + it('treats database and table names as values, not SQL', async () => { + const result = await call({ + database: DATABASE, + table: "events' OR '1'='1", + }); + expect(result.isError).toBe(true); + expect(getFirstText(result)).toMatch(/not found/); + }); + + it("does not use another team's connection", async () => { + const otherConnection = await Connection.create({ + team: new mongoose.Types.ObjectId(), + name: 'Other', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + const result = await callTool(client, 'clickstack_describe_table', { + connectionId: otherConnection._id.toString(), + database: DATABASE, + }); + expect(result.isError).toBe(true); + expect(getFirstText(result)).toMatch(/Connection .* not found/); + }); +}); diff --git a/packages/api/src/mcp/__tests__/describeTable.test.ts b/packages/api/src/mcp/__tests__/describeTable.test.ts new file mode 100644 index 0000000000..bc725ab2f9 --- /dev/null +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -0,0 +1,360 @@ +const mockQuery = jest.fn(); +const mockGetConnectionById = jest.fn(); + +jest.mock('@/models/source', () => ({})); +jest.mock('@/controllers/sources', () => ({})); +jest.mock('@/controllers/connection', () => ({ + getConnectionById: (...args: unknown[]) => mockGetConnectionById(...args), +})); +jest.mock('@/clickhouse', () => ({ + ClickhouseClient: jest.fn(() => ({ query: mockQuery })), +})); + +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; + +import { registerDescribeTable } from '@/mcp/tools/sources/describeTable'; +import { + shapeTableDescription, + shapeTableListEntry, +} from '@/mcp/tools/sources/describeTableQueries'; +import { + describeTableSchema, + encodeTableListCursor, + validateDescribeTableInput, +} from '@/mcp/tools/sources/describeTableSchema'; +import type { RegisterToolFn, ToolResult } from '@/mcp/tools/types'; +import { getErrorCategory, type McpErrorResult } from '@/mcp/utils/errors'; + +const CONNECTION_ID = '65f5e4a3b9e77c001a567890'; + +type QueryArgs = { query: string; query_params: Record }; + +const queryCalls = (): QueryArgs[] => + mockQuery.mock.calls.map(([args]: [QueryArgs]) => args); + +function respondWith(routes: Record) { + mockQuery.mockImplementation(async ({ query: sql }: QueryArgs) => { + const rows = + Object.entries(routes).find(([fragment]) => + sql.includes(fragment), + )?.[1] ?? []; + if (rows instanceof Error) throw rows; + return { json: async () => ({ data: rows }) }; + }); +} + +function getHandler() { + let handler: + | ((args: Record) => Promise) + | undefined; + const registerTool: RegisterToolFn = (_name, config, h) => { + handler = args => h(config.inputSchema.parse(args)); + }; + registerDescribeTable({ + server: new McpServer({ name: 'test', version: '1.0.0' }), + context: { teamId: 'team-1', userId: 'user-1' }, + registerTool, + }); + return handler!; +} + +function errorCategory(result: ToolResult) { + const isMcpError = (r: ToolResult): r is ToolResult & McpErrorResult => + r.isError === true; + return isMcpError(result) ? getErrorCategory(result) : undefined; +} + +const parseText = (result: ToolResult) => JSON.parse(result.content[0].text); + +const TABLE_ROW = { + name: 'events', + engine: 'MergeTree', + sorting_key: 'ServiceName, Timestamp', + primary_key: 'ServiceName', + partition_key: 'toDate(Timestamp)', + comment: '', + total_rows: '1200', + total_bytes: '34567', +}; + +const COLUMN_ROWS = [ + { + name: 'Timestamp', + type: 'DateTime64(9)', + default_kind: '', + default_expression: '', + comment: '', + is_in_sorting_key: 1, + is_in_primary_key: 0, + }, + { + name: 'Day', + type: 'Date', + default_kind: 'MATERIALIZED', + default_expression: 'toDate(Timestamp)', + comment: 'partition day', + is_in_sorting_key: 0, + is_in_primary_key: 0, + }, +]; + +describe('clickstack_describe_table', () => { + beforeEach(() => { + mockQuery.mockReset(); + mockGetConnectionById.mockReset().mockResolvedValue({ + host: 'http://localhost:8123', + username: 'default', + password: '', + }); + }); + + describe('input validation', () => { + it('caps sampleRows at 5 and defaults it to 0', () => { + expect( + describeTableSchema.safeParse({ + connectionId: CONNECTION_ID, + database: 'default', + table: 't', + sampleRows: 6, + }).success, + ).toBe(false); + const parsed = describeTableSchema.parse({ + connectionId: CONNECTION_ID, + database: 'default', + }); + expect(parsed.sampleRows).toBe(0); + expect(parsed.limit).toBe(50); + }); + + it('rejects an empty database name', () => { + expect( + describeTableSchema.safeParse({ + connectionId: CONNECTION_ID, + database: '', + }).success, + ).toBe(false); + }); + + it('rejects sampleRows without a table and cursor with a table', () => { + expect( + validateDescribeTableInput({ sampleRows: 2, cursor: undefined }), + ).toMatch(/sampleRows requires table/); + expect( + validateDescribeTableInput({ table: 't', sampleRows: 0, cursor: 'x' }), + ).toMatch(/cursor only applies when listing/); + expect( + validateDescribeTableInput({ table: 't', sampleRows: 3 }), + ).toBeNull(); + }); + + it('returns a user error for a malformed connectionId without a lookup', async () => { + const result = await getHandler()({ + connectionId: 'not-an-id', + database: 'default', + }); + expect(result.isError).toBe(true); + expect(mockGetConnectionById).not.toHaveBeenCalled(); + }); + + it('returns a user error when the connection is not in the team', async () => { + mockGetConnectionById.mockResolvedValue(null); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + }); + expect(result.isError).toBe(true); + expect(mockGetConnectionById).toHaveBeenCalledWith( + 'team-1', + CONNECTION_ID, + true, + ); + }); + }); + + describe('output shaping', () => { + it('shapes a table description from system table rows', () => { + expect( + shapeTableDescription({ + database: 'default', + table: { ...TABLE_ROW, total_rows: null, total_bytes: null }, + columns: COLUMN_ROWS, + indices: [{ name: 'idx_day', type: 'minmax', expr: 'Day' }], + }), + ).toEqual({ + table: { + database: 'default', + name: 'events', + engine: 'MergeTree', + sortingKey: 'ServiceName, Timestamp', + primaryKey: 'ServiceName', + partitionKey: 'toDate(Timestamp)', + approxTotalRows: null, + approxTotalBytes: null, + }, + columns: [ + { + name: 'Timestamp', + type: 'DateTime64(9)', + inSortingKey: true, + inPrimaryKey: false, + }, + { + name: 'Day', + type: 'Date', + defaultKind: 'MATERIALIZED', + defaultExpression: 'toDate(Timestamp)', + comment: 'partition day', + inSortingKey: false, + inPrimaryKey: false, + }, + ], + dataSkippingIndices: [ + { name: 'idx_day', type: 'minmax', expression: 'Day' }, + ], + }); + }); + + it('converts quoted 64-bit row counts to numbers', () => { + expect( + shapeTableListEntry({ + name: 'events', + engine: 'MergeTree', + total_rows: '42', + }), + ).toEqual({ name: 'events', engine: 'MergeTree', approxRows: 42 }); + }); + + it('binds database and table names as query parameters', async () => { + respondWith({ 'FROM system.tables': [TABLE_ROW] }); + await getHandler()({ + connectionId: CONNECTION_ID, + database: "evil'db", + table: 'events`; DROP TABLE x', + }); + for (const { query: sql, query_params } of queryCalls()) { + expect(sql).not.toContain('evil'); + expect(sql).not.toContain('DROP'); + expect(Object.values(query_params)).toContain("evil'db"); + } + }); + + it('pages the table list with an opaque cursor', async () => { + respondWith({ + 'FROM system.tables': [ + { name: 'a', engine: 'MergeTree', total_rows: '1' }, + { name: 'b', engine: 'View', total_rows: null }, + { name: 'c', engine: 'MergeTree', total_rows: '3' }, + ], + }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + limit: 2, + }), + ); + expect(output.tables).toEqual([ + { name: 'a', engine: 'MergeTree', approxRows: 1 }, + { name: 'b', engine: 'View', approxRows: null }, + ]); + expect(output.nextCursor).toBe(encodeTableListCursor({ lastName: 'b' })); + }); + + it('rejects a malformed cursor', async () => { + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + cursor: 'garbage', + }); + expect(result.isError).toBe(true); + expect(result.content[0].text).toMatch(/Invalid cursor/); + }); + + it('reports a missing database as a user error', async () => { + respondWith({}); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'nope', + }); + expect(result.isError).toBe(true); + expect(errorCategory(result)).toBe('user'); + expect(result.content[0].text).toMatch(/Database "nope" not found/); + }); + + it('reports a missing table as a user error', async () => { + respondWith({}); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'nope', + }); + expect(result.isError).toBe(true); + expect(errorCategory(result)).toBe('user'); + expect(result.content[0].text).toMatch(/default\.nope" not found/); + }); + + it('maps permission errors to user errors with a hint', async () => { + respondWith({ + 'FROM system.columns': new Error( + 'Code: 497. DB::Exception: default: Not enough privileges. (ACCESS_DENIED)', + ), + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }); + expect(result.isError).toBe(true); + expect(errorCategory(result)).toBe('user'); + expect(result.content[0].text).toMatch(/lacks the privileges/); + }); + + it('skips sampling by default', async () => { + respondWith({ 'FROM system.tables': [TABLE_ROW] }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }), + ); + expect(output.sampleRows).toBeUndefined(); + expect(queryCalls().some(q => q.query.includes('SELECT *'))).toBe(false); + }); + + it('trims oversized sample rows', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + 'SELECT *': [{ Body: 'x'.repeat(50_000) }, { Body: 'short' }], + }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 2, + }), + ); + expect(output.sampleRows).toHaveLength(2); + expect(output.sampleRows[0].__hdx_trimmed).toBe(true); + expect(output.sampleRows[1]).toEqual({ Body: 'short' }); + }); + + it('keeps the schema when sampling fails', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + 'SELECT *': new Error('ACCESS_DENIED'), + }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 1, + }), + ); + expect(output.table.name).toBe('events'); + expect(output.sampleRowsError).toBe('ACCESS_DENIED'); + }); + }); +}); diff --git a/packages/api/src/mcp/tools/query/sql.ts b/packages/api/src/mcp/tools/query/sql.ts index 28ac32a851..4d09b27874 100644 --- a/packages/api/src/mcp/tools/query/sql.ts +++ b/packages/api/src/mcp/tools/query/sql.ts @@ -94,7 +94,9 @@ export function registerSql({ context, registerTool }: ToolRegistrar) { 'fall back to SQL if it cannot express what you need.\n\n' + 'Requires connectionId — call clickstack_list_sources to find connections. Pass sourceId ' + 'too when querying a source, so $__sourceTable resolves to its table. ' + - 'Call clickstack_describe_source to discover column names before writing SQL.\n\n' + + 'Call clickstack_describe_source to discover column names before writing SQL. ' + + 'For tables that are not sources, call clickstack_describe_table instead of ' + + 'querying system.columns / system.tables directly.\n\n' + 'Results are always returned as table rows — for time-series semantics, ' + 'include a time column and ORDER BY it in your SQL.', inputSchema: sqlSchema, diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts new file mode 100644 index 0000000000..7232296c89 --- /dev/null +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -0,0 +1,277 @@ +import { trace } from '@opentelemetry/api'; + +import { ClickhouseClient } from '@/clickhouse'; +import { getConnectionById } from '@/controllers/connection'; +import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; +import { + clickHouseErrorResult, + MCP_REQUEST_TIMEOUT, +} from '@/mcp/tools/query/helpers'; +import type { ToolRegistrar, ToolResult } from '@/mcp/tools/types'; +import { + mcpServerError, + mcpUserError, + sanitizeFetchError, + validateObjectId, +} from '@/mcp/utils/errors'; +import { MCP_TOOL_TIMEOUT_MS, runWithTimeout } from '@/mcp/utils/timeout'; +import logger from '@/utils/logger'; +import { trimToolResponse } from '@/utils/trimToolResponse'; + +import { + databaseExists, + fetchSampleRows, + fetchTableDetails, + fetchTablePage, + shapeTableDescription, + shapeTableListEntry, + type TableQueryContext, +} from './describeTableQueries'; +import { + decodeTableListCursor, + type DescribeTableInput, + describeTableSchema, + encodeTableListCursor, + MAX_SAMPLE_ROWS, + validateDescribeTableInput, +} from './describeTableSchema'; + +// Sample rows can hold arbitrarily large values; keep them a small share of +// the response so they never crowd out the schema. +const SAMPLE_ROWS_MAX_CHARS = 10_000; + +const quoteName = (name: string) => `\`${name.replace(/`/g, '\\`')}\``; + +function jsonResult(payload: unknown): ToolResult { + return { content: [{ type: 'text', text: JSON.stringify(payload) }] }; +} + +/** + * Not-found and permission failures are the caller's to fix, so they go back + * as user errors with a hint; only infrastructure failures stay server errors. + */ +function tableQueryErrorResult(e: unknown, prefix: string) { + const message = e instanceof Error ? e.message : String(e); + let suffix: string | undefined; + if ( + /ACCESS_DENIED|NOT_ENOUGH_PRIVILEGES|Not enough privileges/i.test(message) + ) { + suffix = + "The ClickHouse user on this connection lacks the privileges needed. Ask the connection's administrator to grant them, or use another connection."; + } else if (/UNKNOWN_(TABLE|DATABASE)|does not exist/i.test(message)) { + suffix = + 'Call clickstack_describe_table with only database to list the tables it contains.'; + } + return clickHouseErrorResult(e, { prefix, suffix }); +} + +async function listTables( + ctx: TableQueryContext, + input: DescribeTableInput, +): Promise { + const cursor = input.cursor ? decodeTableListCursor(input.cursor) : null; + if (input.cursor && !cursor) { + return mcpUserError( + 'Invalid cursor. Omit cursor to start over, or pass the exact `nextCursor` value returned by a previous call.', + ); + } + + let rows; + try { + rows = await fetchTablePage(ctx, { + database: input.database, + afterName: cursor?.lastName, + limit: input.limit, + }); + } catch (e) { + return tableQueryErrorResult(e, 'Failed to list tables'); + } + + // system.tables returns nothing for a missing database rather than erroring. + if (rows.length === 0 && !cursor) { + try { + if (!(await databaseExists(ctx, input.database))) { + return mcpUserError( + `Database "${input.database}" not found or not visible to this connection's ClickHouse user.`, + ); + } + } catch (e) { + return tableQueryErrorResult(e, 'Failed to look up database'); + } + } + + const page = rows.slice(0, input.limit); + const nextCursor = + rows.length > input.limit + ? encodeTableListCursor({ lastName: page[page.length - 1].name }) + : undefined; + + trace + .getActiveSpan() + ?.setAttribute('mcp.describe_table.table_count', page.length); + + return jsonResult({ + database: input.database, + tables: page.map(shapeTableListEntry), + ...(nextCursor && { nextCursor }), + nextSteps: { + describe: + 'Call clickstack_describe_table again with table set to get its columns, keys, and indices.', + }, + }); +} + +async function describeTable( + ctx: TableQueryContext, + input: DescribeTableInput & { table: string }, +): Promise { + const { database, table } = input; + let details; + try { + details = await fetchTableDetails(ctx, { database, table }); + } catch (e) { + return tableQueryErrorResult(e, 'Failed to describe table'); + } + + if (!details.table) { + return mcpUserError( + `Table "${database}.${table}" not found or not visible to this connection's ClickHouse user. ` + + 'Call clickstack_describe_table with only database to list its tables.', + ); + } + + const description = shapeTableDescription({ + database, + table: details.table, + columns: details.columns, + indices: details.indices, + }); + + let sample: Record = {}; + if (input.sampleRows > 0) { + try { + const rows = await fetchSampleRows(ctx, { + database, + table, + limit: Math.min(input.sampleRows, MAX_SAMPLE_ROWS), + }); + sample = { + sampleRows: trimToolResponse(rows, SAMPLE_ROWS_MAX_CHARS).data, + }; + } catch (e) { + // The schema is still useful when the user can read system tables but + // not the table itself. + sample = { sampleRowsError: sanitizeFetchError(e) }; + } + } + + trace.getActiveSpan()?.setAttributes({ + 'mcp.describe_table.column_count': description.columns.length, + 'mcp.describe_table.index_count': description.dataSkippingIndices.length, + 'mcp.describe_table.engine': description.table.engine, + }); + + const { data, isTrimmed } = trimToolResponse({ + ...description, + ...sample, + nextSteps: { + query: + `Query it with clickstack_sql using connectionId "${ctx.connectionId}" and ` + + `FROM ${quoteName(database)}.${quoteName(table)}. ` + + 'To query it with the builder tools, register it as a source with clickstack_save_source.', + }, + }); + + return jsonResult( + isTrimmed + ? { + ...data, + note: 'Result was trimmed for context size. Some columns or sample rows may be omitted.', + } + : data, + ); +} + +export function registerDescribeTable({ + context, + registerTool, +}: ToolRegistrar): void { + const teamId = context.teamId.toString(); + + registerTool( + 'clickstack_describe_table', + { + title: 'Describe ClickHouse Table', + annotations: READ_ONLY_TOOL, + description: + 'Inspect ClickHouse tables that are NOT registered as sources. For a source, ' + + 'call clickstack_describe_source instead.\n\n' + + 'Use this instead of querying system.tables / system.columns with clickstack_sql.\n\n' + + 'With table set, returns:\n' + + '- table: engine, sorting key, primary key, partition key, approximate total rows and bytes\n' + + '- columns[]: name, type, default kind/expression, comment, and whether it is in the sorting/primary key\n' + + '- dataSkippingIndices[]: name, type, expression\n' + + `- sampleRows: up to ${MAX_SAMPLE_ROWS} rows when sampleRows is set (off by default)\n\n` + + 'Without table, lists the tables in the database (name, engine, approximate rows), ' + + 'paginated. Pass the returned `nextCursor` back unchanged to fetch the next page.', + inputSchema: describeTableSchema, + }, + async rawInput => { + const input = describeTableSchema.parse(rawInput); + + const invalid = validateDescribeTableInput(input); + if (invalid) return mcpUserError(invalid); + const invalidId = validateObjectId(input.connectionId, 'connectionId'); + if (invalidId) return invalidId; + + trace.getActiveSpan()?.setAttributes({ + 'mcp.describe_table.mode': input.table ? 'describe' : 'list', + 'mcp.describe_table.connection_id': input.connectionId, + 'mcp.describe_table.database': input.database, + 'mcp.describe_table.sample_rows': input.sampleRows, + ...(input.table && { 'mcp.describe_table.table': input.table }), + }); + + const connection = await getConnectionById( + teamId, + input.connectionId, + true, + ); + if (!connection) { + return mcpUserError( + `Connection "${input.connectionId}" not found. Call clickstack_list_sources to see available connections.`, + ); + } + + const clickhouseClient = new ClickhouseClient({ + host: connection.host, + username: connection.username, + password: connection.password, + requestTimeout: MCP_REQUEST_TIMEOUT, + }); + + const outcome = await runWithTimeout( + signal => { + const ctx = { + clickhouseClient, + connectionId: input.connectionId, + signal, + }; + return input.table + ? describeTable(ctx, { ...input, table: input.table }) + : listTables(ctx, input); + }, + { timeoutMs: MCP_TOOL_TIMEOUT_MS }, + ); + if (!outcome.timedOut) return outcome.value; + + logger.warn( + { teamId, connectionId: input.connectionId, database: input.database }, + 'clickstack_describe_table timed out', + ); + return mcpServerError( + 'clickstack_describe_table timed out. The ClickHouse server may be under load; try again.', + ); + }, + ); +} diff --git a/packages/api/src/mcp/tools/sources/describeTableQueries.ts b/packages/api/src/mcp/tools/sources/describeTableQueries.ts new file mode 100644 index 0000000000..a49b3f9d19 --- /dev/null +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -0,0 +1,198 @@ +import { + type ChSql, + chSql, + tableExpr, +} from '@hyperdx/common-utils/dist/clickhouse'; + +import type { ClickhouseClient } from '@/clickhouse'; +import { MCP_CLICKHOUSE_SETTINGS } from '@/mcp/tools/query/helpers'; + +export type TableQueryContext = { + clickhouseClient: Pick; + connectionId: string; + signal: AbortSignal; +}; + +async function queryRows( + { clickhouseClient, connectionId, signal }: TableQueryContext, + sql: ChSql, +): Promise { + const response = await clickhouseClient.query<'JSON'>({ + query: sql.sql, + query_params: sql.params, + format: 'JSON', + connectionId, + clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + abort_signal: signal, + }); + const result = (await response.json()) as { data: T[] }; + return result.data; +} + +/** ClickHouse returns 64-bit integers as strings, and NULL for views. */ +function toApproxCount(value: string | number | null): number | null { + if (value == null) return null; + const n = Number(value); + return Number.isFinite(n) ? n : null; +} + +export type SystemTableRow = { + name: string; + engine: string; + total_rows: string | number | null; +}; + +export type SystemTableDetailRow = SystemTableRow & { + sorting_key: string; + primary_key: string; + partition_key: string; + comment: string; + total_bytes: string | number | null; +}; + +export type SystemColumnRow = { + name: string; + type: string; + default_kind: string; + default_expression: string; + comment: string; + is_in_sorting_key: number; + is_in_primary_key: number; +}; + +export type SystemIndexRow = { + name: string; + type: string; + expr: string; +}; + +export function shapeTableListEntry(row: SystemTableRow) { + return { + name: row.name, + engine: row.engine, + approxRows: toApproxCount(row.total_rows), + }; +} + +export function shapeTableDescription({ + database, + table, + columns, + indices, +}: { + database: string; + table: SystemTableDetailRow; + columns: SystemColumnRow[]; + indices: SystemIndexRow[]; +}) { + return { + table: { + database, + name: table.name, + engine: table.engine, + sortingKey: table.sorting_key, + primaryKey: table.primary_key, + partitionKey: table.partition_key, + ...(table.comment && { comment: table.comment }), + approxTotalRows: toApproxCount(table.total_rows), + approxTotalBytes: toApproxCount(table.total_bytes), + }, + columns: columns.map(c => ({ + name: c.name, + type: c.type, + ...(c.default_kind && { + defaultKind: c.default_kind, + defaultExpression: c.default_expression, + }), + ...(c.comment && { comment: c.comment }), + inSortingKey: Boolean(Number(c.is_in_sorting_key)), + inPrimaryKey: Boolean(Number(c.is_in_primary_key)), + })), + dataSkippingIndices: indices.map(i => ({ + name: i.name, + type: i.type, + expression: i.expr, + })), + }; +} + +/** One extra row past `limit` signals that another page exists. */ +export function fetchTablePage( + ctx: TableQueryContext, + { + database, + afterName, + limit, + }: { database: string; afterName: string | undefined; limit: number }, +) { + return queryRows( + ctx, + chSql` + SELECT name, engine, total_rows + FROM system.tables + WHERE database = ${{ String: database }} + AND NOT is_temporary + ${afterName !== undefined ? chSql`AND name > ${{ String: afterName }}` : ''} + ORDER BY name ASC + LIMIT ${{ Int32: limit + 1 }} + `, + ); +} + +export async function databaseExists(ctx: TableQueryContext, database: string) { + const rows = await queryRows<{ name: string }>( + ctx, + chSql`SELECT name FROM system.databases WHERE name = ${{ String: database }}`, + ); + return rows.length > 0; +} + +export async function fetchTableDetails( + ctx: TableQueryContext, + { database, table }: { database: string; table: string }, +) { + const [tables, columns, indices] = await Promise.all([ + queryRows( + ctx, + chSql` + SELECT name, engine, sorting_key, primary_key, partition_key, comment, + total_rows, total_bytes + FROM system.tables + WHERE database = ${{ String: database }} AND name = ${{ String: table }} + `, + ), + queryRows( + ctx, + chSql` + SELECT name, type, default_kind, default_expression, comment, + is_in_sorting_key, is_in_primary_key + FROM system.columns + WHERE database = ${{ String: database }} AND table = ${{ String: table }} + ORDER BY position + `, + ), + queryRows( + ctx, + chSql` + SELECT name, type, expr + FROM system.data_skipping_indices + WHERE database = ${{ String: database }} AND table = ${{ String: table }} + `, + ), + ]); + return { table: tables[0], columns, indices }; +} + +export function fetchSampleRows( + ctx: TableQueryContext, + { + database, + table, + limit, + }: { database: string; table: string; limit: number }, +) { + return queryRows>( + ctx, + chSql`SELECT * FROM ${tableExpr({ database, table })} LIMIT ${{ Int32: limit }}`, + ); +} diff --git a/packages/api/src/mcp/tools/sources/describeTableSchema.ts b/packages/api/src/mcp/tools/sources/describeTableSchema.ts new file mode 100644 index 0000000000..4e70b85cfc --- /dev/null +++ b/packages/api/src/mcp/tools/sources/describeTableSchema.ts @@ -0,0 +1,83 @@ +import { z } from 'zod'; + +import { + decodeCursor as decodeCursorPayload, + encodeCursor as encodeCursorPayload, +} from '@/utils/pagination'; + +const DEFAULT_TABLE_LIST_LIMIT = 50; +const MAX_TABLE_LIST_LIMIT = 500; +export const MAX_SAMPLE_ROWS = 5; + +const tableListCursorSchema = z.object({ lastName: z.string() }); + +export type TableListCursor = z.infer; + +export function encodeTableListCursor(payload: TableListCursor): string { + return encodeCursorPayload(payload); +} + +export function decodeTableListCursor(raw: string): TableListCursor | null { + return decodeCursorPayload(raw, tableListCursorSchema); +} + +export const describeTableSchema = z.object({ + connectionId: z + .string() + .describe( + 'Connection to inspect. Call clickstack_list_sources to find connection IDs.', + ), + database: z.string().min(1).describe('ClickHouse database name.'), + table: z + .string() + .min(1) + .optional() + .describe( + 'Table to describe. Omit to list the tables in the database instead.', + ), + sampleRows: z + .number() + .int() + .min(0) + .max(MAX_SAMPLE_ROWS) + .optional() + .default(0) + .describe( + `Number of rows to sample from the table (0-${MAX_SAMPLE_ROWS}). Only valid with table. Default: 0.`, + ), + limit: z + .number() + .int() + .min(1) + .max(MAX_TABLE_LIST_LIMIT) + .optional() + .default(DEFAULT_TABLE_LIST_LIMIT) + .describe( + 'Maximum tables returned per page when listing. Use cursor to fetch further pages.', + ), + cursor: z + .string() + .optional() + .describe( + 'Opaque pagination cursor returned by a previous listing call as `nextCursor`. ' + + 'Pass it back unchanged to get the next page.', + ), +}); + +export type DescribeTableInput = z.infer; + +/** + * Checks the option combinations a flat Zod object cannot express. Returns an + * error message, or null when the input is consistent. + */ +export function validateDescribeTableInput( + input: Pick, +): string | null { + if (input.table == null && input.sampleRows > 0) { + return 'sampleRows requires table. Pass a table name, or omit sampleRows to list tables.'; + } + if (input.table != null && input.cursor) { + return 'cursor only applies when listing tables. Omit table to page through the table list, or omit cursor to describe the table.'; + } + return null; +} diff --git a/packages/api/src/mcp/tools/sources/index.ts b/packages/api/src/mcp/tools/sources/index.ts index 2efea649a1..2c36fc9d0d 100644 --- a/packages/api/src/mcp/tools/sources/index.ts +++ b/packages/api/src/mcp/tools/sources/index.ts @@ -3,6 +3,7 @@ import type { ToolDefinition, ToolRegistrar } from '@/mcp/tools/types'; import { registerDeleteSource } from './deleteSource'; import { registerDescribeMetric } from './describeMetric'; import { registerDescribeSource } from './describeSource'; +import { registerDescribeTable } from './describeTable'; import { registerListMetrics } from './listMetrics'; import { registerListSources } from './listSources'; import { registerSaveSource } from './saveSource'; @@ -10,6 +11,7 @@ import { registerSaveSource } from './saveSource'; const sourcesTools: ToolDefinition = (registrar: ToolRegistrar) => { registerListSources(registrar); registerDescribeSource(registrar); + registerDescribeTable(registrar); registerSaveSource(registrar); registerDeleteSource(registrar); registerListMetrics(registrar); From 21370ec05c27b3565d992217b9866ae7c66202cb Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 15:29:41 -0600 Subject: [PATCH 2/7] fix(mcp): bound describe_table samples and resolve Distributed tables (review) - Trim each sample row to its share of the sample budget and cap the bytes ClickHouse returns, so a few wide rows cannot push columns out. - Read keys and skip indices through Metadata so Distributed tables report their local table's keys, and take key-column flags from the local table. - Surface infra failures while sampling as server errors; add a hint for user-class sample failures. - Check the database exists on any empty page, including cursor replays. - Use approxTotalRows in both modes and quoteIdentifierIfNeeded for the hint. --- .../mcp/__tests__/describeTable.int.test.ts | 4 +- .../src/mcp/__tests__/describeTable.test.ts | 187 +++++++++++++++++- .../src/mcp/tools/sources/describeTable.ts | 106 +++++++--- .../mcp/tools/sources/describeTableQueries.ts | 142 +++++++++---- 4 files changed, 355 insertions(+), 84 deletions(-) diff --git a/packages/api/src/mcp/__tests__/describeTable.int.test.ts b/packages/api/src/mcp/__tests__/describeTable.int.test.ts index 864a97986d..4ade69dbfa 100644 --- a/packages/api/src/mcp/__tests__/describeTable.int.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.int.test.ts @@ -165,7 +165,7 @@ describe('clickstack_describe_table', () => { getFirstText(await call({ database: DATABASE, limit: 1 })), ); expect(first.tables).toEqual([ - { name: 'events', engine: 'MergeTree', approxRows: 3 }, + { name: 'events', engine: 'MergeTree', approxTotalRows: 3 }, ]); expect(first.nextCursor).toEqual(expect.any(String)); @@ -175,7 +175,7 @@ describe('clickstack_describe_table', () => { ), ); expect(second.tables).toEqual([ - { name: 'events_view', engine: 'View', approxRows: null }, + { name: 'events_view', engine: 'View', approxTotalRows: null }, ]); expect(second.nextCursor).toBeUndefined(); }); diff --git a/packages/api/src/mcp/__tests__/describeTable.test.ts b/packages/api/src/mcp/__tests__/describeTable.test.ts index bc725ab2f9..42a826966c 100644 --- a/packages/api/src/mcp/__tests__/describeTable.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -1,11 +1,20 @@ const mockQuery = jest.fn(); const mockGetConnectionById = jest.fn(); +const mockGetTableMetadata = jest.fn(); +const mockGetSkipIndices = jest.fn(); jest.mock('@/models/source', () => ({})); jest.mock('@/controllers/sources', () => ({})); jest.mock('@/controllers/connection', () => ({ getConnectionById: (...args: unknown[]) => mockGetConnectionById(...args), })); +jest.mock('@hyperdx/common-utils/dist/core/metadata', () => ({ + ...jest.requireActual('@hyperdx/common-utils/dist/core/metadata'), + getMetadata: () => ({ + getTableMetadata: mockGetTableMetadata, + getSkipIndices: mockGetSkipIndices, + }), +})); jest.mock('@/clickhouse', () => ({ ClickhouseClient: jest.fn(() => ({ query: mockQuery })), })); @@ -27,7 +36,11 @@ import { getErrorCategory, type McpErrorResult } from '@/mcp/utils/errors'; const CONNECTION_ID = '65f5e4a3b9e77c001a567890'; -type QueryArgs = { query: string; query_params: Record }; +type QueryArgs = { + query: string; + query_params: Record; + clickhouse_settings: Record; +}; const queryCalls = (): QueryArgs[] => mockQuery.mock.calls.map(([args]: [QueryArgs]) => args); @@ -106,6 +119,8 @@ describe('clickstack_describe_table', () => { username: 'default', password: '', }); + mockGetTableMetadata.mockReset().mockResolvedValue(TABLE_ROW); + mockGetSkipIndices.mockReset().mockResolvedValue([]); }); describe('input validation', () => { @@ -177,8 +192,9 @@ describe('clickstack_describe_table', () => { shapeTableDescription({ database: 'default', table: { ...TABLE_ROW, total_rows: null, total_bytes: null }, + keys: TABLE_ROW, columns: COLUMN_ROWS, - indices: [{ name: 'idx_day', type: 'minmax', expr: 'Day' }], + indices: [{ name: 'idx_day', type: 'minmax', expression: 'Day' }], }), ).toEqual({ table: { @@ -221,7 +237,7 @@ describe('clickstack_describe_table', () => { engine: 'MergeTree', total_rows: '42', }), - ).toEqual({ name: 'events', engine: 'MergeTree', approxRows: 42 }); + ).toEqual({ name: 'events', engine: 'MergeTree', approxTotalRows: 42 }); }); it('binds database and table names as query parameters', async () => { @@ -230,11 +246,16 @@ describe('clickstack_describe_table', () => { connectionId: CONNECTION_ID, database: "evil'db", table: 'events`; DROP TABLE x', + sampleRows: 1, }); - for (const { query: sql, query_params } of queryCalls()) { + const calls = queryCalls(); + expect(calls.some(c => c.query.includes(':Identifier}'))).toBe(true); + for (const { query: sql, query_params } of calls) { expect(sql).not.toContain('evil'); expect(sql).not.toContain('DROP'); - expect(Object.values(query_params)).toContain("evil'db"); + expect(Object.values(query_params)).toEqual( + expect.arrayContaining(["evil'db", 'events`; DROP TABLE x']), + ); } }); @@ -254,8 +275,8 @@ describe('clickstack_describe_table', () => { }), ); expect(output.tables).toEqual([ - { name: 'a', engine: 'MergeTree', approxRows: 1 }, - { name: 'b', engine: 'View', approxRows: null }, + { name: 'a', engine: 'MergeTree', approxTotalRows: 1 }, + { name: 'b', engine: 'View', approxTotalRows: null }, ]); expect(output.nextCursor).toBe(encodeTableListCursor({ lastName: 'b' })); }); @@ -295,6 +316,7 @@ describe('clickstack_describe_table', () => { it('maps permission errors to user errors with a hint', async () => { respondWith({ + 'FROM system.tables': [TABLE_ROW], 'FROM system.columns': new Error( 'Code: 497. DB::Exception: default: Not enough privileges. (ACCESS_DENIED)', ), @@ -319,13 +341,15 @@ describe('clickstack_describe_table', () => { }), ); expect(output.sampleRows).toBeUndefined(); - expect(queryCalls().some(q => q.query.includes('SELECT *'))).toBe(false); + expect(queryCalls().some(q => q.query.includes(':Identifier}'))).toBe( + false, + ); }); it('trims oversized sample rows', async () => { respondWith({ 'FROM system.tables': [TABLE_ROW], - 'SELECT *': [{ Body: 'x'.repeat(50_000) }, { Body: 'short' }], + ':Identifier}': [{ Body: 'x'.repeat(50_000) }, { Body: 'short' }], }); const output = parseText( await getHandler()({ @@ -343,7 +367,7 @@ describe('clickstack_describe_table', () => { it('keeps the schema when sampling fails', async () => { respondWith({ 'FROM system.tables': [TABLE_ROW], - 'SELECT *': new Error('ACCESS_DENIED'), + ':Identifier}': new Error('ACCESS_DENIED'), }); const output = parseText( await getHandler()({ @@ -355,6 +379,149 @@ describe('clickstack_describe_table', () => { ); expect(output.table.name).toBe('events'); expect(output.sampleRowsError).toBe('ACCESS_DENIED'); + expect(output.sampleRowsHint).toMatch(/schema above is complete/); + }); + + it('lists an existing empty database without an error', async () => { + respondWith({ 'FROM system.databases': [{ name: 'empty' }] }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'empty', + }); + expect(result.isError).toBeFalsy(); + expect(parseText(result).tables).toEqual([]); + }); + + it('reports a missing database when a cursor is replayed against it', async () => { + respondWith({}); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'dropped', + cursor: encodeTableListCursor({ lastName: 'a' }), + }); + expect(result.isError).toBe(true); + expect(result.content[0].text).toMatch(/Database "dropped" not found/); + }); + + it('reads keys, indices, and key columns from the local table of a Distributed table', async () => { + respondWith({ + 'FROM system.tables': [ + { + ...TABLE_ROW, + engine: 'Distributed', + engine_full: + "Distributed('prod', 'default', 'events_local', rand())", + sorting_key: '', + primary_key: '', + partition_key: '', + total_rows: null, + }, + ], + 'cluster(': [ + { name: 'Timestamp', is_in_sorting_key: 1, is_in_primary_key: 1 }, + ], + 'FROM system.columns': COLUMN_ROWS.map(c => ({ + ...c, + is_in_sorting_key: 0, + is_in_primary_key: 0, + })), + }); + mockGetTableMetadata.mockResolvedValue({ + ...TABLE_ROW, + engine: 'ReplicatedMergeTree', + }); + mockGetSkipIndices.mockResolvedValue([ + { + name: 'idx', + type: 'minmax', + typeFull: 'minmax', + expression: 'Day', + granularity: 1, + }, + ]); + + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }), + ); + expect(output.table).toMatchObject({ + engine: 'Distributed', + localTable: { + cluster: 'prod', + database: 'default', + table: 'events_local', + engine: 'ReplicatedMergeTree', + }, + sortingKey: 'ServiceName, Timestamp', + primaryKey: 'ServiceName', + }); + expect(output.columns[0]).toMatchObject({ + name: 'Timestamp', + inSortingKey: true, + inPrimaryKey: true, + }); + expect(output.dataSkippingIndices).toEqual([ + { name: 'idx', type: 'minmax', expression: 'Day' }, + ]); + }); + + it('keeps several medium sample rows within the sample budget', async () => { + const columns = Array.from({ length: 200 }, (_, i) => ({ + ...COLUMN_ROWS[0], + name: `Column_${i}`, + })); + respondWith({ + 'FROM system.tables': [TABLE_ROW], + 'FROM system.columns': columns, + ':Identifier}': Array.from({ length: 5 }, () => ({ + Body: 'x'.repeat(9_900), + })), + }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 5, + }), + ); + expect(JSON.stringify(output.sampleRows).length).toBeLessThan(11_000); + expect(output.columns).toHaveLength(200); + expect(output.note).toMatch(/trimmed/); + }); + + it('caps the bytes ClickHouse returns for the sample', async () => { + respondWith({ 'FROM system.tables': [TABLE_ROW] }); + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 1, + }); + const sampleCall = queryCalls().find(c => + c.query.includes(':Identifier}'), + ); + expect(sampleCall?.clickhouse_settings.max_result_bytes).toBeDefined(); + }); + + it('reports an infrastructure failure while sampling as a server error', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + ':Identifier}': Object.assign(new Error('connect ECONNREFUSED'), { + code: 'ECONNREFUSED', + }), + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 1, + }); + expect(result.isError).toBe(true); + expect(errorCategory(result)).toBe('server'); }); }); }); diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts index 7232296c89..2c3d05e54c 100644 --- a/packages/api/src/mcp/tools/sources/describeTable.ts +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -1,3 +1,8 @@ +import { + getMetadata, + type Metadata, + quoteIdentifierIfNeeded, +} from '@hyperdx/common-utils/dist/core/metadata'; import { trace } from '@opentelemetry/api'; import { ClickhouseClient } from '@/clickhouse'; @@ -5,6 +10,7 @@ import { getConnectionById } from '@/controllers/connection'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import { clickHouseErrorResult, + isServerError, MCP_REQUEST_TIMEOUT, } from '@/mcp/tools/query/helpers'; import type { ToolRegistrar, ToolResult } from '@/mcp/tools/types'; @@ -39,8 +45,7 @@ import { // Sample rows can hold arbitrarily large values; keep them a small share of // the response so they never crowd out the schema. const SAMPLE_ROWS_MAX_CHARS = 10_000; - -const quoteName = (name: string) => `\`${name.replace(/`/g, '\\`')}\``; +const RESPONSE_MAX_CHARS = 50_000; function jsonResult(payload: unknown): ToolResult { return { content: [{ type: 'text', text: JSON.stringify(payload) }] }; @@ -88,7 +93,7 @@ async function listTables( } // system.tables returns nothing for a missing database rather than erroring. - if (rows.length === 0 && !cursor) { + if (rows.length === 0) { try { if (!(await databaseExists(ctx, input.database))) { return mcpUserError( @@ -121,75 +126,111 @@ async function listTables( }); } +function trimSampleRows(rows: Record[]) { + const perRowBudget = Math.floor( + SAMPLE_ROWS_MAX_CHARS / Math.max(rows.length, 1), + ); + let isTrimmed = false; + const data = rows.map(row => { + const trimmed = trimToolResponse(row, perRowBudget); + isTrimmed ||= trimmed.isTrimmed; + return trimmed.data; + }); + return { data, isTrimmed }; +} + +async function sampleTableRows( + ctx: TableQueryContext, + args: { database: string; table: string; limit: number }, +) { + try { + const rows = trimSampleRows(await fetchSampleRows(ctx, args)); + return { sampleRows: rows.data, isTrimmed: rows.isTrimmed }; + } catch (e) { + if (isServerError(e)) throw e; + // The schema is still useful when the user can read system tables but + // not the table itself. + return { + sampleRowsError: sanitizeFetchError(e), + sampleRowsHint: + 'Only the row sample failed; the schema above is complete. Retry without sampleRows, ' + + "or check that the connection's ClickHouse user can SELECT from this table.", + isTrimmed: false, + }; + } +} + async function describeTable( ctx: TableQueryContext, + metadata: Metadata, input: DescribeTableInput & { table: string }, ): Promise { const { database, table } = input; let details; try { - details = await fetchTableDetails(ctx, { database, table }); + details = await fetchTableDetails(ctx, metadata, { database, table }); } catch (e) { return tableQueryErrorResult(e, 'Failed to describe table'); } - if (!details.table) { + if (!details) { return mcpUserError( `Table "${database}.${table}" not found or not visible to this connection's ClickHouse user. ` + 'Call clickstack_describe_table with only database to list its tables.', ); } - const description = shapeTableDescription({ + const { columns, ...description } = shapeTableDescription({ database, - table: details.table, - columns: details.columns, - indices: details.indices, + ...details, }); - let sample: Record = {}; + let sample = {}; + let sampleTrimmed = false; if (input.sampleRows > 0) { try { - const rows = await fetchSampleRows(ctx, { + const { isTrimmed, ...result } = await sampleTableRows(ctx, { database, table, - limit: Math.min(input.sampleRows, MAX_SAMPLE_ROWS), + limit: input.sampleRows, }); - sample = { - sampleRows: trimToolResponse(rows, SAMPLE_ROWS_MAX_CHARS).data, - }; + sample = result; + sampleTrimmed = isTrimmed; } catch (e) { - // The schema is still useful when the user can read system tables but - // not the table itself. - sample = { sampleRowsError: sanitizeFetchError(e) }; + return clickHouseErrorResult(e, 'Failed to sample rows'); } } trace.getActiveSpan()?.setAttributes({ - 'mcp.describe_table.column_count': description.columns.length, + 'mcp.describe_table.column_count': columns.length, 'mcp.describe_table.index_count': description.dataSkippingIndices.length, 'mcp.describe_table.engine': description.table.engine, }); - const { data, isTrimmed } = trimToolResponse({ + const rest = { ...description, ...sample, nextSteps: { query: `Query it with clickstack_sql using connectionId "${ctx.connectionId}" and ` + - `FROM ${quoteName(database)}.${quoteName(table)}. ` + + `FROM ${quoteIdentifierIfNeeded(database)}.${quoteIdentifierIfNeeded(table)}. ` + 'To query it with the builder tools, register it as a source with clickstack_save_source.', }, - }); - - return jsonResult( - isTrimmed - ? { - ...data, - note: 'Result was trimmed for context size. Some columns or sample rows may be omitted.', - } - : data, + }; + // Trim only the column list, against what the rest of the response leaves, + // so a large sample cannot push columns out. + const trimmedColumns = trimToolResponse( + columns, + RESPONSE_MAX_CHARS - JSON.stringify(rest).length, ); + + return jsonResult({ + ...rest, + columns: trimmedColumns.data, + ...((sampleTrimmed || trimmedColumns.isTrimmed) && { + note: 'Result was trimmed for context size. Some columns or sample values may be omitted.', + }), + }); } export function registerDescribeTable({ @@ -258,7 +299,10 @@ export function registerDescribeTable({ signal, }; return input.table - ? describeTable(ctx, { ...input, table: input.table }) + ? describeTable(ctx, getMetadata(clickhouseClient), { + ...input, + table: input.table, + }) : listTables(ctx, input); }, { timeoutMs: MCP_TOOL_TIMEOUT_MS }, diff --git a/packages/api/src/mcp/tools/sources/describeTableQueries.ts b/packages/api/src/mcp/tools/sources/describeTableQueries.ts index a49b3f9d19..0e6c77d2b5 100644 --- a/packages/api/src/mcp/tools/sources/describeTableQueries.ts +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -1,8 +1,15 @@ +import type { ClickHouseSettings } from '@clickhouse/client-common'; import { type ChSql, chSql, tableExpr, } from '@hyperdx/common-utils/dist/clickhouse'; +import type { + Metadata, + SkipIndexMetadata, + TableMetadata, +} from '@hyperdx/common-utils/dist/core/metadata'; +import { getDistributedTableArgs } from '@hyperdx/common-utils/dist/core/utils'; import type { ClickhouseClient } from '@/clickhouse'; import { MCP_CLICKHOUSE_SETTINGS } from '@/mcp/tools/query/helpers'; @@ -16,13 +23,14 @@ export type TableQueryContext = { async function queryRows( { clickhouseClient, connectionId, signal }: TableQueryContext, sql: ChSql, + settings?: ClickHouseSettings, ): Promise { const response = await clickhouseClient.query<'JSON'>({ query: sql.sql, query_params: sql.params, format: 'JSON', connectionId, - clickhouse_settings: MCP_CLICKHOUSE_SETTINGS, + clickhouse_settings: { ...MCP_CLICKHOUSE_SETTINGS, ...settings }, abort_signal: signal, }); const result = (await response.json()) as { data: T[] }; @@ -42,57 +50,70 @@ export type SystemTableRow = { total_rows: string | number | null; }; -export type SystemTableDetailRow = SystemTableRow & { - sorting_key: string; - primary_key: string; - partition_key: string; - comment: string; - total_bytes: string | number | null; +// total_rows / total_bytes are NULL for views, unlike TableMetadata's typing. +type TableInfo = Pick & { + total_rows: string | null; + total_bytes: string | null; }; -export type SystemColumnRow = { +type TableKeys = Pick< + TableMetadata, + 'engine' | 'sorting_key' | 'primary_key' | 'partition_key' +>; + +type ColumnKeyFlags = { name: string; - type: string; - default_kind: string; - default_expression: string; - comment: string; is_in_sorting_key: number; is_in_primary_key: number; }; -export type SystemIndexRow = { - name: string; +export type SystemColumnRow = ColumnKeyFlags & { type: string; - expr: string; + default_kind: string; + default_expression: string; + comment: string; }; +type DistributedTarget = { cluster: string; database: string; table: string }; + export function shapeTableListEntry(row: SystemTableRow) { return { name: row.name, engine: row.engine, - approxRows: toApproxCount(row.total_rows), + approxTotalRows: toApproxCount(row.total_rows), }; } +/** + * `keys` comes from the local table when `table` is Distributed, since a + * Distributed table has no keys or indices of its own. + */ export function shapeTableDescription({ database, table, + keys, + distributedTarget, columns, indices, }: { database: string; - table: SystemTableDetailRow; + table: TableInfo; + keys: TableKeys; + distributedTarget?: DistributedTarget; columns: SystemColumnRow[]; - indices: SystemIndexRow[]; + indices: Pick[]; }) { return { table: { database, name: table.name, engine: table.engine, - sortingKey: table.sorting_key, - primaryKey: table.primary_key, - partitionKey: table.partition_key, + ...(distributedTarget && { + localTable: { ...distributedTarget, engine: keys.engine }, + }), + sortingKey: keys.sorting_key, + primaryKey: keys.primary_key, + partitionKey: keys.partition_key, ...(table.comment && { comment: table.comment }), approxTotalRows: toApproxCount(table.total_rows), approxTotalBytes: toApproxCount(table.total_bytes), @@ -111,7 +132,7 @@ export function shapeTableDescription({ dataSkippingIndices: indices.map(i => ({ name: i.name, type: i.type, - expression: i.expr, + expression: i.expression, })), }; } @@ -147,20 +168,33 @@ export async function databaseExists(ctx: TableQueryContext, database: string) { return rows.length > 0; } +/** Returns undefined when the table does not exist or is not visible. */ export async function fetchTableDetails( ctx: TableQueryContext, + metadata: Pick, { database, table }: { database: string; table: string }, ) { - const [tables, columns, indices] = await Promise.all([ - queryRows( - ctx, - chSql` - SELECT name, engine, sorting_key, primary_key, partition_key, comment, - total_rows, total_bytes - FROM system.tables - WHERE database = ${{ String: database }} AND name = ${{ String: table }} - `, - ), + const rows = await queryRows( + ctx, + chSql` + SELECT * FROM system.tables + WHERE database = ${{ String: database }} AND name = ${{ String: table }} + LIMIT 1 + `, + ); + const info = rows.at(0); + if (!info) return undefined; + + const distributedTarget = + info.engine === 'Distributed' ? getDistributedTableArgs(info) : undefined; + const metadataArgs = { + databaseName: database, + tableName: table, + connectionId: ctx.connectionId, + }; + const [keys, indices, columns, localKeyFlags] = await Promise.all([ + metadata.getTableMetadata(metadataArgs), + metadata.getSkipIndices(metadataArgs), queryRows( ctx, chSql` @@ -171,18 +205,43 @@ export async function fetchTableDetails( ORDER BY position `, ), - queryRows( - ctx, - chSql` - SELECT name, type, expr - FROM system.data_skipping_indices - WHERE database = ${{ String: database }} AND table = ${{ String: table }} - `, - ), + distributedTarget + ? queryRows( + ctx, + chSql` + SELECT name, is_in_sorting_key, is_in_primary_key + FROM cluster(${{ String: distributedTarget.cluster }}, system.columns) + WHERE database = ${{ String: distributedTarget.database }} + AND table = ${{ String: distributedTarget.table }} + LIMIT 1 BY name + `, + ) + : undefined, ]); - return { table: tables[0], columns, indices }; + + const flagsByName = new Map(localKeyFlags?.map(f => [f.name, f])); + return { + table: info, + keys: keys ?? info, + distributedTarget, + indices, + columns: columns.map(c => { + const flags = flagsByName.get(c.name); + return flags + ? { + ...c, + is_in_sorting_key: flags.is_in_sorting_key, + is_in_primary_key: flags.is_in_primary_key, + } + : c; + }), + }; } +// Five rows of wide String columns can still be large; cap what ClickHouse +// sends back so the response is bounded before it is buffered. +const SAMPLE_MAX_RESULT_BYTES = 1_000_000; + export function fetchSampleRows( ctx: TableQueryContext, { @@ -194,5 +253,6 @@ export function fetchSampleRows( return queryRows>( ctx, chSql`SELECT * FROM ${tableExpr({ database, table })} LIMIT ${{ Int32: limit }}`, + { max_result_bytes: SAMPLE_MAX_RESULT_BYTES.toString() }, ); } From 7a394447dcc9dda463e89686551ea211c4189cc9 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 15:44:15 -0600 Subject: [PATCH 3/7] fix(mcp): read describe_table keys and indices uncached (review) The Metadata cache has no TTL for table metadata or skip indices, so an ALTER showed new columns next to stale keys and indices, and those lookups skipped the MCP readonly settings and abort signal. Read the keys and indices with the tool's own queries, following Distributed tables to the local table via cluster(). A missing grant on data_skipping_indices now adds a hint instead of failing the describe. Add integration coverage for a Distributed table and for an index added after a describe. --- .../mcp/__tests__/describeTable.int.test.ts | 55 +++++++++++++ .../src/mcp/__tests__/describeTable.test.ts | 76 ++++++++++++------ .../src/mcp/tools/sources/describeTable.ts | 21 +++-- .../mcp/tools/sources/describeTableQueries.ts | 78 ++++++++++++++----- 4 files changed, 175 insertions(+), 55 deletions(-) diff --git a/packages/api/src/mcp/__tests__/describeTable.int.test.ts b/packages/api/src/mcp/__tests__/describeTable.int.test.ts index 4ade69dbfa..3c67659f89 100644 --- a/packages/api/src/mcp/__tests__/describeTable.int.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.int.test.ts @@ -13,6 +13,8 @@ import Connection from '@/models/connection'; import { callTool, createTestClient, getFirstText } from './mcpTestUtils'; const DATABASE = 'mcp_describe_table_test'; +// Kept apart so the Distributed table does not change DATABASE's listing. +const DIST_DATABASE = 'mcp_describe_table_dist_test'; describe('clickstack_describe_table', () => { const server = getServer(); @@ -23,7 +25,9 @@ describe('clickstack_describe_table', () => { await server.start(); const ch = await getTestFixtureClickHouseClient(); await ch.command({ query: `DROP DATABASE IF EXISTS ${DATABASE}` }); + await ch.command({ query: `DROP DATABASE IF EXISTS ${DIST_DATABASE}` }); await ch.command({ query: `CREATE DATABASE ${DATABASE}` }); + await ch.command({ query: `CREATE DATABASE ${DIST_DATABASE}` }); await ch.command({ query: ` CREATE TABLE ${DATABASE}.events ( @@ -43,6 +47,9 @@ describe('clickstack_describe_table', () => { await ch.command({ query: `CREATE VIEW ${DATABASE}.events_view AS SELECT ServiceName FROM ${DATABASE}.events`, }); + await ch.command({ + query: `CREATE TABLE ${DIST_DATABASE}.events AS ${DATABASE}.events ENGINE = Distributed(hdx_cluster, ${DATABASE}, events)`, + }); await ch.insert({ table: `${DATABASE}.events`, values: [ @@ -77,6 +84,7 @@ describe('clickstack_describe_table', () => { afterAll(async () => { const ch = await getTestFixtureClickHouseClient(); + await ch.command({ query: `DROP DATABASE IF EXISTS ${DIST_DATABASE}` }); await ch.command({ query: `DROP DATABASE IF EXISTS ${DATABASE}` }); await server.stop(); }); @@ -139,6 +147,53 @@ describe('clickstack_describe_table', () => { expect(output.sampleRows).toBeUndefined(); }); + it('reads keys and indices of a Distributed table from its local table', async () => { + const result = await call({ database: DIST_DATABASE, table: 'events' }); + expect(result.isError).toBeFalsy(); + const output = JSON.parse(getFirstText(result)); + expect(output.table).toMatchObject({ + engine: 'Distributed', + localTable: { + cluster: 'hdx_cluster', + database: DATABASE, + table: 'events', + engine: 'MergeTree', + }, + sortingKey: 'ServiceName, Timestamp', + primaryKey: 'ServiceName', + partitionKey: 'toDate(Timestamp)', + }); + expect( + output.columns.find((c: any) => c.name === 'ServiceName'), + ).toMatchObject({ + inSortingKey: true, + inPrimaryKey: true, + }); + expect(output.dataSkippingIndices).toEqual([ + { name: 'idx_body', type: 'tokenbf_v1', expression: 'Body' }, + ]); + }); + + it('reflects an index added after a previous describe', async () => { + await call({ database: DATABASE, table: 'events' }); + const ch = await getTestFixtureClickHouseClient(); + await ch.command({ + query: `ALTER TABLE ${DATABASE}.events ADD INDEX idx_svc ServiceName TYPE set(10) GRANULARITY 1`, + }); + try { + const output = JSON.parse( + getFirstText(await call({ database: DATABASE, table: 'events' })), + ); + expect( + output.dataSkippingIndices.map((i: { name: string }) => i.name), + ).toContain('idx_svc'); + } finally { + await ch.command({ + query: `ALTER TABLE ${DATABASE}.events DROP INDEX idx_svc`, + }); + } + }); + it('returns a capped sample of rows when asked', async () => { const result = await call({ database: DATABASE, diff --git a/packages/api/src/mcp/__tests__/describeTable.test.ts b/packages/api/src/mcp/__tests__/describeTable.test.ts index 42a826966c..af3b101b72 100644 --- a/packages/api/src/mcp/__tests__/describeTable.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -1,20 +1,11 @@ const mockQuery = jest.fn(); const mockGetConnectionById = jest.fn(); -const mockGetTableMetadata = jest.fn(); -const mockGetSkipIndices = jest.fn(); jest.mock('@/models/source', () => ({})); jest.mock('@/controllers/sources', () => ({})); jest.mock('@/controllers/connection', () => ({ getConnectionById: (...args: unknown[]) => mockGetConnectionById(...args), })); -jest.mock('@hyperdx/common-utils/dist/core/metadata', () => ({ - ...jest.requireActual('@hyperdx/common-utils/dist/core/metadata'), - getMetadata: () => ({ - getTableMetadata: mockGetTableMetadata, - getSkipIndices: mockGetSkipIndices, - }), -})); jest.mock('@/clickhouse', () => ({ ClickhouseClient: jest.fn(() => ({ query: mockQuery })), })); @@ -119,8 +110,6 @@ describe('clickstack_describe_table', () => { username: 'default', password: '', }); - mockGetTableMetadata.mockReset().mockResolvedValue(TABLE_ROW); - mockGetSkipIndices.mockReset().mockResolvedValue([]); }); describe('input validation', () => { @@ -417,7 +406,11 @@ describe('clickstack_describe_table', () => { total_rows: null, }, ], - 'cluster(': [ + ', system.tables)': [{ ...TABLE_ROW, engine: 'ReplicatedMergeTree' }], + ', system.data_skipping_indices)': [ + { name: 'idx', type: 'minmax', expression: 'Day' }, + ], + ', system.columns)': [ { name: 'Timestamp', is_in_sorting_key: 1, is_in_primary_key: 1 }, ], 'FROM system.columns': COLUMN_ROWS.map(c => ({ @@ -426,19 +419,6 @@ describe('clickstack_describe_table', () => { is_in_primary_key: 0, })), }); - mockGetTableMetadata.mockResolvedValue({ - ...TABLE_ROW, - engine: 'ReplicatedMergeTree', - }); - mockGetSkipIndices.mockResolvedValue([ - { - name: 'idx', - type: 'minmax', - typeFull: 'minmax', - expression: 'Day', - granularity: 1, - }, - ]); const output = parseText( await getHandler()({ @@ -466,6 +446,52 @@ describe('clickstack_describe_table', () => { expect(output.dataSkippingIndices).toEqual([ { name: 'idx', type: 'minmax', expression: 'Day' }, ]); + const localQueries = queryCalls().filter(c => + c.query.includes('cluster('), + ); + expect(localQueries).toHaveLength(3); + for (const { query_params } of localQueries) { + expect(Object.values(query_params)).toEqual( + expect.arrayContaining(['prod', 'events_local']), + ); + } + }); + + it('still describes the table when skip indices are not readable', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + 'FROM system.data_skipping_indices': new Error( + 'Code: 497. DB::Exception: Not enough privileges. (ACCESS_DENIED)', + ), + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }); + expect(result.isError).toBeFalsy(); + const output = parseText(result); + expect(output.dataSkippingIndices).toEqual([]); + expect(output.dataSkippingIndicesHint).toMatch(/cannot read/); + }); + + it('reports a sample over the byte cap as sampleRowsError', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + ':Identifier}': new Error( + 'Code: 396. DB::Exception: Limit for result exceeded, max bytes: 976.56 KiB. (TOO_MANY_ROWS_OR_BYTES)', + ), + }); + const output = parseText( + await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 5, + }), + ); + expect(output.columns).toBeDefined(); + expect(output.sampleRowsError).toMatch(/TOO_MANY_ROWS_OR_BYTES/); }); it('keeps several medium sample rows within the sample budget', async () => { diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts index 2c3d05e54c..20951193d3 100644 --- a/packages/api/src/mcp/tools/sources/describeTable.ts +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -1,8 +1,4 @@ -import { - getMetadata, - type Metadata, - quoteIdentifierIfNeeded, -} from '@hyperdx/common-utils/dist/core/metadata'; +import { quoteIdentifierIfNeeded } from '@hyperdx/common-utils/dist/core/metadata'; import { trace } from '@opentelemetry/api'; import { ClickhouseClient } from '@/clickhouse'; @@ -162,13 +158,12 @@ async function sampleTableRows( async function describeTable( ctx: TableQueryContext, - metadata: Metadata, input: DescribeTableInput & { table: string }, ): Promise { const { database, table } = input; let details; try { - details = await fetchTableDetails(ctx, metadata, { database, table }); + details = await fetchTableDetails(ctx, { database, table }); } catch (e) { return tableQueryErrorResult(e, 'Failed to describe table'); } @@ -180,9 +175,10 @@ async function describeTable( ); } + const { indicesUnavailable, ...shapeInput } = details; const { columns, ...description } = shapeTableDescription({ database, - ...details, + ...shapeInput, }); let sample = {}; @@ -209,6 +205,10 @@ async function describeTable( const rest = { ...description, + ...(indicesUnavailable && { + dataSkippingIndicesHint: + "The connection's ClickHouse user cannot read system.data_skipping_indices, so indices are not listed.", + }), ...sample, nextSteps: { query: @@ -299,10 +299,7 @@ export function registerDescribeTable({ signal, }; return input.table - ? describeTable(ctx, getMetadata(clickhouseClient), { - ...input, - table: input.table, - }) + ? describeTable(ctx, { ...input, table: input.table }) : listTables(ctx, input); }, { timeoutMs: MCP_TOOL_TIMEOUT_MS }, diff --git a/packages/api/src/mcp/tools/sources/describeTableQueries.ts b/packages/api/src/mcp/tools/sources/describeTableQueries.ts index 0e6c77d2b5..3947beec0a 100644 --- a/packages/api/src/mcp/tools/sources/describeTableQueries.ts +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -5,7 +5,6 @@ import { tableExpr, } from '@hyperdx/common-utils/dist/clickhouse'; import type { - Metadata, SkipIndexMetadata, TableMetadata, } from '@hyperdx/common-utils/dist/core/metadata'; @@ -168,10 +167,28 @@ export async function databaseExists(ctx: TableQueryContext, database: string) { return rows.length > 0; } -/** Returns undefined when the table does not exist or is not visible. */ +type SystemTableName = 'tables' | 'columns' | 'data_skipping_indices'; + +/** A system table, read across the cluster when `cluster` is set. */ +function systemTable(name: SystemTableName, cluster?: string): ChSql { + return cluster + ? chSql`cluster(${{ String: cluster }}, system.${name})` + : chSql`system.${name}`; +} + +function isPrivilegeError(e: unknown) { + const message = e instanceof Error ? e.message : String(e); + return /ACCESS_DENIED|NOT_ENOUGH_PRIVILEGES|Not enough privileges/i.test( + message, + ); +} + +/** + * Returns undefined when the table does not exist or is not visible. Every + * lookup is uncached so the result reflects ALTERs made since the last call. + */ export async function fetchTableDetails( ctx: TableQueryContext, - metadata: Pick, { database, table }: { database: string; table: string }, ) { const rows = await queryRows( @@ -185,16 +202,39 @@ export async function fetchTableDetails( const info = rows.at(0); if (!info) return undefined; + // A Distributed table has no keys or indices of its own; read them from + // the local table it routes to. const distributedTarget = info.engine === 'Distributed' ? getDistributedTableArgs(info) : undefined; - const metadataArgs = { - databaseName: database, - tableName: table, - connectionId: ctx.connectionId, - }; - const [keys, indices, columns, localKeyFlags] = await Promise.all([ - metadata.getTableMetadata(metadataArgs), - metadata.getSkipIndices(metadataArgs), + const keyTable = distributedTarget ?? { database, table, cluster: undefined }; + + const [localKeys, indices, columns, localKeyFlags] = await Promise.all([ + distributedTarget + ? queryRows( + ctx, + chSql` + SELECT engine, sorting_key, primary_key, partition_key + FROM ${systemTable('tables', keyTable.cluster)} + WHERE database = ${{ String: keyTable.database }} + AND name = ${{ String: keyTable.table }} + LIMIT 1 + `, + ) + : undefined, + queryRows>( + ctx, + chSql` + SELECT name, type, expr AS expression + FROM ${systemTable('data_skipping_indices', keyTable.cluster)} + WHERE database = ${{ String: keyTable.database }} + AND table = ${{ String: keyTable.table }} + LIMIT 1 BY name + `, + ).catch((e: unknown) => { + // Indices are optional detail; a missing grant should not hide the schema. + if (isPrivilegeError(e)) return undefined; + throw e; + }), queryRows( ctx, chSql` @@ -210,9 +250,9 @@ export async function fetchTableDetails( ctx, chSql` SELECT name, is_in_sorting_key, is_in_primary_key - FROM cluster(${{ String: distributedTarget.cluster }}, system.columns) - WHERE database = ${{ String: distributedTarget.database }} - AND table = ${{ String: distributedTarget.table }} + FROM ${systemTable('columns', keyTable.cluster)} + WHERE database = ${{ String: keyTable.database }} + AND table = ${{ String: keyTable.table }} LIMIT 1 BY name `, ) @@ -222,9 +262,10 @@ export async function fetchTableDetails( const flagsByName = new Map(localKeyFlags?.map(f => [f.name, f])); return { table: info, - keys: keys ?? info, + keys: localKeys?.at(0) ?? info, distributedTarget, - indices, + indices: indices ?? [], + indicesUnavailable: indices === undefined, columns: columns.map(c => { const flags = flagsByName.get(c.name); return flags @@ -238,8 +279,9 @@ export async function fetchTableDetails( }; } -// Five rows of wide String columns can still be large; cap what ClickHouse -// sends back so the response is bounded before it is buffered. +// Five rows of wide String columns can still be large. Past this cap +// ClickHouse fails the query, which surfaces as sampleRowsError, so nothing +// unbounded is buffered. const SAMPLE_MAX_RESULT_BYTES = 1_000_000; export function fetchSampleRows( From 22093743cebdd2c8bafbf62522b79c65dcaddf06 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 16:04:24 -0600 Subject: [PATCH 4/7] fix(mcp): keep describe_table working without REMOTE grant or on sample timeout (review) - The cluster() lookups for a Distributed table's local keys and key flags now fall back to the Distributed table on a privilege error and add a hint, instead of failing the describe. - A sample that hits max_execution_time becomes sampleRowsError; only connection failures still fail the call. - Share isPrivilegeError between the two files. --- .../src/mcp/__tests__/describeTable.test.ts | 59 +++++++++++++++++++ .../src/mcp/tools/sources/describeTable.ts | 16 +++-- .../mcp/tools/sources/describeTableQueries.ts | 45 +++++++++----- 3 files changed, 99 insertions(+), 21 deletions(-) diff --git a/packages/api/src/mcp/__tests__/describeTable.test.ts b/packages/api/src/mcp/__tests__/describeTable.test.ts index af3b101b72..b495534a88 100644 --- a/packages/api/src/mcp/__tests__/describeTable.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -10,6 +10,7 @@ jest.mock('@/clickhouse', () => ({ ClickhouseClient: jest.fn(() => ({ query: mockQuery })), })); +import { ClickHouseError } from '@clickhouse/client-common'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { registerDescribeTable } from '@/mcp/tools/sources/describeTable'; @@ -457,6 +458,64 @@ describe('clickstack_describe_table', () => { } }); + it('falls back to the Distributed table when cluster() is not granted', async () => { + const denied = new Error( + 'Code: 497. DB::Exception: Not enough privileges. To execute this query, it is necessary to have the grant REMOTE ON *.*. (ACCESS_DENIED)', + ); + respondWith({ + 'FROM system.tables': [ + { + ...TABLE_ROW, + engine: 'Distributed', + engine_full: + "Distributed('prod', 'default', 'events_local', rand())", + sorting_key: '', + primary_key: '', + partition_key: '', + }, + ], + ', system.tables)': denied, + ', system.data_skipping_indices)': denied, + ', system.columns)': denied, + 'FROM system.columns': COLUMN_ROWS, + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }); + expect(result.isError).toBeFalsy(); + const output = parseText(result); + expect(output.table).toMatchObject({ + engine: 'Distributed', + sortingKey: '', + }); + expect(output.columns).toHaveLength(2); + expect(output.localTableHint).toMatch(/REMOTE grant/); + }); + + it('keeps the schema when the sample times out', async () => { + const timeout = new Error('query failed', { + cause: new ClickHouseError({ + message: 'Timeout exceeded: elapsed 30s', + code: '159', + type: 'TIMEOUT_EXCEEDED', + }), + }); + respondWith({ + 'FROM system.tables': [TABLE_ROW], + ':Identifier}': timeout, + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + sampleRows: 1, + }); + expect(result.isError).toBeFalsy(); + expect(parseText(result).sampleRowsError).toBeDefined(); + }); + it('still describes the table when skip indices are not readable', async () => { respondWith({ 'FROM system.tables': [TABLE_ROW], diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts index 20951193d3..c13f73410f 100644 --- a/packages/api/src/mcp/tools/sources/describeTable.ts +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -6,6 +6,7 @@ import { getConnectionById } from '@/controllers/connection'; import { READ_ONLY_TOOL } from '@/mcp/tools/annotations'; import { clickHouseErrorResult, + isQueryOutOfTime, isServerError, MCP_REQUEST_TIMEOUT, } from '@/mcp/tools/query/helpers'; @@ -25,6 +26,7 @@ import { fetchSampleRows, fetchTableDetails, fetchTablePage, + isPrivilegeError, shapeTableDescription, shapeTableListEntry, type TableQueryContext, @@ -54,9 +56,7 @@ function jsonResult(payload: unknown): ToolResult { function tableQueryErrorResult(e: unknown, prefix: string) { const message = e instanceof Error ? e.message : String(e); let suffix: string | undefined; - if ( - /ACCESS_DENIED|NOT_ENOUGH_PRIVILEGES|Not enough privileges/i.test(message) - ) { + if (isPrivilegeError(e)) { suffix = "The ClickHouse user on this connection lacks the privileges needed. Ask the connection's administrator to grant them, or use another connection."; } else if (/UNKNOWN_(TABLE|DATABASE)|does not exist/i.test(message)) { @@ -143,7 +143,9 @@ async function sampleTableRows( const rows = trimSampleRows(await fetchSampleRows(ctx, args)); return { sampleRows: rows.data, isTrimmed: rows.isTrimmed }; } catch (e) { - if (isServerError(e)) throw e; + // A slow View or Distributed table can time out on SELECT *; only a + // connection failure should cost the schema already fetched. + if (isServerError(e) && !isQueryOutOfTime(e)) throw e; // The schema is still useful when the user can read system tables but // not the table itself. return { @@ -175,7 +177,7 @@ async function describeTable( ); } - const { indicesUnavailable, ...shapeInput } = details; + const { indicesUnavailable, localTableUnavailable, ...shapeInput } = details; const { columns, ...description } = shapeTableDescription({ database, ...shapeInput, @@ -205,6 +207,10 @@ async function describeTable( const rest = { ...description, + ...(localTableUnavailable && { + localTableHint: + "The connection's ClickHouse user cannot run cluster() (REMOTE grant), so keys and key flags come from the Distributed table itself and may be empty.", + }), ...(indicesUnavailable && { dataSkippingIndicesHint: "The connection's ClickHouse user cannot read system.data_skipping_indices, so indices are not listed.", diff --git a/packages/api/src/mcp/tools/sources/describeTableQueries.ts b/packages/api/src/mcp/tools/sources/describeTableQueries.ts index 3947beec0a..fd719766c0 100644 --- a/packages/api/src/mcp/tools/sources/describeTableQueries.ts +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -176,13 +176,21 @@ function systemTable(name: SystemTableName, cluster?: string): ChSql { : chSql`system.${name}`; } -function isPrivilegeError(e: unknown) { +export function isPrivilegeError(e: unknown) { const message = e instanceof Error ? e.message : String(e); return /ACCESS_DENIED|NOT_ENOUGH_PRIVILEGES|Not enough privileges/i.test( message, ); } +/** Resolves to null instead of rejecting when the user lacks a grant. */ +function nullOnPrivilegeError(query: Promise): Promise { + return query.catch((e: unknown) => { + if (isPrivilegeError(e)) return null; + throw e; + }); +} + /** * Returns undefined when the table does not exist or is not visible. Every * lookup is uncached so the result reflects ALTERs made since the last call. @@ -210,31 +218,33 @@ export async function fetchTableDetails( const [localKeys, indices, columns, localKeyFlags] = await Promise.all([ distributedTarget - ? queryRows( - ctx, - chSql` + ? nullOnPrivilegeError( + queryRows( + ctx, + chSql` SELECT engine, sorting_key, primary_key, partition_key FROM ${systemTable('tables', keyTable.cluster)} WHERE database = ${{ String: keyTable.database }} AND name = ${{ String: keyTable.table }} LIMIT 1 `, + ), ) : undefined, - queryRows>( - ctx, - chSql` + // Indices and the local-table lookups are optional detail, and cluster() + // needs the REMOTE grant; a missing grant should not hide the schema. + nullOnPrivilegeError( + queryRows>( + ctx, + chSql` SELECT name, type, expr AS expression FROM ${systemTable('data_skipping_indices', keyTable.cluster)} WHERE database = ${{ String: keyTable.database }} AND table = ${{ String: keyTable.table }} LIMIT 1 BY name `, - ).catch((e: unknown) => { - // Indices are optional detail; a missing grant should not hide the schema. - if (isPrivilegeError(e)) return undefined; - throw e; - }), + ), + ), queryRows( ctx, chSql` @@ -246,15 +256,17 @@ export async function fetchTableDetails( `, ), distributedTarget - ? queryRows( - ctx, - chSql` + ? nullOnPrivilegeError( + queryRows( + ctx, + chSql` SELECT name, is_in_sorting_key, is_in_primary_key FROM ${systemTable('columns', keyTable.cluster)} WHERE database = ${{ String: keyTable.database }} AND table = ${{ String: keyTable.table }} LIMIT 1 BY name `, + ), ) : undefined, ]); @@ -265,7 +277,8 @@ export async function fetchTableDetails( keys: localKeys?.at(0) ?? info, distributedTarget, indices: indices ?? [], - indicesUnavailable: indices === undefined, + indicesUnavailable: indices === null, + localTableUnavailable: localKeys === null || localKeyFlags === null, columns: columns.map(c => { const flags = flagsByName.get(c.name); return flags From bfd941d4d993a6f356a99815ce83a12deffead1d Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Tue, 3 Nov 2026 16:14:02 -0600 Subject: [PATCH 5/7] fix(mcp): give the describe_table sample its own time limit (review) - Cap the sample at max_execution_time=10s so a slow sample fails as sampleRowsError inside the 30s tool budget instead of timing out the call. - Fall back to the Distributed table on any cluster() failure (unknown cluster, unreachable shard), not only a missing grant. - Omit localTable.engine when the local table could not be read. --- .../src/mcp/__tests__/describeTable.test.ts | 39 ++++++++++++++++- .../src/mcp/tools/sources/describeTable.ts | 4 +- .../mcp/tools/sources/describeTableQueries.ts | 42 +++++++++++++++---- 3 files changed, 75 insertions(+), 10 deletions(-) diff --git a/packages/api/src/mcp/__tests__/describeTable.test.ts b/packages/api/src/mcp/__tests__/describeTable.test.ts index b495534a88..bb7ef19e21 100644 --- a/packages/api/src/mcp/__tests__/describeTable.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -492,6 +492,41 @@ describe('clickstack_describe_table', () => { }); expect(output.columns).toHaveLength(2); expect(output.localTableHint).toMatch(/REMOTE grant/); + expect(output.table.localTable).toEqual({ + cluster: 'prod', + database: 'default', + table: 'events_local', + }); + }); + + it('falls back to the Distributed table when its cluster is unknown', async () => { + const unknownCluster = new Error( + "Code: 701. DB::Exception: Requested cluster 'prod' not found. (CLUSTER_DOESNT_EXIST)", + ); + respondWith({ + 'FROM system.tables': [ + { + ...TABLE_ROW, + engine: 'Distributed', + engine_full: + "Distributed('prod', 'default', 'events_local', rand())", + }, + ], + ', system.tables)': unknownCluster, + ', system.data_skipping_indices)': unknownCluster, + ', system.columns)': unknownCluster, + 'FROM system.columns': COLUMN_ROWS, + }); + const result = await getHandler()({ + connectionId: CONNECTION_ID, + database: 'default', + table: 'events', + }); + expect(result.isError).toBeFalsy(); + const output = parseText(result); + expect(output.columns).toHaveLength(2); + expect(output.localTableHint).toMatch(/unknown cluster/); + expect(output.dataSkippingIndicesHint).toBeDefined(); }); it('keeps the schema when the sample times out', async () => { @@ -531,7 +566,7 @@ describe('clickstack_describe_table', () => { expect(result.isError).toBeFalsy(); const output = parseText(result); expect(output.dataSkippingIndices).toEqual([]); - expect(output.dataSkippingIndicesHint).toMatch(/cannot read/); + expect(output.dataSkippingIndicesHint).toMatch(/Could not read/); }); it('reports a sample over the byte cap as sampleRowsError', async () => { @@ -590,6 +625,8 @@ describe('clickstack_describe_table', () => { c.query.includes(':Identifier}'), ); expect(sampleCall?.clickhouse_settings.max_result_bytes).toBeDefined(); + // Below the 30s tool budget, so a slow sample fails on its own. + expect(sampleCall?.clickhouse_settings.max_execution_time).toBe(10); }); it('reports an infrastructure failure while sampling as a server error', async () => { diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts index c13f73410f..21bca26c1b 100644 --- a/packages/api/src/mcp/tools/sources/describeTable.ts +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -209,11 +209,11 @@ async function describeTable( ...description, ...(localTableUnavailable && { localTableHint: - "The connection's ClickHouse user cannot run cluster() (REMOTE grant), so keys and key flags come from the Distributed table itself and may be empty.", + 'Could not read the local table through cluster() (missing REMOTE grant, unknown cluster, or unreachable shard), so keys and key flags come from the Distributed table itself and may be empty.', }), ...(indicesUnavailable && { dataSkippingIndicesHint: - "The connection's ClickHouse user cannot read system.data_skipping_indices, so indices are not listed.", + 'Could not read system.data_skipping_indices for this table (missing grant, or the local table was unreachable), so indices are not listed.', }), ...sample, nextSteps: { diff --git a/packages/api/src/mcp/tools/sources/describeTableQueries.ts b/packages/api/src/mcp/tools/sources/describeTableQueries.ts index fd719766c0..85b00dd8cf 100644 --- a/packages/api/src/mcp/tools/sources/describeTableQueries.ts +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -12,6 +12,7 @@ import { getDistributedTableArgs } from '@hyperdx/common-utils/dist/core/utils'; import type { ClickhouseClient } from '@/clickhouse'; import { MCP_CLICKHOUSE_SETTINGS } from '@/mcp/tools/query/helpers'; +import logger from '@/utils/logger'; export type TableQueryContext = { clickhouseClient: Pick; @@ -92,6 +93,7 @@ export function shapeTableDescription({ table, keys, distributedTarget, + localTableEngine, columns, indices, }: { @@ -99,6 +101,8 @@ export function shapeTableDescription({ table: TableInfo; keys: TableKeys; distributedTarget?: DistributedTarget; + /** Unset when the local table could not be read. */ + localTableEngine?: string; columns: SystemColumnRow[]; indices: Pick[]; }) { @@ -108,7 +112,10 @@ export function shapeTableDescription({ name: table.name, engine: table.engine, ...(distributedTarget && { - localTable: { ...distributedTarget, engine: keys.engine }, + localTable: { + ...distributedTarget, + ...(localTableEngine && { engine: localTableEngine }), + }, }), sortingKey: keys.sorting_key, primaryKey: keys.primary_key, @@ -191,6 +198,18 @@ function nullOnPrivilegeError(query: Promise): Promise { }); } +/** + * Resolves to null on any failure. The cluster() lookups can fail for reasons + * besides grants (unknown cluster, unreachable shard), and the Distributed + * table's own schema is still worth returning. + */ +function nullOnFailure(query: Promise): Promise { + return query.catch((e: unknown) => { + logger.warn({ error: e }, 'Distributed local-table lookup failed'); + return null; + }); +} + /** * Returns undefined when the table does not exist or is not visible. Every * lookup is uncached so the result reflects ALTERs made since the last call. @@ -218,7 +237,7 @@ export async function fetchTableDetails( const [localKeys, indices, columns, localKeyFlags] = await Promise.all([ distributedTarget - ? nullOnPrivilegeError( + ? nullOnFailure( queryRows( ctx, chSql` @@ -233,7 +252,7 @@ export async function fetchTableDetails( : undefined, // Indices and the local-table lookups are optional detail, and cluster() // needs the REMOTE grant; a missing grant should not hide the schema. - nullOnPrivilegeError( + (distributedTarget ? nullOnFailure : nullOnPrivilegeError)( queryRows>( ctx, chSql` @@ -256,7 +275,7 @@ export async function fetchTableDetails( `, ), distributedTarget - ? nullOnPrivilegeError( + ? nullOnFailure( queryRows( ctx, chSql` @@ -271,14 +290,17 @@ export async function fetchTableDetails( : undefined, ]); + const localKeyRow = localKeys?.at(0); const flagsByName = new Map(localKeyFlags?.map(f => [f.name, f])); return { table: info, - keys: localKeys?.at(0) ?? info, + keys: localKeyRow ?? info, distributedTarget, + localTableEngine: localKeyRow?.engine, indices: indices ?? [], indicesUnavailable: indices === null, - localTableUnavailable: localKeys === null || localKeyFlags === null, + localTableUnavailable: + distributedTarget != null && (!localKeyRow || localKeyFlags === null), columns: columns.map(c => { const flags = flagsByName.get(c.name); return flags @@ -296,6 +318,9 @@ export async function fetchTableDetails( // ClickHouse fails the query, which surfaces as sampleRowsError, so nothing // unbounded is buffered. const SAMPLE_MAX_RESULT_BYTES = 1_000_000; +// Well inside the tool's budget, which the schema queries have already used +// part of, so a slow sample fails on its own instead of timing out the call. +const SAMPLE_MAX_EXECUTION_SEC = 10; export function fetchSampleRows( ctx: TableQueryContext, @@ -308,6 +333,9 @@ export function fetchSampleRows( return queryRows>( ctx, chSql`SELECT * FROM ${tableExpr({ database, table })} LIMIT ${{ Int32: limit }}`, - { max_result_bytes: SAMPLE_MAX_RESULT_BYTES.toString() }, + { + max_result_bytes: SAMPLE_MAX_RESULT_BYTES.toString(), + max_execution_time: SAMPLE_MAX_EXECUTION_SEC, + }, ); } From 92d36283e3ecf096b99a40191ad4483548b79247 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Wed, 4 Nov 2026 12:51:13 -0600 Subject: [PATCH 6/7] fix(mcp): point describe_table to describe_source for source-backed tables --- .../mcp/__tests__/describeTable.int.test.ts | 66 +++++++++++++++++ .../src/mcp/tools/sources/describeTable.ts | 74 ++++++++++++++++--- 2 files changed, 131 insertions(+), 9 deletions(-) diff --git a/packages/api/src/mcp/__tests__/describeTable.int.test.ts b/packages/api/src/mcp/__tests__/describeTable.int.test.ts index 3c67659f89..c1c4344ab3 100644 --- a/packages/api/src/mcp/__tests__/describeTable.int.test.ts +++ b/packages/api/src/mcp/__tests__/describeTable.int.test.ts @@ -1,3 +1,4 @@ +import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import mongoose from 'mongoose'; @@ -9,6 +10,7 @@ import { } from '@/fixtures'; import { McpContext } from '@/mcp/tools/types'; import Connection from '@/models/connection'; +import { Source } from '@/models/source'; import { callTool, createTestClient, getFirstText } from './mcpTestUtils'; @@ -194,6 +196,70 @@ describe('clickstack_describe_table', () => { } }); + it('points to clickstack_describe_source when the table backs a source', async () => { + const source = await Source.create({ + kind: SourceKind.Log, + team: connection.team, + from: { databaseName: DATABASE, tableName: 'events' }, + timestampValueExpression: 'Timestamp', + connection: connection._id, + name: 'Events', + }); + const metricSource = await Source.create({ + kind: SourceKind.Metric, + team: connection.team, + from: { databaseName: DATABASE, tableName: '' }, + metricTables: { gauge: 'events' }, + timestampValueExpression: 'Timestamp', + resourceAttributesExpression: 'ResourceAttributes', + connection: connection._id, + name: 'Event metrics', + }); + + const output = JSON.parse( + getFirstText(await call({ database: DATABASE, table: 'events' })), + ); + expect(output.sources).toEqual( + expect.arrayContaining([ + { id: source._id.toString(), name: 'Events', kind: SourceKind.Log }, + { + id: metricSource._id.toString(), + name: 'Event metrics', + kind: SourceKind.Metric, + }, + ]), + ); + expect(output.nextSteps.describeSource).toMatch( + /clickstack_describe_source/, + ); + expect(output.nextSteps.query).toBeUndefined(); + expect(output.columns.length).toBeGreaterThan(0); + }); + + it('does not match a source on another connection', async () => { + const otherConnection = await Connection.create({ + team: connection.team, + name: 'Other', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + await Source.create({ + kind: SourceKind.Log, + team: connection.team, + from: { databaseName: DATABASE, tableName: 'events' }, + timestampValueExpression: 'Timestamp', + connection: otherConnection._id, + name: 'Events', + }); + + const output = JSON.parse( + getFirstText(await call({ database: DATABASE, table: 'events' })), + ); + expect(output.sources).toBeUndefined(); + expect(output.nextSteps.query).toMatch(/clickstack_save_source/); + }); + it('returns a capped sample of rows when asked', async () => { const result = await call({ database: DATABASE, diff --git a/packages/api/src/mcp/tools/sources/describeTable.ts b/packages/api/src/mcp/tools/sources/describeTable.ts index 21bca26c1b..63290a75f1 100644 --- a/packages/api/src/mcp/tools/sources/describeTable.ts +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -1,4 +1,5 @@ import { quoteIdentifierIfNeeded } from '@hyperdx/common-utils/dist/core/metadata'; +import { SourceKind } from '@hyperdx/common-utils/dist/types'; import { trace } from '@opentelemetry/api'; import { ClickhouseClient } from '@/clickhouse'; @@ -18,6 +19,7 @@ import { validateObjectId, } from '@/mcp/utils/errors'; import { MCP_TOOL_TIMEOUT_MS, runWithTimeout } from '@/mcp/utils/timeout'; +import { Source } from '@/models/source'; import logger from '@/utils/logger'; import { trimToolResponse } from '@/utils/trimToolResponse'; @@ -158,14 +160,55 @@ async function sampleTableRows( } } +/** Sources on this connection that read the table, so callers can be sent to clickstack_describe_source. */ +async function findSourcesForTable(args: { + teamId: string; + connectionId: string; + database: string; + table: string; +}) { + try { + const sources = await Source.find({ + team: args.teamId, + connection: args.connectionId, + 'from.databaseName': args.database, + kind: { $ne: SourceKind.Promql }, + }).lean(); + return sources + .filter( + s => + s.from?.tableName === args.table || + ('metricTables' in s && + Object.values(s.metricTables ?? {}).includes(args.table)), + ) + .map(s => ({ id: s._id.toString(), name: s.name, kind: s.kind })); + } catch (e) { + // The schema is still worth returning without the source hint. + logger.warn( + { err: e, teamId: args.teamId, connectionId: args.connectionId }, + 'clickstack_describe_table source lookup failed', + ); + return []; + } +} + async function describeTable( ctx: TableQueryContext, - input: DescribeTableInput & { table: string }, + input: DescribeTableInput & { table: string; teamId: string }, ): Promise { const { database, table } = input; let details; + let sources; try { - details = await fetchTableDetails(ctx, { database, table }); + [details, sources] = await Promise.all([ + fetchTableDetails(ctx, { database, table }), + findSourcesForTable({ + teamId: input.teamId, + connectionId: ctx.connectionId, + database, + table, + }), + ]); } catch (e) { return tableQueryErrorResult(e, 'Failed to describe table'); } @@ -203,10 +246,28 @@ async function describeTable( 'mcp.describe_table.column_count': columns.length, 'mcp.describe_table.index_count': description.dataSkippingIndices.length, 'mcp.describe_table.engine': description.table.engine, + 'mcp.describe_table.source_count': sources.length, }); + const qualifiedName = `${quoteIdentifierIfNeeded(database)}.${quoteIdentifierIfNeeded(table)}`; + const nextSteps = + sources.length > 0 + ? { + describeSource: + `This table backs ${sources.length === 1 ? 'a source' : 'sources'}. ` + + `Call clickstack_describe_source with sourceId "${sources[0].id}" for attribute keys and sampled values, ` + + 'then query it with the builder tools (clickstack_table, clickstack_timeseries, clickstack_search) rather than clickstack_sql.', + } + : { + query: + `Query it with clickstack_sql using connectionId "${ctx.connectionId}" and ` + + `FROM ${qualifiedName}. ` + + 'To query it with the builder tools, register it as a source with clickstack_save_source.', + }; + const rest = { ...description, + ...(sources.length > 0 && { sources }), ...(localTableUnavailable && { localTableHint: 'Could not read the local table through cluster() (missing REMOTE grant, unknown cluster, or unreachable shard), so keys and key flags come from the Distributed table itself and may be empty.', @@ -216,12 +277,7 @@ async function describeTable( 'Could not read system.data_skipping_indices for this table (missing grant, or the local table was unreachable), so indices are not listed.', }), ...sample, - nextSteps: { - query: - `Query it with clickstack_sql using connectionId "${ctx.connectionId}" and ` + - `FROM ${quoteIdentifierIfNeeded(database)}.${quoteIdentifierIfNeeded(table)}. ` + - 'To query it with the builder tools, register it as a source with clickstack_save_source.', - }, + nextSteps, }; // Trim only the column list, against what the rest of the response leaves, // so a large sample cannot push columns out. @@ -305,7 +361,7 @@ export function registerDescribeTable({ signal, }; return input.table - ? describeTable(ctx, { ...input, table: input.table }) + ? describeTable(ctx, { ...input, table: input.table, teamId }) : listTables(ctx, input); }, { timeoutMs: MCP_TOOL_TIMEOUT_MS }, From 37c7c41b75f5bb97d49515ef14516a58f52b8014 Mon Sep 17 00:00:00 2001 From: Brandon Pereira Date: Wed, 4 Nov 2026 13:06:01 -0600 Subject: [PATCH 7/7] fix(mcp): keep clickstack_sql description under the 2048-byte limit main added the batch-queries paragraph, which pushed the merged description to 2138 bytes. Drop the sentence that repeated the 'only use raw SQL' rule and tighten the describe_table pointer. --- packages/api/src/mcp/tools/query/sql.ts | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/api/src/mcp/tools/query/sql.ts b/packages/api/src/mcp/tools/query/sql.ts index 3f5f2b1a5d..82a33f7a8a 100644 --- a/packages/api/src/mcp/tools/query/sql.ts +++ b/packages/api/src/mcp/tools/query/sql.ts @@ -109,13 +109,12 @@ export function registerSql({ context, registerTool }: ToolRegistrar) { 'ONLY use raw SQL when the query genuinely cannot be expressed by a builder tool — ' + `i.e. it requires ${SQL_FALLBACK_CRITERIA}. ` + 'A single-table aggregation, top-N, time-series, or row browse is ALWAYS a ' + - 'builder-tool job, never raw SQL. If you are unsure, try the builder tool first and only ' + - 'fall back to SQL if it cannot express what you need.\n\n' + + 'builder-tool job, never raw SQL.\n\n' + 'Requires connectionId — call clickstack_list_sources to find connections. Pass sourceId ' + 'too when querying a source, so $__sourceTable resolves to its table. ' + - 'Call clickstack_describe_source to discover column names before writing SQL. ' + - 'For tables that are not sources, call clickstack_describe_table instead of ' + - 'querying system.columns / system.tables directly.\n\n' + + 'Before writing SQL, get column names from clickstack_describe_source, or ' + + 'clickstack_describe_table for tables that are not sources (instead of ' + + 'querying system.columns / system.tables).\n\n' + 'Results are always returned as table rows — for time-series semantics, ' + 'include a time column and ORDER BY it in your SQL.\n\n' + 'To run several independent queries at once (e.g. comparing sources), pass ' +