[GitHub] [airflow] mik-laj commented on a change in pull request #20124: Added function in AWSAthenaHook to get s3 output query results file URI

[GitBox <gi...@apache.org>]

mik-laj commented on a change in pull request #20124:
URL: https://github.com/apache/airflow/pull/20124#discussion_r764543412



##########
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")

Review comment:
       ```suggestion
                       self.log.error("Error retrieving OutputLocation")
                       raise
   ```




-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org