Building a Bedrock Agent That Queries Athena Data, On Behalf of the User
Build a Bedrock agent that queries Amazon Athena using IAM Identity Center trusted identity propagation, so Lake Formation enforces row-level access per user.
A regional manager and a sales rep both ask the same agent the exact same question: "show me revenue by region." The manager sees every region. The rep sees only their own. Same agent, same action group, same Athena query under the hood, two completely different answers, and nobody wrote a single
if user.role == "manager" line to make that happen.That's the pitch for wiring Amazon Athena into a Bedrock agent the right way instead of the easy way. The easy way gives the agent's Lambda function one IAM role, and every person who talks to that agent inherits whatever that role can see in the data lake. The right way makes the agent run each query as the actual signed-in person, so AWS Lake Formation's own fine-grained permissions do the enforcement they were already built to do. AWS calls this trusted identity propagation, and it runs on IAM Identity Center. It's the same idea as On-Behalf-Of passthrough on other platforms, just with
CreateTokenWithIAM doing the token exchange instead of Entra or Google OAuth.This is a hands-on build guide. We'll define a row-level data filter in Lake Formation, wire a Bedrock agent action group to query Athena through an identity-enhanced session, and run the agent as two different signed-in users to prove the permission boundary actually holds.
The mental model
Before touching code, it's worth separating two things people conflate: who the agent's Lambda function runs as, and who the agent runs for.
Who the Lambda runs as is a question of execution role. By default, a Lambda function backing a Bedrock agent action group has one IAM execution role, and every invocation of that function, regardless of which end user triggered it, authenticates to Athena and the Glue Data Catalog as that same role. That's fine for a tool with no per-user data boundary. It's a liability the moment two different users of the agent are supposed to see different rows of the same table.
Who the agent runs for is the question trusted identity propagation answers. Instead of the Lambda authenticating purely as its execution role, it exchanges the calling user's own identity, propagated from your identity provider through IAM Identity Center, for a short-lived, identity-enhanced AWS session. That session is used for the Athena call. If Lake Formation has a data filter on the table that restricts rows based on the caller's identity, that restriction applies exactly as if the person had opened Athena's query editor and run the SQL themselves.
The agent's action group code doesn't change based on who's asking. The tool definition is identical for every user. The only thing that changes per request is which identity gets attached to the session, and that part is handled by IAM Identity Center's token exchange, not by anything you write into the query.
Prerequisites
A few things need to be true before any of the code below will work:
- IAM Identity Center enabled for your AWS organization, with users and groups provisioned from your identity provider
- A customer-managed application registered in IAM Identity Center for your Bedrock agent's Lambda function, so it can call
CreateTokenWithIAM - At least one Lake Formation data filter defined on the table you want to query, so there's something for identity passthrough to actually enforce
- Lake Formation permissions granted to the relevant Identity Center users or groups for the catalog, database, and table
- An Athena workgroup and S3 results location the identity-enhanced session is permitted to use
- A real IdP token is required. Trusted identity propagation is for interactive, user-facing agents where a signed-in person is in the loop for every call. A background job running under a plain execution role gets none of this, and that's by design.
That last point is worth sitting with before you build anything. If your agent needs to run unattended, this entire pattern doesn't apply to that path, and you should query Athena directly under a scoped service role instead.
Step 1: Define the Lake Formation data filter
Everything downstream depends on Lake Formation actually having something to enforce. If the table grants full access to every user and group, identity propagation will faithfully carry the caller's identity all the way to a query that returns every row regardless of who's asking.
1
2
3
4
5
6
7
8
9
{
"Name": "sales-rep-region-filter",
"DatabaseName": "sales",
"TableName": "orders",
"RowFilter": {
"FilterExpression": "region = current_principal_tag('region')"
},
"ColumnWildcard": {}
}
1
2
3
4
5
6
7
aws lakeformation create-data-cells-filter --table-data \
file://sales-rep-region-filter.json
aws lakeformation grant-permissions \
--principal DataLakePrincipalIdentifier="arn:aws:identitystore:::user/sales-rep-group" \
--resource '{"DataCellsFilter":{"TableCatalogId":"111122223333","DatabaseName":"sales","TableName":"orders","Name":"sales-rep-region-filter"}}' \
--permissions "SELECT"
Once this is live, open Athena's query editor as two different Identity Center users and run the same
SELECT * FROM sales.orders to confirm the filter behaves the way you expect, before a Bedrock agent is anywhere near it. This step has nothing to do with Bedrock or Lambda. It's Lake Formation's own fine-grained access control, and it's the entire reason the rest of this works.Step 2: Register the application in IAM Identity Center
Trusted identity propagation needs an IAM Identity Center customer-managed application representing your agent, so the Lambda function can exchange a caller's identity for AWS credentials scoped to that specific person.
Set up the trust relationship through the console or API: register the application, define the IAM role it assumes, and grant that role
sts:AssumeRole permission with an identity context from Identity Center. The key piece is the ProvidedContexts parameter on AssumeRole, which is what carries the caller's identity into the resulting session.1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import boto3
def get_identity_enhanced_session(user_id_token, access_role_arn, idc_application_arn):
idc = boto3.client("sso-oidc", region_name="us-east-1")
token_response = idc.create_token_with_iam(
clientId=idc_application_arn,
grantType="urn:ietf:params:oauth:grant-type:jwt-bearer",
assertion=user_id_token,
)
identity_context = token_response["idToken"]
sts = boto3.client("sts")
creds = sts.assume_role(
RoleArn=access_role_arn,
RoleSessionName="bedrock-agent-athena-session",
ProvidedContexts=[{
"ProviderArn": "arn:aws:iam::aws:contextProvider/IdentityCenter",
"ContextAssertion": identity_context,
}],
)["Credentials"]
return creds
Nothing in this exchange touches Athena yet. It's purely about turning a caller's IdP token into AWS credentials that carry their workforce identity.
Step 3: Wire the Athena tool into the Bedrock agent action group
With the identity exchange in place, the Lambda function backing your agent's action group uses the resulting credentials for the actual Athena call instead of its own execution role.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import boto3
import time
def lambda_handler(event, context):
user_id_token = event["sessionAttributes"]["user_id_token"]
creds = get_identity_enhanced_session(
user_id_token,
access_role_arn="arn:aws:iam::111122223333:role/athena-tip-access-role",
idc_application_arn="arn:aws:sso::111122223333:application/ssoins-xxxx/apl-xxxx",
)
athena = boto3.client(
"athena",
aws_access_key_id=creds["AccessKeyId"],
aws_secret_access_key=creds["SecretAccessKey"],
aws_session_token=creds["SessionToken"],
)
query = athena.start_query_execution(
QueryString=event["parameters"]["sql"],
QueryExecutionContext={"Database": "sales"},
ResultConfiguration={"OutputLocation": "s3://my-athena-results/"},
WorkGroup="bedrock-agent-workgroup",
)
query_id = query["QueryExecutionId"]
while True:
status = athena.get_query_execution(QueryExecutionId=query_id)
state = status["QueryExecution"]["Status"]["State"]
if state in ("SUCCEEDED", "FAILED", "CANCELLED"):
break
time.sleep(1)
results = athena.get_query_results(QueryExecutionId=query_id)
return {"results": results["ResultSet"]["Rows"]}
The action group's schema and the agent's instructions stay exactly as they'd be for a plain, service-role Athena tool. The only addition is pulling the caller's identity out of the session attributes Bedrock passes into the Lambda invocation, and using it to mint a scoped session before the query runs.
Step 4: Run it, and run it as two different users
This is the step that actually proves trusted identity propagation works, and it's also the step people skip because testing everything under one broadly-permissioned developer identity is faster.
Invoke the agent as a user in a group with full regional access, and ask for revenue by region. You'll get every region back. Now invoke the same agent, same action group, as a user in the sales rep group scoped by the Lake Formation data filter from step 1. Ask the identical question. The Lambda code path is the same. The rows that come back only cover that person's region.
Nothing in the Lambda function or the agent's instructions mentions roles or regions. The restriction is enforced entirely inside Lake Formation, by the data cells filter, against the identity carried in the
ProvidedContexts of that specific caller's session.Here's the shape of that single call, end to end:
Where this fits in the bigger picture
Zoom out past a single request and the interesting part is what stays fixed. The Bedrock agent is the same agent, the action group's Athena tool is the same tool, regardless of who's using it. What changes per request is a token exchange that happens automatically, between IAM Identity Center and the Lambda function, before the query ever reaches Athena or Lake Formation.
That's the architectural payoff. You build the agent once. You don't maintain a permissions table inside your Lambda code that has to stay in sync with whatever's configured in Lake Formation, and you don't stand up a separate agent per role or region. The permissions model lives in exactly one place, which is also the place your data platform team already manages it with
create-data-cells-filter and grant-permissions.Production considerations
A handful of things are worth knowing before this goes past a demo.
Service roles and identity-enhanced sessions don't mix mid-request. Decide per tool call which mode you're in. A background job that refreshes a cached dashboard can run under the Lambda's own execution role. Anything an individual user is directly asking should run under their propagated identity, not a blend of both.
Data cells filters are table-specific and composable, but they add query planning overhead. A table with many overlapping filters across roles can see real query cost and latency increases. Test with representative data volumes and realistic filter counts before assuming the enforcement layer is free.
Token refresh is automatic, revocation still isn't.
CreateTokenWithIAM sessions expire on a short clock and get refreshed transparently while a user is active. That doesn't substitute for offboarding. If someone leaves a team or changes roles, their Lake Formation grants need to be updated directly, not assumed to lapse with a session token.Test with a genuinely scoped account, not your own. It's tempting to build and demo everything under an admin or data-platform-owner identity, where every query looks correct because nothing is actually filtered. The entire value of this pattern is invisible until you test with an account that Lake Formation actually restricts. Budget real time for that before calling the integration done.
CloudTrail already tracks this for you. Every Data API and Athena call made through an identity-enhanced session logs the propagated identity in the
onBehalfOf section of the CloudTrail event. That audit trail is close to free once trusted identity propagation is wired up, and it's worth pointing your security team at it early.Where this leaves you
The code to run an Athena query as the actual signed-in user is small: a token exchange, an
AssumeRole call with a provided context, and a data filter definition. The part worth taking seriously is everything that small piece of code is standing on: fine-grained access control Lake Formation has supported for years, an identity propagation model that means your Lambda function never has to encode a permissions table, and a guarantee that whoever's asking the question only ever sees what they were already allowed to see.If you're building an internal agent on top of data your organization already governs with Lake Formation, this is very likely less work than building and maintaining your own access control layer on top of a generic retrieval tool. Define the filter, register the identity exchange, and let Lake Formation keep doing the job it was already doing.
References
Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article