Boto3 Athena Query Example

execute_athena_query.py

def execute_athena_query(query, database="default", timeout=30, sleep_time=10):
    context = {"Database": database}
    config = {"OutputLocation": os.environ["ATHENA_BUCKET"]}

    # Execute query
    request = athena.start_query_execution(QueryString=query, QueryExecutionContext=context, ResultConfiguration=config)

    # Wait for query result
    num_tries = int(timeout / sleep_time)
    status = athena.get_query_execution(QueryExecutionId=request["QueryExecutionId"])
    while status["QueryExecution"]["Status"]["State"] in ["RUNNING", "QUEUED"]:
        time.sleep(sleep_time)
        status = athena.get_query_execution(QueryExecutionId=request["QueryExecutionId"])
        num_tries -= 1
        if num_tries <= 0:
            raise TimeoutError("Query failed due to timeout")

    # Get result set
    if status["QueryExecution"]["Status"]["State"] != "SUCCEEDED":
        raise ValueError(f"Query failed: {status}")

    result = athena.get_query_results(QueryExecutionId=request["QueryExecutionId"])
    return result["ResultSet"]["Rows"]