How to check Athena query ran successfully in Python?

Solution 1:

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

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")