ashb commented on a change in pull request #6075:
URL: https://github.com/apache/airflow/pull/6075#discussion_r450720324
##########
File path: airflow/providers/amazon/aws/hooks/athena.py
##########
@@ -125,7 +133,49 @@ def get_query_results(self, query_execution_id):
elif query_state in self.INTERMEDIATE_STATES or query_state in
self.FAILURE_STATES:
self.log.error('Query is in "%s" state. Cannot fetch results',
query_state)
return None
- return
self.get_conn().get_query_results(QueryExecutionId=query_execution_id)
+ result_params = {
+ 'QueryExecutionId': query_execution_id,
+ 'MaxResults': max_results
+ }
+ if next_token_id:
+ result_params['NextToken'] = next_token_id
+ return self.get_conn().get_query_results(**result_params)
+
+ def get_query_results_paginator(self, query_execution_id, max_items=None,
+ page_size=None, starting_token=None):
+ """
+ Fetch submitted athena query results. returns none if query is in
intermediate state or
+ failed/cancelled state else a paginator to iterate through pages of
results. If you
+ wish to get all results at once, call build_full_result() on the
returned PageIterator
+
+ :param query_execution_id: Id of submitted athena query
+ :type query_execution_id: str
+ :param max_items: The total number of items to return.
+ :type max_items: int
+ :param page_size: The size of each page.
+ :type page_size: int
+ :param starting_token: A token to specify where to start paginating.
+ :type starting_token: str
+ :return: PageIterator
+ """
+ query_state = self.check_query_status(query_execution_id)
+ if query_state is None:
+ self.log.error('Invalid Query state')
+ return None
+ elif query_state in self.INTERMEDIATE_STATES or query_state in
self.FAILURE_STATES:
Review comment:
```suggestion
if query_state in self.INTERMEDIATE_STATES or query_state in
self.FAILURE_STATES:
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
