[python] boto3 athena 실행

boto3 아테나 클라이언트에서 start_query_execution를 실행하면, 쿼리를 시작하고 끝나지 않아도 바로 response가 되돌아온다.
QueryExecutionId로 쿼리 진행상태를 확인하고 다음 과정을 이어간다.

아래 함수는 CREATE TABLE, DROP TABLE을 할때 쓸 수 있는 함수이다. (SELECT 등은 고려되지 않음)
- Rate exceeded : 너무 많은 요청을 했을때 메시지이므로 쉬어간다.
- Could not find results : DROP TABLE을 실행
- Query has not yet finished : 쿼리가 끝나지 않았으므로 기다린다.

import boto3

client = boto3.client('athena')

def athena_query(query: str):
    response = start_query_execution(query)
    while True:
        try:
            sleep(5)
            return client.get_query_results(QueryExecutionId=response['QueryExecutionId'])
        except Exception as e:
            err_response = getattr(e, 'response', None)
            if err_response is None:
                raise e
            if err_response['Error']['Message'].__contains__("Could not find results"):
                return
            elif err_response['Error']['Message'].__contains__("Query has not yet finished") \
                    or err_response['Error']['Message'].__contains__("Rate exceeded"):
                sleep(5)
                continue
            raise e
                

def start_query_execution(query: str):
    response = client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': 'test_database',
        },
        ResultConfiguration={
            'OutputLocation': 's3://test-bucket/test-path/',
        },
        WorkGroup='test-workgroup'
    )
    return response