Taragolis commented on code in PR #35090:
URL: https://github.com/apache/airflow/pull/35090#discussion_r1374380605
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -121,7 +122,26 @@ def run_query(
self.log.info("Query execution id: %s", query_execution_id)
return query_execution_id
- def check_query_status(self, query_execution_id: str) -> str | None:
+ def get_query_execution(self, query_execution_id: str,
use_cached_response: bool = False) -> dict:
+ """Get information about a single execution of a query.
+
+ .. seealso::
+ - :external+boto3:py:meth:`Athena.Client.get_query_execution`
+
+ :param query_execution_id: Id of submitted athena query
+ :param use_cached_response: If True, use execution information cache
+ """
+ # second check if to satisfy mypy
+ if use_cached_response or not hasattr(self._get_query_execution,
"__wrapped__"):
+ return
self._get_query_execution(query_execution_id=query_execution_id)
+ return self._get_query_execution.__wrapped__(self,
query_execution_id=query_execution_id)
Review Comment:
Looks a bit hacky for me, maybe we could specify in constructor private
attribute and store information there. WDYT?
```python
def __init__(...):
...
self.__query_results: dict[str, Any] = {}
def get_query_execution(self, query_execution_id: str, use_cache: bool =
False):
if use_cache and query_execution_id in self.__query_results:
return self.__query_results[query_execution_id]
response =
self.conn.get_query_execution(QueryExecutionId=query_execution_id)
if use_cache:
self.__query_results[query_execution_id] = response
return response
```
--
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]
