AWS Athena : Execute query in python boto3

To execute an Amazon Athena query using the boto3 library in Python, you can follow these steps:

1. Install Boto3: If you haven’t already, install the boto3 library using pip:

pip install boto3

2. Import Libraries: Import the necessary libraries and create a boto3 client for Athena:

import boto3
# Create a Boto3 client for Athena
athena_client = boto3.client('athena', region_name='your-region-name')

3. Execute the Query: Use the start_query_execution method to submit a query to Athena:

def execute_athena_query(query_string, database, output_location):
    response = athena_client.start_query_execution(
        QueryString=query_string,
        QueryExecutionContext={
            'Database': database
        },
        ResultConfiguration={
            'OutputLocation': output_location
        }
    )
    return response['QueryExecutionId']

Replace 'your-region-name', query_string, database, and output_location with your actual values.

4. Check Query Status: You can use the get_query_execution method to check the status of your query:

def get_query_status(query_execution_id):
    response = athena_client.get_query_execution(
        QueryExecutionId=query_execution_id
    )
    return response['QueryExecution']['Status']['State']

5. Retrieve Query Results: If the query is successful, you can retrieve the results using the get_query_results method:

def get_query_results(query_execution_id):
    response = athena_client.get_query_results(
        QueryExecutionId=query_execution_id
    )
    # Process and print/query the results
    for row in response['ResultSet']['Rows']:
        print([field['VarCharValue'] for field in row['Data']])

Here’s an example of how to use these functions:

query_string = "SELECT * FROM your_table LIMIT 10;"
database = "your_database"
output_location = "s3://your-bucket-name/athena-query-results/"
query_execution_id = execute_athena_query(query_string, database, output_location)
while get_query_status(query_execution_id) == 'RUNNING':
    print("Query is still running...")
    time.sleep(5)  # Wait for a few seconds before checking again
if get_query_status(query_execution_id) == 'SUCCEEDED':
    print("Query succeeded!")
    get_query_results(query_execution_id)
else:
    print("Query failed or was canceled.")

Remember to replace placeholders like 'your-region-name', 'your_table', 'your_database', 'your-bucket-name', and adapt the code to your specific query and environment.

Make sure you have the necessary AWS credentials configured properly (using environment variables, configuration files, or IAM roles) to authenticate your Python script when interacting with AWS services.

Happy learning ..!!