When Your Data Lake Gets a Brain: Agentic AI Querying with Amazon Bedrock Title Querying Your Data Lake with Amazon Bedrock Agents: An Architecture Deep Dive

┌─────────────────────────────────────────────────────────────────┐
│ User / Application │
└────────────────────────────┬────────────────────────────────────┘
│ Natural language question
▼
┌─────────────────────────────────────────────────────────────────┐
│ Amazon Bedrock Agent │
│ (Claude 3.5 Sonnet / Haiku) │
│ │
│ Orchestration Layer — decides which tools to call, in what │
│ order, and how to interpret results before responding │
└──────┬─────────────────┬──────────────────┬─────────────────────┘
│ │ │
▼ ▼ ▼
┌────────────┐ ┌──────────────┐ ┌──────────────────┐
│ Tool 1 │ │ Tool 2 │ │ Tool 3 │
│ Schema │ │ Query │ │ Result │
│ Lookup │ │ Builder │ │ Interpreter │
│ │ │ │ │ │
│ Glue Data │ │ Constructs + │ │ Formats + │
│ Catalog │ │ validates │ │ summarizes │
│ API calls │ │ Athena SQL │ │ query output │
└─────┬──────┘ └──────┬───────┘ └──────┬───────────┘
│ │ │
▼ ▼ │
┌───────────┐ ┌─────────────┐ │
│ AWS Glue │ │ Amazon │ │
│ Data │ │ Athena │ │
│ Catalog │ │ │ │
└───────────┘ └──────┬──────┘ │
│ │
▼ │
┌─────────────┐ │
│ Amazon S3 │ │
│ Data Lake ├───────────┘
└─────────────┘

import boto3
import json

glue = boto3.client('glue')

def get_schema_context(database: str, question_keywords: list[str]) -> dict:
"""
Returns relevant table schemas from the Glue Data Catalog
based on keyword relevance to the user's question.
"""
tables_response = glue.get_tables(DatabaseName=database)
tables = tables_response['TableList']

relevant_tables = []

for table in tables:
table_name = table['Name']
columns = table['StorageDescriptor']['Columns']
partition_keys = table.get('PartitionKeys', [])

# Score relevance based on table name + column name overlap with question
score = sum(
1 for kw in question_keywords
if kw.lower() in table_name.lower() or
any(kw.lower() in col['Name'].lower() for col in columns)
)

if score > 0:
relevant_tables.append({
'table_name': table_name,
'columns': [
{'name': col['Name'], 'type': col['Type']}
for col in columns
],
'partition_keys': [
{'name': pk['Name'], 'type': pk['Type']}
for pk in partition_keys
],
'location': table['StorageDescriptor']['Location'],
'relevance_score': score
})

return {
'database': database,
'relevant_tables': sorted(
relevant_tables,
key=lambda x: x['relevance_score'],
reverse=True
)[:5] # Cap at top 5 to stay within context limits
}

import boto3
import time
import re

athena = boto3.client('athena')
S3_OUTPUT = 's3://your-athena-results-bucket/query-results/'

def validate_query(sql: str, available_tables: list[str]) -> dict:
"""
Pre-execution validation: check referenced tables exist,
no destructive operations, partition filter present on large tables.
"""
sql_upper = sql.upper().strip()

# Block destructive operations
destructive = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'ALTER', 'TRUNCATE']
for keyword in destructive:
if re.search(rf'\b{keyword}\b', sql_upper):
return {'valid': False, 'reason': f'Destructive operation detected: {keyword}'}

# Ensure query is SELECT only
if not sql_upper.startswith('SELECT'):
return {'valid': False, 'reason': 'Only SELECT queries are permitted'}

# Check referenced tables exist in catalog
referenced = re.findall(r'FROM\s+(\w+)', sql_upper)
referenced += re.findall(r'JOIN\s+(\w+)', sql_upper)
for ref in referenced:
if ref.lower() not in [t.lower() for t in available_tables]:
return {'valid': False, 'reason': f'Table not found in catalog: {ref}'}

return {'valid': True}

def execute_athena_query(sql: str, database: str) -> dict:
"""
Execute validated SQL against Athena, poll for completion,
return results with column headers.
"""
response = athena.start_query_execution(
QueryString=sql,
QueryExecutionContext={'Database': database},
ResultConfiguration={'OutputLocation': S3_OUTPUT},
WorkGroup='data-agent-workgroup' # separate workgroup for cost tracking
)

