syedahsn commented on code in PR #36586:
URL: https://github.com/apache/airflow/pull/36586#discussion_r1456419985
##########
airflow/providers/amazon/aws/hooks/redshift_data.py:
##########
@@ -201,3 +216,54 @@ def get_table_primary_key(
break
return pk_columns or None
+
+ async def check_query_is_finished_async(
+ self, statement_id: str, poll_interval: int = 10
+ ) -> dict[str, str]:
+ """Async function to check statement is finished.
+
+ It takes statement_id, makes async connection to redshift data to get
the query status
+ by statement_id and returns the query status.
+
+ :param statement_id: the UUID of the statement
+ :param poll_interval: how often in seconds to check the query status
+ """
+ try:
+ while await self.is_still_running(statement_id):
+ await asyncio.sleep(poll_interval)
+
+ resp = await self.async_conn.describe_statement(Id=statement_id)
+ status = resp["Status"]
Review Comment:
Both the `check_query_is_finished_async` and `check_query_is_finished`
should behave similarly, and if that is the case, you could probably take all
the parsing logic and combine it so that it can be used by both methods. i.e.
```
async def check_query_is_finished_async:
resp = await self.async_conn.describe_statement(Id=statement_id)
parse_resp(resp)
def check_query_is_finished:
resp = self.conn.describe_statement(Id=statement_id)
parse_resp(resp)
````
Doing it like this would reduce a lot of code, and if there were any API
changes that changed how the response looked, we would only need to make
changes in one place. That's kind of the logic I've used when I was working on
deferrable operators - reducing duplicate code as much as possible. What are
your thoughts on that? Th
--
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]
