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)
  ↓
Result added to messages as toolResult
  ↓
bedrock_runtime.converse()  ← Generate answer based on result
  ↓
(If needed, generate additional SQL → 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
      └── ... (grows 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 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 the Converse API
    2. If stop_reason is "tool_use" → extract SQL → execute on Athena → add result to messages → resend
    3. If stop_reason is "end_turn" → return text answer
    4. Force exit when max iterations are reached
    """
    system_prompt = load_system_prompt()
    messages = _build_messages(history, user_message)

    for iteration in range(MAX_ITERATIONS):
        # ── Call Converse API ──
        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 → done ──
        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 result as toolResult to the conversation ──
        messages = messages + [
            assistant_message,  # Claude's assistant message (including tool_use)
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": tool_use_id,
                            "content": [
                                {"text": json.dumps(result, ensure_ascii=False)}
                            ],
                        }
                    }
                ],
            },
        ]
        # Call converse() again in the next iteration of the loop

    return "Maximum number of iterations reached."

def _extract_tool_use(message: dict) -> dict | None:
    """Extracts 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]:
    """Converts conversation history to the Converse API message format."""
    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 top of loop                  │
│                                                    │
└────────────────────────────────────────────────────┘

You are a data analyst for [domain].

## Answer Rules
1. Quote exact figures from the search results when answering
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 Rules
- Syntax: Presto SQL (Athena-compatible)
- Wrap Japanese column names in double quotes (e.g., "revenue")
- 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 JPY) | — |
| `revenue_actual` | BIGINT | Actual revenue (millions JPY) | — |

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

What is 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'