execution_id = response['QueryExecutionId']

# Poll with backoff
for attempt in range(20):
status = athena.get_query_execution(QueryExecutionId=execution_id)
state = status['QueryExecution']['Status']['State']

if state == 'SUCCEEDED':
break
elif state in ['FAILED', 'CANCELLED']:
reason = status['QueryExecution']['Status'].get('StateChangeReason', 'Unknown')
return {'success': False, 'error': reason, 'query': sql}

time.sleep(min(2 ** attempt * 0.5, 10)) # exponential backoff, max 10s

# Fetch results
results = athena.get_query_results(QueryExecutionId=execution_id)
rows = results['ResultSet']['Rows']

if not rows:
return {'success': True, 'data': [], 'row_count': 0}

headers = [col['VarCharValue'] for col in rows[0]['Data']]
data = [
{headers[i]: col.get('VarCharValue', None)
for i, col in enumerate(row['Data'])}
for row in rows[1:] # Skip header row
]

# Cost metadata
stats = status['QueryExecution']['Statistics']
data_scanned_mb = stats.get('DataScannedInBytes', 0) / (1024 * 1024)

return {
'success': True,
'data': data[:100], # Cap rows returned to agent context
'row_count': len(data),
'data_scanned_mb': round(data_scanned_mb, 2),
'execution_id': execution_id
}

import boto3

bedrock_agent = boto3.client('bedrock-agent')

agent_instruction = """
You are a data analyst assistant with direct access to a production data lake.

Your capabilities:
1. Look up available tables and their schemas from the Glue Data Catalog
2. Write and execute SQL queries against Amazon Athena
3. Interpret query results and answer the user's question in plain language

How to operate:
- Always start by checking schema context before writing a query
- Use partition filters wherever possible to minimize data scanned
- If a query returns no results, tell the user why (wrong date range, filter too narrow, etc.)
- Never assume column names — always verify against the schema first
- Report the data scanned in MB after each query so the user is aware of cost

What you must NOT do:
- Run queries without partition filters on tables larger than 10GB
- Return raw query dumps — always summarize and interpret the data
- Make up column names or table names that weren't in the schema lookup
"""

response = bedrock_agent.create_agent(
agentName='data-lake-analyst',
agentResourceRoleArn='arn:aws:iam::ACCOUNT_ID:role/BedrockAgentRole',
foundationModel='anthropic.claude-3-5-sonnet-20241022-v2:0',
instruction=agent_instruction,
idleSessionTTLInSeconds=1800
)

{
"name": "execute_data_lake_query",
"description": "Execute a validated SQL query against the Athena data lake and return results",
"parameters": {
"sql": {
"description": "The SELECT SQL query to execute. Must include partition filters where applicable.",
"required": true,
"type": "string"
},
"database": {
"description": "The Glue database name to query against",
"required": true,
"type": "string"
},
"explain_intent": {
"description": "Brief description of what this query is intended to answer",
"required": true,
"type": "string"
}
}
}

Step 1 → TOOL: schema_lookup
Input: keywords=["product", "category", "revenue", "order"]
Output: Returns `orders` table (order_id, product_id, amount, order_date)
and `products` table (product_id, category, name)
Partition key on `orders`: year, month

Step 2 → REASONING:
"I need to join orders + products on product_id,
filter year IN (2023, 2024) AND month IN (10,11,12),
group by category and year, sum amount, rank by revenue"

Step 3 → TOOL: execute_data_lake_query
SQL:
WITH ranked AS (
SELECT
p.category,
YEAR(o.order_date) AS year,
SUM(o.amount) AS revenue,
RANK() OVER (PARTITION BY YEAR(o.order_date) ORDER BY SUM(o.amount) DESC) AS rank
FROM orders o
JOIN products p ON o.product_id = p.product_id
WHERE o.year IN (2023, 2024)
AND o.month IN (10, 11, 12)
GROUP BY p.category, YEAR(o.order_date)
)
SELECT * FROM ranked WHERE rank <= 5
ORDER BY year DESC, rank ASC

Data scanned: 1.34 MB ✓ (partition filter worked)

Step 4 → REASONING + RESPONSE:
Compares 2023 vs 2024 top 5, notes category rank changes,
surfaces any new entrants or dropouts, formats as readable summary