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 2ece9e1d33..724e754630 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..c1c4344ab3 --- /dev/null +++ b/packages/api/src/mcp/__tests__/describeTable.int.test.ts @@ -0,0 +1,338 @@ +import { SourceKind } from '@hyperdx/common-utils/dist/types'; +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 { Source } from '@/models/source'; + +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(); + 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: `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 ( + 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.command({ + query: `CREATE TABLE ${DIST_DATABASE}.events AS ${DATABASE}.events ENGINE = Distributed(hdx_cluster, ${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 ${DIST_DATABASE}` }); + 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('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('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, + 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', approxTotalRows: 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', approxTotalRows: 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..bb7ef19e21 --- /dev/null +++ b/packages/api/src/mcp/__tests__/describeTable.test.ts @@ -0,0 +1,649 @@ +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 { ClickHouseError } from '@clickhouse/client-common'; +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; + clickhouse_settings: 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 }, + keys: TABLE_ROW, + columns: COLUMN_ROWS, + indices: [{ name: 'idx_day', type: 'minmax', expression: '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', approxTotalRows: 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', + sampleRows: 1, + }); + 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)).toEqual( + expect.arrayContaining(["evil'db", 'events`; DROP TABLE x']), + ); + } + }); + + 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', approxTotalRows: 1 }, + { name: 'b', engine: 'View', approxTotalRows: 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.tables': [TABLE_ROW], + '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(':Identifier}'))).toBe( + false, + ); + }); + + it('trims oversized sample rows', async () => { + respondWith({ + 'FROM system.tables': [TABLE_ROW], + ':Identifier}': [{ 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], + ':Identifier}': 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'); + 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, + }, + ], + ', 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 => ({ + ...c, + is_in_sorting_key: 0, + is_in_primary_key: 0, + })), + }); + + 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' }, + ]); + 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('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/); + 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 () => { + 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], + '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(/Could not 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 () => { + 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(); + // 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 () => { + 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/query/sql.ts b/packages/api/src/mcp/tools/query/sql.ts index fd68dbc886..82a33f7a8a 100644 --- a/packages/api/src/mcp/tools/query/sql.ts +++ b/packages/api/src/mcp/tools/query/sql.ts @@ -109,11 +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.\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 ' + 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..63290a75f1 --- /dev/null +++ b/packages/api/src/mcp/tools/sources/describeTable.ts @@ -0,0 +1,380 @@ +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'; +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'; +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 { Source } from '@/models/source'; +import logger from '@/utils/logger'; +import { trimToolResponse } from '@/utils/trimToolResponse'; + +import { + databaseExists, + fetchSampleRows, + fetchTableDetails, + fetchTablePage, + isPrivilegeError, + 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 RESPONSE_MAX_CHARS = 50_000; + +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 (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)) { + 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) { + 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.', + }, + }); +} + +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) { + // 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 { + 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, + }; + } +} + +/** 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; teamId: string }, +): Promise { + const { database, table } = input; + let details; + let sources; + try { + [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'); + } + + 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 { indicesUnavailable, localTableUnavailable, ...shapeInput } = details; + const { columns, ...description } = shapeTableDescription({ + database, + ...shapeInput, + }); + + let sample = {}; + let sampleTrimmed = false; + if (input.sampleRows > 0) { + try { + const { isTrimmed, ...result } = await sampleTableRows(ctx, { + database, + table, + limit: input.sampleRows, + }); + sample = result; + sampleTrimmed = isTrimmed; + } catch (e) { + return clickHouseErrorResult(e, 'Failed to sample rows'); + } + } + + trace.getActiveSpan()?.setAttributes({ + '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.', + }), + ...(indicesUnavailable && { + dataSkippingIndicesHint: + '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, + }; + // 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({ + 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, teamId }) + : 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..85b00dd8cf --- /dev/null +++ b/packages/api/src/mcp/tools/sources/describeTableQueries.ts @@ -0,0 +1,341 @@ +import type { ClickHouseSettings } from '@clickhouse/client-common'; +import { + type ChSql, + chSql, + tableExpr, +} from '@hyperdx/common-utils/dist/clickhouse'; +import type { + 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'; +import logger from '@/utils/logger'; + +export type TableQueryContext = { + clickhouseClient: Pick; + connectionId: string; + signal: AbortSignal; +}; + +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, ...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; +}; + +// total_rows / total_bytes are NULL for views, unlike TableMetadata's typing. +type TableInfo = Pick & { + total_rows: string | null; + total_bytes: string | null; +}; + +type TableKeys = Pick< + TableMetadata, + 'engine' | 'sorting_key' | 'primary_key' | 'partition_key' +>; + +type ColumnKeyFlags = { + name: string; + is_in_sorting_key: number; + is_in_primary_key: number; +}; + +export type SystemColumnRow = ColumnKeyFlags & { + type: 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, + 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, + localTableEngine, + columns, + indices, +}: { + database: string; + table: TableInfo; + keys: TableKeys; + distributedTarget?: DistributedTarget; + /** Unset when the local table could not be read. */ + localTableEngine?: string; + columns: SystemColumnRow[]; + indices: Pick[]; +}) { + return { + table: { + database, + name: table.name, + engine: table.engine, + ...(distributedTarget && { + localTable: { + ...distributedTarget, + ...(localTableEngine && { engine: localTableEngine }), + }, + }), + 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), + }, + 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.expression, + })), + }; +} + +/** 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; +} + +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}`; +} + +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; + }); +} + +/** + * 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. + */ +export async function fetchTableDetails( + ctx: TableQueryContext, + { database, table }: { database: string; table: string }, +) { + 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; + + // 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 keyTable = distributedTarget ?? { database, table, cluster: undefined }; + + const [localKeys, indices, columns, localKeyFlags] = await Promise.all([ + distributedTarget + ? nullOnFailure( + 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, + // Indices and the local-table lookups are optional detail, and cluster() + // needs the REMOTE grant; a missing grant should not hide the schema. + (distributedTarget ? nullOnFailure : 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 + `, + ), + ), + 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 + `, + ), + distributedTarget + ? nullOnFailure( + 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, + ]); + + const localKeyRow = localKeys?.at(0); + const flagsByName = new Map(localKeyFlags?.map(f => [f.name, f])); + return { + table: info, + keys: localKeyRow ?? info, + distributedTarget, + localTableEngine: localKeyRow?.engine, + indices: indices ?? [], + indicesUnavailable: indices === null, + localTableUnavailable: + distributedTarget != null && (!localKeyRow || localKeyFlags === null), + 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. Past this cap +// 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, + { + database, + table, + limit, + }: { database: string; table: string; limit: number }, +) { + return queryRows>( + ctx, + chSql`SELECT * FROM ${tableExpr({ database, table })} LIMIT ${{ Int32: limit }}`, + { + max_result_bytes: SAMPLE_MAX_RESULT_BYTES.toString(), + max_execution_time: SAMPLE_MAX_EXECUTION_SEC, + }, + ); +} 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);