Querying Data with Athena

Welcome back! In the last lesson, you learned how to use an AWS Glue Crawler to catalog your processed Parquet data. Now that your data is registered in the Glue Data Catalog, you are ready to unlock its value by running SQL queries for analytics and reporting. This is where Amazon Athena comes in.

Athena is a serverless query service that lets you analyze your data directly in S3 using standard SQL. Because your data is already cataloged, Athena can quickly discover the schema and let you run queries without any extra setup. In this lesson, you will learn how to use Python and the boto3 library to run Athena queries programmatically. This is a powerful skill for automating analytics and integrating data lake queries into your applications or workflows.

Before you can run queries with Athena, let’s quickly review what you need. This is mostly a reminder, as you have already completed these steps in previous lessons. First, your processed data should be stored in S3 and cataloged in the Glue Data Catalog. You should also have an Athena database that points to your cataloged tables. In this course, we use the database name data_lake_catalog_{SUFFIX}.

Athena needs a place to store the results of your queries. This is usually an S3 bucket and folder, such as s3://library-data-lake-{SUFFIX}/athena_query_results/. In practice, that output prefix should already exist or be created during setup, and your IAM user or role must be able to write query results there in addition to using Athena and the Glue Catalog. If you are working in the CodeSignal environment, the boto3 library is already installed for you. If you are working on your own machine, you will need to install it with pip install boto3.

Let’s walk through how to run a query in Athena using Python and boto3. Below is a script that connects to Athena, runs a SQL query, waits for the query to finish, and then prints the results. This script is designed to be simple and easy to follow.

import boto3
import time

with open(".codesignal/suffix.txt") as f:
    SUFFIX = f.read().strip()
DATABASE_NAME = f'data_lake_catalog_{SUFFIX}'
ATHENA_OUTPUT = f's3://library-data-lake-{SUFFIX}/athena_query_results/'

def execute_athena_query(query):
    athena = boto3.client('athena')
    response = athena.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': DATABASE_NAME},
        ResultConfiguration={'OutputLocation': ATHENA_OUTPUT}
    )
    qid = response['QueryExecutionId']
    while True:
        status = athena.get_query_execution(QueryExecutionId=qid)['QueryExecution']['Status']['State']
        if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
            break
        time.sleep(1)
    results = athena.get_query_results(QueryExecutionId=qid)
    rows = results['ResultSet']['Rows']
    headers = [c.get('VarCharValue', '') for c in rows[0]['Data']]
    print(headers)
    for row in rows[1:11]:
        print([c.get('VarCharValue', '') for c in row['Data']])

if __name__ == "__main__":
    execute_athena_query("SELECT COUNT(*) FROM library;")

Here’s how this script works. First, it sets up the Athena client using boto3. The execute_athena_query function takes a SQL query as input. It starts the query execution in Athena and gets a unique query ID. The script then checks the status of the query every second until it finishes. Once the query is done, it fetches the results, prints the header row first, and then prints up to ten result rows so the output stays readable.

For example, if you run the query SELECT COUNT(*) FROM library;, you might see output like:

['_col0']
['30']

The first row is the column header, and the second row is the result of the count. This approach keeps query output easier to read while still working for any SQL query you want to run on your cataloged data.