I Built an Agent That Automatically Generates and Executes SQL from Natural Language Using Amazon Athena × Bedrock Converse API

User's question
  ↓
Application server (FastAPI, etc.)
  ↓
bedrock_runtime.converse()  ← SQL tool defined via tool_use
  ↓
Claude generates SQL (stop_reason: "tool_use")
  ↓
App receives SQL and executes directly on Athena (no Lambda needed)
  ↓
Results added to messages as toolResult
  ↓
bedrock_runtime.converse()  ← Generate answer based on results
  ↓
(Generate additional SQL if needed → re-execute → loop)
  ↓
stop_reason: "end_turn" → answer complete

import pandas as pd

df = pd.read_excel("data.xlsx", sheet_name="Sheet1")
df.to_parquet("output/data.parquet", engine="pyarrow", compression="snappy")

s3://your-bucket/sql/
  ├── sales/
  │   └── sales.parquet
  ├── employees/
  │   └── employees.parquet
  └── products/
      └── products.parquet

s3://your-bucket/
  ├── sql/                  ← Source data (Parquet)
  │   ├── sales/
  │   └── employees/
  └── athena-results/       ← Auto-generated by Athena (query results)
      ├── <query-id-1>.csv
      ├── <query-id-1>.csv.metadata
      ├── <query-id-2>.csv
      ├── <query-id-2>.csv.metadata
      └── ...（increases with each query execution）

-- Query 1: Create database
CREATE DATABASE IF NOT EXISTS your_database
COMMENT 'Structured data for SQL agent';

-- Query 2: Example table: sales data
CREATE EXTERNAL TABLE your_database.sales (
  `region`         STRING   COMMENT 'Region name',
  `fiscal_year`    INT      COMMENT 'Fiscal year',
  `quarter`        STRING   COMMENT 'Quarter (1Q, 2Q, 3Q, 4Q)',
  `revenue_plan`   BIGINT   COMMENT 'Planned revenue',
  `revenue_actual`  BIGINT  COMMENT 'Actual revenue'
)
STORED AS PARQUET
LOCATION 's3://your-bucket/sql/sales/'
TBLPROPERTIES ('parquet.compression'='SNAPPY');

SELECT * FROM your_database.sales LIMIT 10;

SELECT COUNT(*), MIN(fiscal_year), MAX(fiscal_year) FROM your_database.sales;

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "athena:StartQueryExecution",
                "athena:GetQueryExecution",
                "athena:GetQueryResults",
                "athena:StopQueryExecution",
                "athena:GetWorkGroup"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "glue:GetDatabase",
                "glue:GetTable",
                "glue:GetTables",
                "glue:GetPartition",
                "glue:GetPartitions"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetBucketLocation",
                "s3:ListBucket"
            ],
            "Resource": "arn:aws:s3:::your-bucket"
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:PutObject",
                "s3:GetObject"
            ],
            "Resource": "arn:aws:s3:::your-bucket/*"
        }
    ]
}

"""athena.py — Athena query execution module"""

import asyncio
import re
import time
import boto3

MAX_ROWS = 200
POLL_INTERVAL = 0.5
MAX_WAIT = 30

ATHENA_DATABASE = "your_database"
ATHENA_WORKGROUP = "primary"
ATHENA_OUTPUT_S3 = "s3://your-bucket/athena-results/"

_SELECT_ONLY = re.compile(r"^\s*SELECT\b", re.IGNORECASE)

def _validate_sql(sql: str) -> None:
    """Reject anything other than SELECT statements."""
    if not _SELECT_ONLY.match(sql):
        raise ValueError(
            f"Only SELECT statements are allowed. Received: {sql[:80]!r}"
        )

def _run_query_sync(sql: str) -> dict:
    """Execute a query on Athena and return the results (synchronous)."""
    athena = boto3.client("athena")

    response = athena.start_query_execution(
        QueryString=sql,
        QueryExecutionContext={"Database": ATHENA_DATABASE},
        WorkGroup=ATHENA_WORKGROUP,
        ResultConfiguration={"OutputLocation": ATHENA_OUTPUT_S3},
    )
    execution_id = response["QueryExecutionId"]

    # Poll for completion
    deadline = time.time() + MAX_WAIT
    while time.time() < deadline:
        status = athena.get_query_execution(QueryExecutionId=execution_id)
        state = status["QueryExecution"]["Status"]["State"]
        if state == "SUCCEEDED":
            break
        if state in ("FAILED", "CANCELLED"):
            reason = status["QueryExecution"]["Status"].get(
                "StateChangeReason", "unknown"
            )
            raise RuntimeError(f"Athena query {state}: {reason}")
        time.sleep(POLL_INTERVAL)
    else:
        raise TimeoutError(
            f"Athena query did not complete within {MAX_WAIT}s"
        )

    # Retrieve results
    paginator = athena.get_paginator("get_query_results")
    pages = paginator.paginate(QueryExecutionId=execution_id)

    rows: list[list] = []
    columns: list[str] = []
    for page in pages:
        result = page["ResultSet"]
        if not columns:
            columns = [
                c["Label"]
                for c in result["ResultSetMetadata"]["ColumnInfo"]
            ]
        for row in result["Rows"][1 if not rows else 0:]:
            values = [d.get("VarCharValue") for d in row["Data"]]
            rows.append(values)
            if len(rows) >= MAX_ROWS:
                return {"columns": columns, "rows": rows, "row_count": len(rows)}

    return {"columns": columns, "rows": rows, "row_count": len(rows)}

async def execute_query(sql: str) -> dict:
    """Async wrapper. Executes synchronous processing in a thread pool."""
    _validate_sql(sql)
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _run_query_sync, sql)

