vandonr-amz commented on code in PR #31942:
URL: https://github.com/apache/airflow/pull/31942#discussion_r1234494084
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -229,51 +242,33 @@ def get_query_results_paginator(
return paginator.paginate(**result_params)
def poll_query_status(
- self,
- query_execution_id: str,
- max_polling_attempts: int | None = None,
+ self, query_execution_id: str, max_polling_attempts: int | None =
None, sleep_time: int | None = None
) -> str | None:
"""Poll the state of a submitted query until it reaches final state.
:param query_execution_id: ID of submitted athena query
- :param max_polling_attempts: Number of times to poll for query state
- before function exits
+ :param max_polling_attempts: Number of times to poll for query state
before function exits
+ :param sleep_time: Time (in seconds) to wait between two consecutive
query status checks.
:return: One of the final states
"""
- try_number = 1
- final_query_state = None # Query state when query reaches final state
or max_polling_attempts reached
- while True:
- query_state = self.check_query_status(query_execution_id)
- if query_state is None:
- self.log.info(
- "Query execution id: %s, trial %s: Invalid query state.
Retrying again",
- query_execution_id,
- try_number,
- )
- elif query_state in self.TERMINAL_STATES:
- self.log.info(
- "Query execution id: %s, trial %s: Query execution
completed. Final state is %s",
- query_execution_id,
- try_number,
- query_state,
- )
- final_query_state = query_state
- break
- else:
- self.log.info(
- "Query execution id: %s, trial %s: Query is still in
non-terminal state - %s",
Review Comment:
ok, I used @syedahsn 's code, but as I merged his branch into mine, it's
going to show up in the diff until his PR gets merged
--
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]
