hashing query and returning full response from refresh (#26)

Author: bh2smith

dune_client/client.py

@@ -21,7 +21,7 @@
     ResultsResponse,
     ExecutionState,
 )
-from dune_client.types import DuneRecord
+
 from dune_client.query import Query
 
 log = logging.getLogger(__name__)
@@ -112,7 +112,7 @@ def cancel_execution(self, job_id: str) -> bool:
         except KeyError as err:
             raise DuneError(response_json, "CancellationResponse", err) from err
 
-    def refresh(self, query: Query, ping_frequency: int = 5) -> list[DuneRecord]:
+    def refresh(self, query: Query, ping_frequency: int = 5) -> ResultsResponse:
         """
         Executes a Dune `query`, waits until execution completes,
         fetches and returns the results.
@@ -125,16 +125,8 @@ def refresh(self, query: Query, ping_frequency: int = 5) -> list[DuneRecord]:
             time.sleep(ping_frequency)
             status = self.get_status(job_id)
 
-        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!")
+        full_response = self.get_result(job_id)
+        if status.state == ExecutionState.FAILED:
+            log.error(status)
+            raise Exception(f"{status}. Perhaps your query took too long to run!")
+        return full_response

dune_client/interface.py

@@ -2,10 +2,9 @@
 Abstract class for a basic Dune Interface with refresh method used by Query Runner.
 """
 from abc import ABC
-from typing import List
 
+from dune_client.models import ResultsResponse
 from dune_client.query import Query
-from dune_client.types import DuneRecord
 
 
 # pylint: disable=too-few-public-methods
@@ -14,7 +13,7 @@ class DuneInterface(ABC):
     User Facing Methods for a Dune Client
     """
 
-    def refresh(self, query: Query) -> List[DuneRecord]:
+    def refresh(self, query: Query) -> ResultsResponse:
         """
         Executes a Dune query, waits till query execution completes,
         fetches and returns the results.

dune_client/models.py

@@ -208,11 +208,25 @@ def from_dict(cls, data: dict[str, str | int | ResultData]) -> ResultsResponse:
         assert isinstance(data["execution_id"], str)
         assert isinstance(data["query_id"], int)
         assert isinstance(data["state"], str)
-        assert isinstance(data["result"], dict)
+        result = data.get("result", {})
+        assert isinstance(result, dict)
         return cls(
             execution_id=data["execution_id"],
             query_id=int(data["query_id"]),
             state=ExecutionState(data["state"]),
             times=TimeData.from_dict(data),
-            result=ExecutionResult.from_dict(data["result"]),
+            result=ExecutionResult.from_dict(result) if result else None,
         )
+
+    def get_rows(self) -> list[DuneRecord]:
+        """
+        Absorbs the Optional check and returns the result rows.
+        When execution is a non-complete terminal state, returns empty list.
+        """
+
+        if self.state == ExecutionState.COMPLETED:
+            assert self.result is not None, f"No Results on completed execution {self}"
+            return self.result.rows
+
+        log.info(f"execution {self.state} returning empty list")
+        return []

dune_client/query.py

@@ -12,8 +12,8 @@
 class Query:
     """Basic data structure constituting a Dune Analytics Query."""
 
-    name: str
     query_id: int
+    name: Optional[str] = "unnamed"
     params: Optional[List[QueryParameter]] = None
 
     def base_url(self) -> str:
@@ -33,3 +33,10 @@ def url(self) -> str:
                 [self.base_url(), urllib.parse.quote_plus(params, safe="=&?")]
             )
         return self.base_url()
+
+    def __hash__(self) -> int:
+        """
+        This contains the query ID and the values of relevant parameters.
+        Thus, it is unique for caching purposes
+        """
+        return self.url().__hash__()

tests/e2e/test_client.py

@@ -36,11 +36,13 @@ def test_get_status(self):
         dune = DuneClient(self.valid_api_key)
         job_id = dune.execute(query).execution_id
         status = dune.get_status(job_id)
-        self.assertEqual(status.state, ExecutionState.EXECUTING)
+        self.assertTrue(
+            status.state in [ExecutionState.EXECUTING, ExecutionState.PENDING]
+        )
 
     def test_refresh(self):
         dune = DuneClient(self.valid_api_key)
-        results = dune.refresh(self.query)
+        results = dune.refresh(self.query).get_rows()
         self.assertGreater(len(results), 0)
 
     def test_parameters_recognized(self):
@@ -58,7 +60,7 @@ def test_parameters_recognized(self):
         dune = DuneClient(self.valid_api_key)
         results = dune.refresh(query)
         self.assertEqual(
-            results,
+            results.get_rows(),
             [
                 {
                     "text_field": "different word",
@@ -88,10 +90,14 @@ def test_cancel_execution(self):
             query_id=1229120,
         )
         execution_response = dune.execute(query)
+        job_id = execution_response.execution_id
         # POST Cancellation
-        success = dune.cancel_execution(execution_response.execution_id)
+        success = dune.cancel_execution(job_id)
         self.assertTrue(success)
 
+        results = dune.get_result(job_id)
+        self.assertEqual(results.state, ExecutionState.CANCELLED)
+
     def test_invalid_api_key_error(self):
         dune = DuneClient(api_key="Invalid Key")
         with self.assertRaises(DuneError) as err:

tests/unit/test_query.py

@@ -24,11 +24,27 @@ def test_url(self):
             self.query.url(),
             "https://dune.com/queries/0?Enum=option1&Text=plain+text&Number=12&Date=2021-01-01+12%3A34%3A56",
         )
-        self.assertEqual(Query("", 0, []).url(), "https://dune.com/queries/0")
+        self.assertEqual(Query(0, "", []).url(), "https://dune.com/queries/0")
 
     def test_parameters(self):
         self.assertEqual(self.query.parameters(), self.query_params)
 
+    def test_hash(self):
+        # Same ID, different params
+        query1 = Query(query_id=0, params=[QueryParameter.text_type("Text", "word1")])
+        query2 = Query(query_id=0, params=[QueryParameter.text_type("Text", "word2")])
+        self.assertNotEqual(hash(query1), hash(query2))
+
+        # Different ID, same
+        query1 = Query(query_id=0)
+        query2 = Query(query_id=1)
+        self.assertNotEqual(hash(query1), hash(query2))
+
+        # Different ID different params
+        query1 = Query(query_id=0)
+        query2 = Query(query_id=1, params=[QueryParameter.number_type("num", 1)])
+        self.assertNotEqual(hash(query1), hash(query2))
+
 
 if __name__ == "__main__":
     unittest.main()