feat(context): register MCP sql execution tool

This commit is contained in:
Andrey Avtomonov 2026-05-14 17:56:33 +02:00
parent 06f020dca1
commit c774870346
3 changed files with 105 additions and 0 deletions

View file

@ -143,6 +143,12 @@ const scanArtifactReadSchema = z.object({
path: z.string().min(1),
});
const sqlExecutionSchema = z.object({
connectionId: connectionIdSchema,
sql: z.string().min(1),
maxRows: z.number().int().min(1).max(10_000).default(1000).optional(),
});
export function jsonToolResult<T extends object>(structuredContent: T): KtxMcpToolResult<T> {
return {
content: [{ type: 'text', text: JSON.stringify(structuredContent, null, 2) }],
@ -361,6 +367,34 @@ export function registerKtxContextTools(deps: RegisterKtxContextToolsDeps): void
);
}
if (ports.sqlExecution) {
const sqlExecution = ports.sqlExecution;
registerParsedTool(
server,
'sql_execution',
{
title: 'SQL Execution',
description:
'Execute one parser-validated read-only SQL query against a configured KTX connection and return structured rows.',
inputSchema: sqlExecutionSchema.shape,
},
sqlExecutionSchema,
async (input) => {
try {
return jsonToolResult(
await sqlExecution.execute({
connectionId: input.connectionId,
sql: input.sql,
maxRows: input.maxRows ?? 1000,
}),
);
} catch (error) {
return jsonErrorToolResult(error instanceof Error ? error.message : String(error));
}
},
);
}
if (ports.ingest) {
const ingest = ports.ingest;
registerParsedTool(

View file

@ -11,6 +11,8 @@ import type {
KtxMcpContextPorts,
KtxScanMcpPort,
KtxSemanticLayerMcpPort,
KtxSqlExecutionMcpPort,
KtxSqlExecutionResponse,
MemoryCapturePort,
} from './types.js';
@ -64,6 +66,63 @@ describe('createKtxMcpServer', () => {
});
});
it('registers parser-gated sql_execution when the host provides a SQL execution port', async () => {
const fake = makeFakeServer();
const response: KtxSqlExecutionResponse = {
headers: ['status', 'count'],
headerTypes: ['text', 'bigint'],
rows: [['paid', 42]],
rowCount: 1,
};
const sqlExecution: KtxSqlExecutionMcpPort = {
execute: vi.fn<KtxSqlExecutionMcpPort['execute']>().mockResolvedValue(response),
};
createKtxMcpServer({
server: fake.server,
userContext: { userId: 'local-user' },
contextTools: {
sqlExecution,
},
});
expect(fake.tools.map((tool) => tool.name)).toEqual(['sql_execution']);
await expect(
getTool(fake.tools, 'sql_execution').handler({
connectionId: 'warehouse',
sql: 'select status, count(*) from public.orders group by status',
maxRows: 50,
}),
).resolves.toEqual({
content: [
{
type: 'text',
text: JSON.stringify(
{
headers: ['status', 'count'],
headerTypes: ['text', 'bigint'],
rows: [['paid', 42]],
rowCount: 1,
},
null,
2,
),
},
],
structuredContent: {
headers: ['status', 'count'],
headerTypes: ['text', 'bigint'],
rows: [['paid', 42]],
rowCount: 1,
},
});
expect(sqlExecution.execute).toHaveBeenCalledWith({
connectionId: 'warehouse',
sql: 'select status, count(*) from public.orders group by status',
maxRows: 50,
});
});
it('registers memory capture tools without host app dependencies', async () => {
const fake = makeFakeServer();
const capture: MemoryCapturePort = {

View file

@ -312,10 +312,22 @@ export interface KtxScanMcpPort {
readArtifact?(input: { runId: string; path: string }): Promise<KtxScanArtifactReadResponse | null>;
}
export interface KtxSqlExecutionResponse {
headers: string[];
headerTypes?: string[];
rows: unknown[][];
rowCount: number;
}
export interface KtxSqlExecutionMcpPort {
execute(input: { connectionId: string; sql: string; maxRows: number }): Promise<KtxSqlExecutionResponse>;
}
export interface KtxMcpContextPorts {
connections?: KtxConnectionsMcpPort;
knowledge?: KtxKnowledgeMcpPort;
semanticLayer?: KtxSemanticLayerMcpPort;
sqlExecution?: KtxSqlExecutionMcpPort;
ingest?: KtxIngestMcpPort;
scan?: KtxScanMcpPort;
}