EXECUTE_SQL_TOOL = {
    "toolSpec": {
        "name": "execute_sql_query",
        "description": (
            "Executes a Presto SQL SELECT query on Athena and returns the results as JSON."
        ),
        "inputSchema": {
            "json": {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The Presto SQL SELECT statement to execute",
                    }
                },
                "required": ["sql"],
            }
        },
    }
}

"""agent_loop.py — SQL agent loop using the Converse API"""

import json
import boto3
from athena import execute_query

MODEL_ID = "ap-northeast-1.anthropic.claude-sonnet-4-6"
MAX_ITERATIONS = 4  # Maximum number of SQL execution iterations

bedrock = boto3.client("bedrock-runtime", region_name="ap-northeast-1")

# Tool definition (created in Step 4)
TOOL_CONFIG = {"tools": [EXECUTE_SQL_TOOL]}

def load_system_prompt() -> str:
    with open("system_prompt.md", encoding="utf-8") as f:
        return f.read()

async def run_agent(user_message: str, history: list[dict]) -> str:
    """
    Runs the agent loop for a user's question and returns the final answer.

    Flow:
    1. Send messages + tool definition to Converse API
    2. If stop_reason is "tool_use" → extract SQL → execute on Athena → add results to messages → re-send
    3. If stop_reason is "end_turn" → return text answer
    4. Force termination if max iteration count is reached
    """
    system_prompt = load_system_prompt()
    messages = _build_messages(history, user_message)

    for iteration in range(MAX_ITERATIONS):
        # ── Converse API call ──
        response = bedrock.converse(
            modelId=MODEL_ID,
            system=[{"text": system_prompt}],
            messages=messages,
            toolConfig=TOOL_CONFIG,
            inferenceConfig={"maxTokens": 4096, "temperature": 0},
        )

        assistant_message = response["output"]["message"]
        stop_reason = response["stopReason"]

        # ── If Claude returns a text answer → complete ──
        if stop_reason != "tool_use":
            for block in assistant_message.get("content", []):
                if "text" in block:
                    return block["text"]
            return ""

        # ── If Claude returns a tool call ──
        tool_use_block = _extract_tool_use(assistant_message)
        if not tool_use_block:
            return "Failed to generate SQL."

        sql = tool_use_block["input"]["sql"]
        tool_use_id = tool_use_block["toolUseId"]

        print(f"[iteration {iteration + 1}] SQL: {sql}")

        # ── Execute SQL on Athena ──
        try:
            result = await execute_query(sql)
        except (ValueError, RuntimeError, TimeoutError) as exc:
            result = {"error": str(exc)}

        # ── Add results to the conversation as toolResult ──
        messages = messages + [
            assistant_message,  # Claude's assistant message (containing tool_use)
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": tool_use_id,
                            "content": [
                                {"text": json.dumps(result, ensure_ascii=False)}
                            ],
                        }
                    }
                ],
            },
        ]
        # Call converse() again in the next loop iteration

    return "Maximum iteration count reached."

def _extract_tool_use(message: dict) -> dict | None:
    """Extract the execute_sql_query toolUse block from an assistant message."""
    for block in message.get("content", []):
        if "toolUse" in block and block["toolUse"]["name"] == "execute_sql_query":
            return block["toolUse"]
    return None

def _build_messages(history: list[dict], current: str) -> list[dict]:
    """Convert conversation history to the message format for the Converse API."""
    messages = []
    for turn in history:
        role = turn.get("role")
        content = turn.get("content", "")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": [{"text": content}]})
    messages.append({"role": "user", "content": [{"text": current}]})
    return messages

messages = [user_message]

┌─── Loop start (max 4 times) ───────────────────────┐
│                                                    │
│  response = converse(messages, toolConfig)          │
│                                                    │
│  stop_reason == "end_turn"?  ──→ Return text answer │
│        │ No                                        │
│        ↓                                           │
│  stop_reason == "tool_use"                         │
│        │                                           │
│        ↓                                           │
│  Extract SQL → Execute on Athena                    │
│        │                                           │
│        ↓                                           │
│  messages += [assistant_msg, toolResult_msg]        │
│        │                                           │
│        └──→ Return to loop start                   │
│                                                    │
└────────────────────────────────────────────────────┘

You are a data analyst for [domain].

## Answer Rules
1. Cite exact figures from search results in your answers
2. Always include units with numerical values
3. If data is not found, respond with "No matching data was found"
4. Do not make assumptions or estimates

## SQL Query Writing Rules
- Syntax: Presto SQL (Athena compatible)
- Wrap Japanese column names in double quotes (e.g., "売上高")
- Use single quotes for string values (e.g., '1Q')

## Table Definitions

### Table name: sales
| Column | Type | Description | Example values |
|--------|-----|------|--------|
| `region` | STRING | Region name | Tokyo, Osaka |
| `fiscal_year` | INT | Fiscal year | 2021–2026 |
| `quarter` | STRING | Quarter | 1Q, 2Q, 3Q, 4Q, Full year |
| `revenue_plan` | BIGINT | Planned revenue (millions of yen) | — |
| `revenue_actual` | BIGINT | Actual revenue (millions of yen) | — |

## Data Notes
- [Filter conditions to avoid double-counting during aggregation]
- [Constraints on specific column value combinations]
- [Fiscal years/periods with missing data]
- [Rules for determining ambiguous terminology]

Please tell me the total actual sales across all regions for Q1 FY2025.

[iteration 1] SQL: SELECT SUM("revenue_actual") AS total FROM sales WHERE "fiscal_year" = 2025 AND "quarter" = '1Q'