logging

Signed-off-by: Prithvi Kannan <[email protected]>
databricks · Oct 22, 2024 · 4cadcbc · 4cadcbc
1 parent 1b17d84
commit 4cadcbc
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py
@@ -1,3 +1,4 @@
+import logging
 import time
 from datetime import datetime
 from typing import Union
@@ -82,12 +83,12 @@ def poll_result():
                 )
                 if resp["status"] == "EXECUTING_QUERY":
                     sql = next(r for r in resp["attachments"] if "query" in r)["query"]["query"]
-                    # print(f"SQL: {sql}")
+                    logging.debug(f"SQL: {sql}")
                     return poll_query_results()
                 elif resp["status"] == "COMPLETED":
                     return next(r for r in resp["attachments"] if "text" in r)["text"]["content"]
                 else:
-                    # print(f"Waiting...: {resp['status']}")
+                    logging.debug(f"Waiting...: {resp['status']}")
                     time.sleep(5)
 
         def poll_query_results():
@@ -101,10 +102,10 @@ def poll_query_results():
                 if state == "SUCCEEDED":
                     return _parse_query_result(resp)
                 elif state == "RUNNING" or state == "PENDING":
-                    # print(f"Waiting for query result...")
+                    logging.debug("Waiting for query result...")
                     time.sleep(5)
                 else:
-                    # print(f"No query result: {resp['state']}")
+                    logging.debug(f"No query result: {resp['state']}")
                     return None
 
         return poll_result()