From e4f7acc9694c38a80615febf49aafe450ee09197 Mon Sep 17 00:00:00 2001
From: Deepu Mohan Puthrote <git@deepumohan.com>
Date: Sat, 9 Dec 2023 17:45:28 +0000
Subject: [PATCH] fix: Log athena query id when it is available. Closes #530
- Log the query id as soon as it is available
- Adds query_id to warning and debug logs in AthenaCursor
---
dbt/adapters/athena/connections.py | 9 +++++++--
1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/dbt/adapters/athena/connections.py b/dbt/adapters/athena/connections.py
index f12a5c47..41147167 100644
--- a/dbt/adapters/athena/connections.py
+++ b/dbt/adapters/athena/connections.py
@@ -121,7 +121,7 @@ def _poll(self, query_id: str) -> AthenaQueryExecution:
query_execution = self.__poll(query_id)
except KeyboardInterrupt as e:
if self._kill_on_interrupt:
- logger.warning("Query canceled by user.")
+ logger.warning(f"Query canceled by user: {query_id}.")
self._cancel(query_id)
query_execution = self.__poll(query_id)
else:
@@ -139,7 +139,9 @@ def __poll(self, query_id: str) -> AthenaQueryExecution:
return query_execution
if self.connection.cursor_kwargs.get("debug_query_state", False):
- logger.debug(f"Query state is: {query_execution.state}. Sleeping for {self._poll_interval}...")
+ logger.debug(
+ f"Query state for {query_id} is: {query_execution.state}. Sleeping for {self._poll_interval}..."
+ )
time.sleep(self._poll_interval)
def execute( # type: ignore
@@ -163,6 +165,9 @@ def inner() -> AthenaCursor:
cache_size=cache_size,
cache_expiration_time=cache_expiration_time,
)
+
+ logger.debug(f"Athena query ID {query_id}")
+
query_execution = self._executor.submit(self._collect_result_set, query_id).result()
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
self.result_set = self._result_set_class(