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: