通过AWS Athena查询结果分页时如何跳过标头

时间:2018-08-03 10:36:50

标签: python amazon-web-services boto3 amazon-athena aws-glue

我有一个Angular 6应用,该应用从AWS Lambda请求数据。数据本身存储在Glue数据库中,并通过AWS Athena查询。 AWS Glue数据库设置了skip.header.line.count=1选项,当我在控制台中运行Athena查询时,我得到一个没有标题的响应。 当我尝试使用boto3检索数据时,会出现问题。我有一个函数,可以运行查询,然后对结果进行分页:

def run_query_paged(self, query, page_token=None, page_size=10):
    """
    Run query.
    """
    request = self.athena_client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': self.database
            },
        ResultConfiguration={
            'OutputLocation': self.s3_output,
            }
        )
    execution_id = request['QueryExecutionId']

    if execution_id:
        while True:
            stats = self.athena_client.get_query_execution(QueryExecutionId=execution_id)
            status = stats['QueryExecution']['Status']['State']
            if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
                break
            time.sleep(0.2)  # 200ms

        if status == 'SUCCEEDED':
            paginator = self.athena_client.get_paginator('get_query_results')
            pagination_config = {
                'MaxItems': page_size,
                'PageSize': page_size,
            }
            if page_token:
                pagination_config['StartingToken'] = page_token

            response_iterator = paginator.paginate(
                QueryExecutionId=execution_id,
                PaginationConfig=pagination_config
            )

            for page in response_iterator:
                next_token = page.get('NextToken', '')
                results = page
                break

            return {
                'rows': process_results(results),
                'nextToken': next_token
            }
        if status == 'FAILED':
            raise Exception(stats['QueryExecution']['Status']['StateChangeReason'])

    return None

process_results函数将响应转换为考虑到列类型的列表:

def process_results(response):
    """
    Processes the result of get_query_results function
    """
    rows = response['ResultSet']['Rows']
    meta = response['ResultSet']['ResultSetMetadata']['ColumnInfo']
    result = []
    for row in rows:
        parsed_row = {}
        for idx, val in enumerate(row['Data']):
            field = val
            column_info = meta[idx]
            if 'VarCharValue' in val:
                value = val['VarCharValue']
            else:
                value = ''
            parsed_row[column_info['Name']] = process_row_value(value, column_info)
        result.append(parsed_row)
    return result

问题在于分页响应的第一页的标题带有如下列名:

{
    "foo": "foo",
    "bar": "bar"
},
{
    "foo": 1,
    "bar": 2
},
...

,而其他所有页面都没有。当我从客户端应用程序请求第一页时,我得到一个标题加上9行(页面大小为10),当我使用NextToken请求下一页时,我又得到了10行而没有标题。在首页上显示9个项目,随后显示10个项目,这很尴尬。

如何跳过标题跳过结果?

2 个答案:

答案 0 :(得分:0)

我没有找到任何跳过标头的选项,并通过在第一个请求中请求page_size + 1的结果来破解它,然后对其余请求使用page_size进行入侵。

def _build_response(self, execution_id: str, starting_token: Optional[str], page_size: int) -> AthenaPagedResult:
    """
    Returns the query result for the provided page as well as a token to the next page if there are more
    results to retrieve for the query.
    """
    paginator = self.athena_client.get_paginator('get_query_results')

    # The first page of response contains header. Increase the page size for a first page and then
    # remove header so that all the pages would have the same size.
    if starting_token:
        skip_header = False
    else:
        page_size += 1
        skip_header = True
    max_items = page_size * 2

    pagination_config = {
        'MaxItems': min(max_items, MAXIMUM_ALLOWED_ITEMS_NUMBER),
        'PageSize': min(page_size, MAXIMUM_ALLOWED_ITEMS_NUMBER)
    }
    if starting_token:
        pagination_config['StartingToken'] = starting_token

    response_iterator = paginator.paginate(QueryExecutionId=execution_id, PaginationConfig=pagination_config)


    iterator_index = 0
    results = EMPTY_ATHENA_RESPONSE
    next_token = None

    # Retrieve only a single page and return the next token for the caller to iterate the response.
    for page in response_iterator:
        if iterator_index > 0:
            if len(page['ResultSet']['Rows']) == 0:
                next_token = None
            break
        next_token = page.get('NextToken')
        results = page
        iterator_index += 1

    # ... process and return results

答案 1 :(得分:0)

图我要加我的。我将其分为三个部分-启动查询,对query-result-pages进行分页,以及将结果标准化为列表/词典列表:

import boto3
import logging
from time import sleep

def query_athena_table(sql_query, database, **kwargs):
    client = boto3.client('athena')
    query_started = client.start_query_execution(
        QueryString=sql_query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={"OutputLocation": f"s3://your-specific-athena-query-results-bucket"}
    )

    timeout_value = kwargs.get("timeout", 15) * 1000 # bc its in milliseconds
    finished = False
    logging.info("Started Athena Query")

    while not finished:
        query_in_flight = client.get_query_execution(QueryExecutionId=query_started["QueryExecutionId"])
        query_status = query_in_flight["QueryExecution"]["Status"]["State"]

        if query_status == 'SUCCEEDED':
            finished = True
        elif query_status in ['FAILED', 'CANCELLED']:
            logging.error(query_in_flight['QueryExecution']['Status']['StateChangeReason'])
            return None
        elif timeout_value < ez_get(query_in_flight, "QueryExecution", "Statistics", "TotalExecutionTimeInMillis"):
            logging.warning(f"Query timed out with no response (timeout val: {timeout_value})")
            return None
        else:
            sleep(kwargs.get("wait_interval", 0.1))

    return paginate_athena_response(client, query_started["QueryExecutionId"], **kwargs)

# about 4s per 10k rows, with a floor of ~0.33s if only one page
def paginate_athena_response(client, execution_id: str, **kwargs):

    paginator = client.get_paginator('get_query_results')

    response_iterator = paginator.paginate(
        QueryExecutionId=execution_id, 
        PaginationConfig={
            'MaxItems': kwargs.get("max_results", 100000),
            'PageSize': 1000,
            'StartingToken': kwargs.get("pagination_starting_token", None),
    })

    results = []

    # Iterate through pages. The NextToken logic is handled for you.
    for n, page in enumerate(response_iterator):
        logging.info(f"Now on page {n}, rows on this page: {len(page['ResultSet']['Rows'])}")

        if n > 0 and len(page['ResultSet']['Rows']) == 0: # probably redundant
            break

        results += standardize_athena_query_result(page, **kwargs)
        kwargs["headers"] = list(results[0].keys()) # prevent parser from .pop(0) after 1st page

    return results

def standardize_athena_query_result(results, **kwargs):
    results = [x["Data"] for x in results['ResultSet']['Rows']]
    for n, row in enumerate(results):
        results[n] = [x['VarCharValue'] for x in row]

    if kwargs.get("output_lod"):
        headers = kwargs.get("headers") or results.pop(0)

        output_lod = []
        for n, result_row in enumerate(results):
            output_lod.append({headers[i]:result_row[i] for i in range(0, len(result_row))})
        return output_lod

    return results