From ded27473cebbc01c29ecff69b4194485191b886e Mon Sep 17 00:00:00 2001
From: Serhii Dimchenko <39801237+svdimchenko@users.noreply.github.com>
Date: Mon, 4 Sep 2023 14:32:42 +0200
Subject: [PATCH] feat: Implement configurable debug_query_state parameter for
 AthenaCursor (#395)

---
 README.md                          | 1 +
 dbt/adapters/athena/connections.py | 7 +++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 133290a6..15b5a6c9 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,7 @@ A dbt profile can be configured to run against AWS Athena using the following co
 | schema                | Specify the schema (Athena database) to build models into (lowercase **only**) | Required  | `dbt`                                      |
 | database              | Specify the database (Data catalog) to build models into (lowercase **only**)  | Required  | `awsdatacatalog`                           |
 | poll_interval         | Interval in seconds to use for polling the status of query results in Athena   | Optional  | `5`                                        |
+| debug_query_state     | Flag if debug message with Athena query state is needed                        | Optional  | `false`                                    |
 | aws_access_key_id     | Access key ID of the user performing requests.                                 | Optional  | `AKIAIOSFODNN7EXAMPLE`                     |
 | aws_secret_access_key | Secret access key of the user performing requests                              | Optional  | `wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY` |
 | aws_profile_name      | Profile to use from your AWS shared credentials file.                          | Optional  | `my-profile`                               |
diff --git a/dbt/adapters/athena/connections.py b/dbt/adapters/athena/connections.py
index 8645f8bc..89363d92 100644
--- a/dbt/adapters/athena/connections.py
+++ b/dbt/adapters/athena/connections.py
@@ -54,6 +54,7 @@ class AthenaCredentials(Credentials):
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
     poll_interval: float = 1.0
+    debug_query_state: bool = False
     _ALIASES = {"catalog": "database"}
     num_retries: Optional[int] = 5
     s3_data_dir: Optional[str] = None
@@ -81,7 +82,7 @@ def _connection_keys(self) -> Tuple[str, ...]:
             "endpoint_url",
             "s3_data_dir",
             "s3_data_naming",
-            "lf_tags",
+            "debug_query_state",
         )
 
 
@@ -122,7 +123,8 @@ def __poll(self, query_id: str) -> AthenaQueryExecution:
             ]:
                 return query_execution
             else:
-                logger.debug(f"Query state is: {query_execution.state}. Sleeping for {self._poll_interval}...")
+                if self.connection.cursor_kwargs.get("debug_query_state", False):
+                    logger.debug(f"Query state is: {query_execution.state}. Sleeping for {self._poll_interval}...")
                 time.sleep(self._poll_interval)
 
     def execute(  # type: ignore
@@ -215,6 +217,7 @@ def open(cls, connection: Connection) -> Connection:
                 schema_name=creds.schema,
                 work_group=creds.work_group,
                 cursor_class=AthenaCursor,
+                cursor_kwargs={"debug_query_state": creds.debug_query_state},
                 formatter=AthenaParameterFormatter(),
                 poll_interval=creds.poll_interval,
                 session=get_boto3_session(connection),