How to check Athena query ran successfully in Python?

Question

I want to use the boto3 athena client function start_query_execution to run a query from Python. Function docs here: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena.html#Athena.Client.start_query_execution

Usage would be something like:

query = "SELECT * FROM TABLE"
athena_client = boto3.client("athena")
start_response = athena_client.start_query_execution(
    QueryString=query,
    QueryExecutionContext={"Database": ATHENA.database_name},
    ResultConfiguration={
        "OutputLocation": s3_output,
    },
)

I'm looking for a function / wrapper that ensures this query runs successfully and only returns once it has run to completion. Couldn't find an aws wrapper from my search.

Answer 1

I implemented a generic function that executes a particular query and also ensures it runs successfully by polling the query ID in intervals:

import time
import logging
import boto3

def run_query(query: str, s3_output: str) -> None:
    """Generic function to run athena query and ensures it is successfully completed

    Parameters
    ----------
    query : str
        formatted string containing athena sql query
    s3_output : str
        query output path
    """
    athena_client = boto3.client("athena")
    start_response = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={"Database": ATHENA.database_name},
        ResultConfiguration={
            "OutputLocation": s3_output,
        },
    )
    query_id = start_response["QueryExecutionId"]

    while True:
        finish_state = athena_client.get_query_execution(QueryExecutionId=query_id)[
            "QueryExecution"
        ]["Status"]["State"]
        if finish_state == "RUNNING" or finish_state == "QUEUED":
            time.sleep(10)
        else:
            break

    assert finish_state == "SUCCEEDED", f"query state is {finish_state}"
    logging.info(f"Query {query_id} complete")