Query AWS Athena table using a Cross-Account IAM Role

aws-query-athena.md

Query Athena using a Cross-Account IAM Role

This script demonstrates how to query an AWS Athena table from a different AWS account by assuming an IAM Role. It uses the AWS SDK for JavaScript v3.

Description

The script performs the following actions:

1. Assumes a Cross-Account Role: It uses the AWS Security Token Service (STS) to obtain temporary security credentials by assuming an IAM role in a target AWS account. This is a secure way to grant temporary access to your AWS resources to users from other AWS accounts.
2. Executes an Athena Query: Using the temporary credentials, it connects to Amazon Athena and executes a SQL query against a specified database and table.
3. Polls for Results: It periodically checks the status of the query until it completes.
4. Fetches and Displays Results: Once the query is successful, it retrieves the results from the S3 output location, parses them into a JSON format, and prints them to the console.

Prerequisites

• Node.js and Bun (or npm/yarn) installed.
• AWS credentials (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY) for an IAM user with permission to assume the target role.
• An IAM Role in the target account that this script can assume. The role must have:
  • A trust relationship that allows the source IAM user to assume it.
  • An ExternalId condition in the trust policy for added security.
  • Permissions to run Athena queries and read the target table.
  • Permissions to write query results to the specified S3 bucket (s3:PutObject, s3:GetObject).

Installation

1. Clone the Gist or download the index.js and package.json files.
2. Install the dependencies:

   bun install

Configuration

The script is configured using environment variables. You can create a .env file in the root of the project or export them in your shell session.

# Credentials for the IAM user in the source account
AWS_ACCESS_KEY_ID="YOUR_SOURCE_ACCESS_KEY"
AWS_SECRET_ACCESS_KEY="YOUR_SOURCE_SECRET_KEY"
AWS_REGION="us-east-1"

# Configuration for the cross-account role to assume
AWS_ASSUME_ROLE_EXTERNAL_ID="YOUR_UNIQUE_EXTERNAL_ID"
CLIENT_ORG_ID="TARGET_AWS_ACCOUNT_ID"
CLIENT_ROLE_NAME="NAME_OF_THE_ROLE_TO_ASSUME"

# S3 Bucket for Athena query results in the target account
CLIENT_BUCKET="your-athena-results-bucket"
CLIENT_BUCKET_PATH="path/for/results"

# Athena configuration in the target account
CLIENT_ATHENA_DB="your_athena_database"
CLIENT_ATHENA_TABLE="your_athena_table"

How to Run

Once the dependencies are installed and the environment variables are set, run the script with:

bun start

The script will then print the query results to the console in JSON format.

Code Breakdown

getTemporaryCredentials()

This asynchronous function is responsible for the role assumption.

• It initializes the STSClient.
• It sends an AssumeRoleCommand with the RoleArn, a RoleSessionName, and the ExternalId.
• The RoleArn is constructed dynamically using the CLIENT_ORG_ID and CLIENT_ROLE_NAME environment variables.
• If successful, it returns the temporary accessKeyId, secretAccessKey, and sessionToken.

runAthenaQuery()

This is the main function that orchestrates the process.

1. It calls getTemporaryCredentials() to get the temporary credentials.
2. It initializes the AthenaClient using these temporary credentials. This ensures all subsequent Athena API calls are made with the assumed role's permissions.
3. It sends a StartQueryExecutionCommand. The SQL query is hardcoded in this example, but it could be made dynamic. It also specifies the output location for the results in S3.
4. It enters a while loop to poll the query's status using GetQueryExecutionCommand every 3 seconds.
5. If the query state becomes FAILED or CANCELLED, it throws an error.
6. Once the query SUCCEEDED, it fetches the data using GetQueryResultsCommand.
7. The results from Athena are in a specific ResultSet format. The script parses this format, using the first row as headers and subsequent rows as data, to build an array of JSON objects.
8. It includes specific error handling for InvalidRequestException which is common when the assumed role lacks s3:PutObject permissions on the output S3 bucket.

index.js

import {
	STSClient,
	AssumeRoleCommand
} from "@aws-sdk/client-sts";

import {
	AthenaClient,
	StartQueryExecutionCommand,
	GetQueryExecutionCommand,
	GetQueryResultsCommand
} from "@aws-sdk/client-athena";


