run_query

Instructions

Implementation Reference

• src/index.ts:44-72 (registration)

  Registration of the 'run_query' tool in the MCP ListTools response, including name, description, and input schema.

  {
    name: "run_query",
    description: "Execute a SQL query using AWS Athena. Returns full results if query completes before timeout, otherwise returns queryExecutionId.",
    inputSchema: {
      type: "object",
      properties: {
        database: {
          type: "string",
          description: "The Athena database to query",
        },
        query: {
          type: "string",
          description: "SQL query to execute",
        },
        maxRows: {
          type: "number",
          description: "Maximum number of rows to return (default: 1000)",
          minimum: 1,
          maximum: 10000,
        },
        timeoutMs: {
          type: "number",
          description: "Timeout in milliseconds (default: 60000)",
          minimum: 1000,
        },
      },
      required: ["database", "query"],
    },
  },

• src/index.ts:150-177 (handler)

  MCP CallToolRequest handler case for 'run_query': validates arguments, prepares QueryInput, invokes AthenaService.executeQuery, and returns JSON-formatted result.

  case "run_query": {
    if (!request.params.arguments ||
        typeof request.params.arguments.database !== 'string' ||
        typeof request.params.arguments.query !== 'string') {
      throw new McpError(
        ErrorCode.InvalidParams,
        "Missing or invalid required parameters: database (string) and query (string)"
      );
    }
  
    const queryInput: QueryInput = {
      database: request.params.arguments.database,
      query: request.params.arguments.query,
      maxRows: typeof request.params.arguments.maxRows === 'number' ?
        request.params.arguments.maxRows : undefined,
      timeoutMs: typeof request.params.arguments.timeoutMs === 'number' ?
        request.params.arguments.timeoutMs : undefined,
    };
    const result = await this.athenaService.executeQuery(queryInput);
    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(result, null, 2),
        },
      ],
    };
  }

• src/athena.ts:38-90 (handler)

  Core execution logic for run_query: starts Athena query, waits/polls for completion with configurable timeout, fetches results if complete or returns queryExecutionId on timeout.

  async executeQuery(input: QueryInput): Promise<QueryResult | { queryExecutionId: string }> {
    try {
      // Start query execution
      const startResponse = await this.client.send(
        new StartQueryExecutionCommand({
          QueryString: input.query,
          QueryExecutionContext: {
            Database: input.database,
          },
          ResultConfiguration: {
            OutputLocation: this.outputLocation,
          },
          ...(this.workGroup && { WorkGroup: this.workGroup })
        })
      );
  
      if (!startResponse.QueryExecutionId) {
        throw new Error("Failed to start query execution");
      }
  
      const timeoutMs = input.timeoutMs || 60000; // Default 60 second timeout
      const startTime = Date.now();
  
      try {
        // Wait for query completion or timeout
        const queryExecution = await this.waitForQueryCompletion(
          startResponse.QueryExecutionId,
          100,
          timeoutMs
        );
  
        // If we got here, query completed before timeout
        return await this.getQueryResults(startResponse.QueryExecutionId, input.maxRows);
      } catch (error) {
        if (error && typeof error === "object" && "code" in error) {
          const athenaError = error as AthenaError;
          if (athenaError.code === "TIMEOUT") {
            // Return just the execution ID on timeout
            return { queryExecutionId: startResponse.QueryExecutionId };
          }
        }
        throw error;
      }
    } catch (error) {
      if (error instanceof InvalidRequestException) {
        throw {
          message: error.message,
          code: "INVALID_REQUEST",
        };
      }
      throw error;
    }
  }

• src/types.ts:1-6 (schema)

  TypeScript interface QueryInput used for input validation and typing in the run_query tool.

  export interface QueryInput {
    database: string;
    query: string;
    maxRows?: number;
    timeoutMs?: number;
  }