Merge pull request #7 from railsware/CPL-20411/use-execution-id

[CPL-20411] use execution ID

Author: oshchyhol

README.md

@@ -43,12 +43,14 @@ Remove this line, if you're offline or if you specifically want to use the image
 
 ## Tools
 ### Data flows
-- **get-data** - Gets the result of a data flow run as a SQLite file and executes a read-only query on it.
+- **get-data** - Gets the result of a data flow run as a SQLite file and executes a read-only query on it. Currently, only data flows built from a dashboard or dataset template are supported.
   - `dataflowId`: Data flow ID (`string`, **required**)
+  - `executionId`: Data flow run ID (`string`, **required**)
   - `query`: Query to run on the data flow SQLite file (`string`, **required**)
 
-- **get-schema** - Gets the data flow schema file.
+- **get-schema** - Gets the data flow schema file. Currently, only data flows built from a dashboard or dataset template are supported.
   - `dataflowId`: Data flow ID (`string`, **required**)
+  - `executionId`: Data flow run ID (`string`, **required**)
 
 ## Development
 

package-lock.json

@@ -16,12 +16,15 @@
   "dependencies": {
     "@modelcontextprotocol/sdk": "1.11.4",
     "better-sqlite3": "11.10.0",
+    "json-schema-to-zod": "2.6.1",
     "lodash": "4.17.21",
     "pino": "9.7.0",
     "tsx": "4.19.4",
     "url-template": "3.1.1",
     "znv": "0.5.0",
-    "zod": "3.24.4"
+    "zod": "3.24.4",
+    "zod-to-json-schema": "3.24.5",
+    "zod-validation-error": "3.4.1"
   },
   "devDependencies": {
     "@eslint/js": "9.27.0",

package.json

@@ -0,0 +1,4 @@
+export interface SignedUrlDto {
+  file: 'sqlite' | 'schema',
+  signed_url: string,
+}

src/lib/couplerio-client/dataflows/signed_url.ts

@@ -50,6 +50,7 @@ describe('getData', () => {
 
     const toolResult = await handler({
       dataflowId: 'test-dataflow-id',
+      executionId: 'test-execution-id',
       query: 'SELECT * FROM data'
     })
 
@@ -69,6 +70,7 @@ describe('getData', () => {
 
     const toolResult = await handler({
       dataflowId: 'test-dataflow-id',
+      executionId: 'test-execution-id',
       query: 'SELECT * FROM data'
     })
 
@@ -89,6 +91,7 @@ describe('getData', () => {
 
     const toolResult = await handler({
       dataflowId: 'test-dataflow-id',
+      executionId: 'test-execution-id',
       query: 'SELECT * FROM data'
     })
 
@@ -108,6 +111,7 @@ describe('getData', () => {
 
     const toolResult = await handler({
       dataflowId: 'test-dataflow-id',
+      executionId: 'test-execution-id',
       query: 'SELECT INVALID SQL QUERY'
     })
 
@@ -129,7 +133,7 @@ describe('with invalid params', () => {
       isError: true,
       content: [{
         type: 'text',
-        text: 'This tool requires parameters'
+        text: 'Invalid parameters for get-data tool. Validation error: Required'
       }]
     })
   })
@@ -141,34 +145,35 @@ describe('with invalid params', () => {
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: dataflowId must be a non-empty string'
+        text: 'Invalid parameters for get-data tool. Validation error: Required at "dataflowId"; Required at "executionId"'
       }]
     })
   })
 
   it('returns error on missing query', async () => {
-    const toolResult = await handler({ dataflowId: 'test-dataflow-id' })
+    const toolResult = await handler({ dataflowId: 'test-dataflow-id', executionId: 'test-execution-id' })
 
     expect(toolResult).toEqual({
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: query must be a non-empty string'
+        text: 'Invalid parameters for get-data tool. Validation error: Required at "query"'
       }]
     })
   })
 
   it('returns error on invalid dataflowId', async () => {
     const toolResult = await handler({
       dataflowId: 123,
+      executionId: true,
       query: 'SELECT * FROM data'
     })
 
     expect(toolResult).toEqual({
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: dataflowId must be a non-empty string'
+        text: 'Invalid parameters for get-data tool. Validation error: Expected string, received number at "dataflowId"; Expected string, received boolean at "executionId"'
       }]
     })
   })
@@ -183,22 +188,23 @@ describe('with invalid params', () => {
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: query must be a non-empty string'
+        text: 'Invalid parameters for get-data tool. Validation error: Required at "executionId"; Expected string, received number at "query"'
       }]
     })
   })
 
   it('returns error on non-SELECT query', async () => {
     const toolResult = await handler({
       dataflowId: 'test-dataflow-id',
+      executionId: 'test-execution-id',
       query: 'INSERT INTO data (col_0, col_1) VALUES (2, "Test 2")'
     })
 
     expect(toolResult).toEqual({
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: query must start with SELECT'
+        text: 'Invalid parameters for get-data tool. Validation error: must start with "SELECT" at "query"'
       }]
     })
   })

src/tools/get-data/handler.test.ts

@@ -1,37 +1,39 @@
 import Database from 'better-sqlite3'
 import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
+import { fromError } from 'zod-validation-error'
 
+import { logger } from '@/logger'
 import { textResponse } from '@/util/tool-response'
 import { FileManager } from '@/tools/shared/file-manager'
 
-import { validateParams } from './validate-params'
+import { zodSchema } from './input-schema'
 
 export const handler = async (params?: Record<string, unknown>): Promise<CallToolResult> => {
-  try {
-    validateParams(params)
-  } catch (e) {
+  const validationResult = zodSchema.safeParse(params)
+
+  if (!validationResult.success) {
+    const error = fromError(validationResult.error)
+    logger.error(`Invalid parameters for get-data tool: ${error.toString()}`)
+
     return textResponse({
-      text: (e as Error).message,
+      text: `Invalid parameters for get-data tool. ${error.toString()}`,
       isError: true,
     })
   }
 
-  const dataflowId = params!.dataflowId as string
-  const query = params!.query as string
-
-  const fileManager = new FileManager({ dataflowId })
+  const fileManager = new FileManager(validationResult.data)
 
   let sqlitePath: string
   try {
     sqlitePath = await fileManager.getFile('sqlite')
   } catch (e) {
-    return textResponse({ text: `Failed to get data flow ${dataflowId} sqlite file. ${e}`, isError: true })
+    return textResponse({ text: `Failed to get data flow ${validationResult.data.dataflowId} sqlite file. ${e}`, isError: true })
   }
 
   const db = new Database(sqlitePath)
   let statement, queryResult
   try {
-    statement = db.prepare(query)
+    statement = db.prepare(validationResult.data.query)
     queryResult = statement.all()
   } catch (e) {
     return textResponse({ text: `Failed to execute query: ${e}`, isError: true })

src/tools/get-data/handler.ts

@@ -1,16 +1,19 @@
-export const inputSchema = {
-  type: 'object',
-  properties: {
-    dataflowId: {
-      type: 'string',
-      description: 'The ID of the data flow with a successful run',
-      pattern: '^\\S+'
-    },
-    query: {
-      type: 'string',
-      description: 'The SQL query to run on the data flow sqlite file',
-      pattern: '^SELECT.*?'
-    }
-  },
-  required: ['dataflowId', 'query'],
-}
+import { z } from 'zod'
+import { zodToJsonSchema } from 'zod-to-json-schema'
+
+export const zodSchema = z.object({
+  dataflowId: z.string()
+    .min(1, 'dataflowId is required')
+    .regex(/^\S+$/, 'dataflowId must not contain whitespace')
+    .describe('The ID of the data flow with a successful run'),
+  executionId: z.string()
+    .min(1, 'executionId is required')
+    .regex(/^\S+$/, 'executionId must be a non-empty string')
+    .describe('The ID of the last successful run (execution) of the data flow.'),
+  query: z.string()
+    .min(1, 'query is required')
+    .regex(/^SELECT.*?/, 'must start with "SELECT"')
+    .describe('The SQL query to run on the data flow sqlite file.'),
+}).strict()
+
+export const inputSchema = zodToJsonSchema(zodSchema)

src/tools/get-data/input-schema.ts

@@ -18,7 +18,7 @@ const mockSchema = {
 // Response mocks
 const mockCreateSignedUrl = createMockResponse(
   async () => new Response(
-    JSON.stringify({ file: 'schema', url: 'https://example.com/schema.json' })
+    JSON.stringify({ file: 'schema', signed_url: 'https://example.com/schema.json' })
   )
 )
 const mockGetSchemaFile = createMockResponse(
@@ -48,7 +48,7 @@ describe('getSchema', () => {
       .mockImplementationOnce(mockCreateSignedUrl)
       .mockImplementationOnce(mockGetSchemaFile)
 
-    const toolResult = await handler({ dataflowId: 'test-dataflow-id' })
+    const toolResult = await handler({ dataflowId: 'test-dataflow-id', executionId: 'test-execution-id' })
 
     expect(toolResult).toEqual({
       isError: false,
@@ -70,7 +70,7 @@ describe('getSchema', () => {
         .mockImplementationOnce(mockCreateSignedUrlError)
         .mockImplementationOnce(mockGetSchemaFile)
 
-      const toolResult = await handler({ dataflowId: 'test-dataflow-id' })
+      const toolResult = await handler({ dataflowId: 'test-dataflow-id', executionId: 'test-execution-id' })
 
       expect(toolResult).toEqual({
         isError: true,
@@ -88,7 +88,7 @@ describe('getSchema', () => {
         .mockImplementationOnce(mockCreateSignedUrl)
         .mockImplementationOnce(mockGetSchemaFileError)
 
-      const toolResult = await handler({ dataflowId: 'test-dataflow-id' })
+      const toolResult = await handler({ dataflowId: 'test-dataflow-id', executionId: 'test-execution-id' })
 
       expect(toolResult).toEqual({
         isError: true,
@@ -109,7 +109,7 @@ describe('with invalid params', () => {
       isError: true,
       content: [{
         type: 'text',
-        text: 'This tool requires parameters',
+        text: 'Invalid parameters for get-schema tool. Validation error: Required',
       }]
     })
   })
@@ -121,19 +121,19 @@ describe('with invalid params', () => {
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: dataflowId must be a non-empty string',
+        text: 'Invalid parameters for get-schema tool. Validation error: Required at "dataflowId"; Required at "executionId"',
       }]
     })
   })
 
   it('returns error on invalid dataflowId', async () => {
-    const toolResult = await handler({ dataflowId: 123 })
+    const toolResult = await handler({ dataflowId: 123, executionId: true })
 
     expect(toolResult).toEqual({
       isError: true,
       content: [{
         type: 'text',
-        text: 'Missing or invalid required parameter: dataflowId must be a non-empty string',
+        text: 'Invalid parameters for get-schema tool. Validation error: Expected string, received number at "dataflowId"; Expected string, received boolean at "executionId"',
       }]
     })
   })