const AWS_ACCESS_KEY_ID = process.env.AWS_ACCESS_KEY_ID;
const AWS_SECRET_ACCESS_KEY = process.env.AWS_SECRET_ACCESS_KEY;
const AWS_REGION = process.env.AWS_REGION;
const AWS_ASSUME_ROLE_EXTERNAL_ID = process.env.AWS_ASSUME_ROLE_EXTERNAL_ID;
const CLIENT_ORG_ID = process.env.CLIENT_ORG_ID;
const CLIENT_ROLE_NAME = process.env.CLIENT_ROLE_NAME;
const CLIENT_BUCKET = process.env.CLIENT_BUCKET;
const CLIENT_BUCKET_PATH = process.env.CLIENT_BUCKET_PATH;
const CLIENT_ATHENA_DB = process.env.CLIENT_ATHENA_DB;
const CLIENT_ATHENA_TABLE = process.env.CLIENT_ATHENA_TABLE;

// Config
const ROLE_ARN = `arn:aws:iam::${CLIENT_ORG_ID}:role/${CLIENT_ROLE_NAME}`;
const ATHENA_OUTPUT_S3 = `s3://${CLIENT_BUCKET}/${CLIENT_BUCKET_PATH}`;
const DATABASE = CLIENT_ATHENA_DB;
const TABLE = CLIENT_ATHENA_TABLE;

// Step 1: Assume the cross-account role
async function getTemporaryCredentials() {
	const sts = new STSClient({
		region: AWS_REGION,
		credentials: {
			accessKeyId: AWS_ACCESS_KEY_ID,
			secretAccessKey: AWS_SECRET_ACCESS_KEY
		}
	});

	const assumed = await sts.send(
		new AssumeRoleCommand({
			RoleArn: ROLE_ARN,
			RoleSessionName: "AthenaQuerySession",
			ExternalId: AWS_ASSUME_ROLE_EXTERNAL_ID
		})
	);

	return {
		accessKeyId: assumed.Credentials.AccessKeyId,
		secretAccessKey: assumed.Credentials.SecretAccessKey,
		sessionToken: assumed.Credentials.SessionToken
	};
}

// Step 2: Run Athena query and fetch results
async function runAthenaQuery() {
	const tempCreds = await getTemporaryCredentials();
	console.log("🚀 ~ runAthenaQuery ~ tempCreds:", tempCreds)

	const athena = new AthenaClient({
		region: AWS_REGION,
		credentials: tempCreds
	});

	let queryExecutionId;

	try {
		const startQuery = await athena.send(
			new StartQueryExecutionCommand({
				QueryString: `SELECT * FROM "${DATABASE}"."${TABLE}" WHERE user_id = '1234567890' order by created_at desc LIMIT 10;`,
				QueryExecutionContext: { Database: DATABASE },
				ResultConfiguration: { OutputLocation: ATHENA_OUTPUT_S3 }
			})
		);

		queryExecutionId = startQuery.QueryExecutionId;

		// Polling for query status
		let state = "RUNNING";
		while (state === "RUNNING" || state === "QUEUED") {
			await new Promise((resolve) => setTimeout(resolve, 3000));
			const status = await athena.send(
				new GetQueryExecutionCommand({ QueryExecutionId: queryExecutionId })
			);
			state = status.QueryExecution.Status.State;
			if (state === "FAILED" || state === "CANCELLED") {
				throw new Error(
					`Query ${state}: ${status.QueryExecution.Status.StateChangeReason}`
				);
			}
		}

		// Fetch results
		const results = await athena.send(
			new GetQueryResultsCommand({ QueryExecutionId: queryExecutionId })
		);

		// console.log("🚀 ~ runAthenaQuery ~ results:", results)

		console.log("✅ Query Results:");
		const json = [];
		let headers = null;
		
		for (const row of results.ResultSet.Rows) {
			if (!headers) {
				// First row contains headers
				headers = row.Data.map(cell => cell.VarCharValue);
				continue;
			}
			
			const obj = {};
			row.Data.forEach((cell, index) => {
				obj[headers[index]] = cell.VarCharValue;
			});
			json.push(obj);
		}
		console.log(json);
	} catch (err) {
		if (
			err.name === "InvalidRequestException" &&
			err.message.includes("output location")
		) {
			console.error("❌ Error: Missing write permission for result output location:");
			console.error("   → " + ATHENA_OUTPUT_S3);
			console.error("   Please ensure the role has 's3:PutObject' permission on this path.");
		} else {
			console.error("❌ General error:", err.message);
		}
	}
}

runAthenaQuery().catch((err) => {
	console.error("❌ Unexpected error:", err.message);
});

package.json

{
  "name": "athena-read",
  "version": "1.0.0",
  "main": "index.js",
  "scripts": {
    "start": "bun index.js"
  },
  "keywords": [],
  "author": "",
  "license": "ISC",
  "description": "",
  "dependencies": {
    "@aws-sdk/client-athena": "^3.826.0",
    "@aws-sdk/client-sts": "^3.826.0"
  }
}