add failed state

Signed-off-by: Prithvi Kannan <[email protected]>
databricks · Oct 30, 2024 · 035effe · 035effe
1 parent 63dd12d
commit 035effe
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py
@@ -116,6 +116,9 @@ def poll_result():
                     return poll_query_results()
                 elif resp["status"] == "COMPLETED":
                     return next(r for r in resp["attachments"] if "text" in r)["text"]["content"]
+                elif resp["status"] == "FAILED":
+                    logging.debug("Genie failed to execute the query")
+                    return None
                 else:
                     logging.debug(f"Waiting...: {resp['status']}")
                     time.sleep(5)