From 55d5f92f5436ad5b4fc18d7c050a01f463fce017 Mon Sep 17 00:00:00 2001 From: Benjamin Smith Date: Thu, 27 Oct 2022 15:59:01 +0200 Subject: [PATCH] hashing query and returning full response from refresh (#26) --- dune_client/client.py | 22 +++++++--------------- dune_client/interface.py | 5 ++--- dune_client/models.py | 18 ++++++++++++++++-- dune_client/query.py | 9 ++++++++- tests/e2e/test_client.py | 14 ++++++++++---- tests/unit/test_query.py | 18 +++++++++++++++++- 6 files changed, 60 insertions(+), 26 deletions(-) diff --git a/dune_client/client.py b/dune_client/client.py index d20080b..53586b8 100644 --- a/dune_client/client.py +++ b/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 diff --git a/dune_client/interface.py b/dune_client/interface.py index 3dfdf4f..bf81f10 100644 --- a/dune_client/interface.py +++ b/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. diff --git a/dune_client/models.py b/dune_client/models.py index 06d65af..402ada8 100644 --- a/dune_client/models.py +++ b/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 [] diff --git a/dune_client/query.py b/dune_client/query.py index c000f11..1fdf7e9 100644 --- a/dune_client/query.py +++ b/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__() diff --git a/tests/e2e/test_client.py b/tests/e2e/test_client.py index a78f8fd..9499443 100644 --- a/tests/e2e/test_client.py +++ b/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: diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index 661e9dd..904bfeb 100644 --- a/tests/unit/test_query.py +++ b/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()