aws-athena-mcp

Instructions

Implementation Reference

• src/athena.ts:127-199 (handler)

  Core implementation of the get_result tool: validates query status, fetches results from AWS Athena using GetQueryResultsCommand, processes columns and rows into structured output.

  async getQueryResults(queryExecutionId: string, maxRows?: number): Promise<QueryResult> {
    try {
      // Check query state first
      const status = await this.getQueryStatus(queryExecutionId);
  
      if (status.state === QueryExecutionState.RUNNING || status.state === QueryExecutionState.QUEUED) {
        throw {
          message: "Query is still running",
          code: "QUERY_STILL_RUNNING",
          queryExecutionId,
        };
      }
  
      if (status.state === QueryExecutionState.FAILED) {
        throw {
          message: status.stateChangeReason || "Query failed",
          code: "QUERY_FAILED",
          queryExecutionId,
        };
      }
  
      if (status.state !== QueryExecutionState.SUCCEEDED) {
        throw {
          message: `Unexpected query state: ${status.state}`,
          code: "UNEXPECTED_STATE",
          queryExecutionId,
        };
      }
  
      const results = await this.client.send(
        new GetQueryResultsCommand({
          QueryExecutionId: queryExecutionId,
          MaxResults: maxRows || 1000,
        })
      );
  
      if (!results.ResultSet) {
        throw new Error("No results returned from query");
      }
  
      const columns = results.ResultSet.ResultSetMetadata?.ColumnInfo?.map(
        (col) => col.Name || ""
      ) || [];
  
      const rows = (results.ResultSet.Rows || [])
        .slice(status.substatementType === 'SELECT' ? 1 : 0) // Skip header row if query is SELECT
        .map((row) => {
          const obj: Record<string, unknown> = {};
          row.Data?.forEach((data, index) => {
            if (columns[index]) {
              obj[columns[index]] = data.VarCharValue;
            }
          });
          return obj;
        });
  
      return {
        columns,
        rows,
        queryExecutionId,
        bytesScanned: status.statistics?.dataScannedInBytes || 0,
        executionTime: status.statistics?.engineExecutionTimeInMillis || 0,
      };
    } catch (error) {
      if (error instanceof InvalidRequestException) {
        throw {
          message: "Query execution not found",
          code: "QUERY_NOT_FOUND",
        };
      }
      throw error;
    }
  }

• src/index.ts:179-202 (handler)

  MCP CallToolRequest handler for get_result: validates input parameters and calls AthenaService.getQueryResults, returns JSON response.

  case "get_result": {
    if (!request.params.arguments?.queryExecutionId ||
        typeof request.params.arguments.queryExecutionId !== 'string') {
      throw new McpError(
        ErrorCode.InvalidParams,
        "Missing or invalid required parameter: queryExecutionId (string)"
      );
    }
  
    const maxRows = typeof request.params.arguments.maxRows === 'number' ?
      request.params.arguments.maxRows : undefined;
    const result = await this.athenaService.getQueryResults(
      request.params.arguments.queryExecutionId,
      maxRows
    );
    return {
      content: [
        {
          type: "text",
          text: JSON.stringify(result, null, 2),
        },
      ],
    };
  }

• src/index.ts:73-92 (schema)

  Tool schema and registration: defines name, description, input schema with queryExecutionId (required) and optional maxRows for the get_result tool in ListTools response.

  {
    name: "get_result",
    description: "Get results for a completed query. Returns error if query is still running.",
    inputSchema: {
      type: "object",
      properties: {
        queryExecutionId: {
          type: "string",
          description: "The query execution ID",
        },
        maxRows: {
          type: "number",
          description: "Maximum number of rows to return (default: 1000)",
          minimum: 1,
          maximum: 10000,
        },
      },
      required: ["queryExecutionId"],
    },
  },

• src/types.ts:8-14 (schema)

  TypeScript interface defining the output structure for query results returned by get_result.

  export interface QueryResult {
    columns: string[];
    rows: Record<string, unknown>[];
    queryExecutionId: string;
    bytesScanned: number;
    executionTime: number;
  }