import time


def query(query, database, workgroup, athena_client, max_checks=30):
    """Executes an Athena query, waits for success or failure, and returns the first page
    of the query results.

    Waits up to max_checks * 10 seconds for the query to complete before raising.
    """
    resp = athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={"Database": database},
        WorkGroup=workgroup,
    )
    qid = resp["QueryExecutionId"]
    for i in range(max_checks):
        resp = athena_client.get_query_execution(QueryExecutionId=qid)
        state = resp["QueryExecution"]["Status"]["State"]
        if state == "SUCCEEDED":
            return qid
        elif state == "FAILED":
            raise RuntimeError(f"Failed query execution: {query}")
        # Continue to wait
        time.sleep(10)
    else:
        raise TimeoutError("Reached max_checks")