From 569f4a3caf6a2fbb52cf6667c0f28fddfd7fe79b Mon Sep 17 00:00:00 2001
From: Benjamin Smith <bh2smith@users.noreply.github.com>
Date: Fri, 7 Oct 2022 13:24:30 +0200
Subject: [PATCH] update query refresh to include terminal states (#21)

---
 dune_client/client.py | 28 +++++++++++++++++-----------
 dune_client/models.py | 19 +++++++++++++++++++
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/dune_client/client.py b/dune_client/client.py
index 2e2bf12..d20080b 100644
--- a/dune_client/client.py
+++ b/dune_client/client.py
@@ -119,16 +119,22 @@ def refresh(self, query: Query, ping_frequency: int = 5) -> list[DuneRecord]:
         Sleeps `ping_frequency` seconds between each status request.
         """
         job_id = self.execute(query).execution_id
-        state = self.get_status(job_id).state
-        while state != ExecutionState.COMPLETED:
-            log.info(
-                f"waiting for query execution {job_id} to complete: current state {state}"
-            )
+        status = self.get_status(job_id)
+        while status.state not in ExecutionState.terminal_states():
+            log.info(f"waiting for query execution {job_id} to complete: {status}")
             time.sleep(ping_frequency)
-            state = self.get_status(job_id).state
+            status = self.get_status(job_id)
 
-        full_response = self.get_result(job_id)
-        assert (
-            full_response.result is not None
-        ), f"Expected Results on completed execution status {full_response}"
-        return full_response.result.rows
+        if status.state == ExecutionState.COMPLETED:
+            full_response = self.get_result(job_id)
+            assert (
+                full_response.result is not None
+            ), f"Expected Results on completed execution status {full_response}"
+            return full_response.result.rows
+
+        if status.state == ExecutionState.CANCELLED:
+            log.info("Execution Cancelled, returning empty record set")
+            return []
+
+        log.error(status)
+        raise Exception(f"{status}. Perhaps your query took too long to run!")
diff --git a/dune_client/models.py b/dune_client/models.py
index 2153659..06d65af 100644
--- a/dune_client/models.py
+++ b/dune_client/models.py
@@ -42,6 +42,14 @@ class ExecutionState(Enum):
     EXECUTING = "QUERY_STATE_EXECUTING"
     PENDING = "QUERY_STATE_PENDING"
     CANCELLED = "QUERY_STATE_CANCELLED"
+    FAILED = "QUERY_STATE_FAILED"
+
+    @classmethod
+    def terminal_states(cls) -> set[ExecutionState]:
+        """
+        Returns the terminal states (i.e. when a query execution is no longer executing
+        """
+        return {cls.COMPLETED, cls.CANCELLED, cls.FAILED}
 
 
 @dataclass
@@ -116,6 +124,17 @@ def from_dict(cls, data: dict[str, Any]) -> ExecutionStatusResponse:
             times=TimeData.from_dict(data),  # Sending the entire data dict
         )
 
+    def __str__(self) -> str:
+        if self.state == ExecutionState.PENDING:
+            return f"{self.state} (queue position: {self.queue_position})"
+        if self.state == ExecutionState.FAILED:
+            return (
+                f"{self.state}: execution_id={self.execution_id}, "
+                f"query_id={self.query_id}, times={self.times}"
+            )
+
+        return f"{self.state}"
+
 
 @dataclass
 class ResultMetadata: