mik-laj commented on a change in pull request #20124:
URL: https://github.com/apache/airflow/pull/20124#discussion_r764543794
##########
File path: airflow/providers/amazon/aws/hooks/athena.py
##########
@@ -227,6 +227,29 @@ def poll_query_status(self, query_execution_id: str,
max_tries: Optional[int] =
sleep(self.sleep_time)
return final_query_state
+ def get_output_location(self, query_execution_id: str) -> str:
+ """
+ Function to get the output location of the query results
+ in s3 uri format.
+
+ :param query_execution_id: Id of submitted athena query
+ :type query_execution_id: str
+ :return: str
+ """
+ output_location = None
+ if query_execution_id:
+ response =
self.get_conn().get_query_execution(QueryExecutionId=query_execution_id)
+
+ if response:
+ try:
+ output_location =
response['QueryExecution']['ResultConfiguration']['OutputLocation']
+ except KeyError:
+ self.log.error("Error retrieving OutputLocation")
+ else:
+ self.log.error("Invalid Query execution id")
Review comment:
```suggestion
raise ValueError("Invalid Query execution id")
```
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
