feat: add async support for csv and dataframe methods

dune_client/client.py

@@ -33,10 +33,7 @@ class DuneClient(DuneInterface, BaseDuneClient):
  combining the use of endpoints (e.g. refresh)
  """
 
- def _handle_response(
- self,
- response: Response,
- ) -> Any:
+ def _handle_response(self, response: Response) -> Any:
  try:
  # Some responses can be decoded and converted to DuneErrors
  response_json = response.json()
@@ -48,17 +45,24 @@ def _handle_response(
  raise ValueError("Unreachable since previous line raises") from err
 
  def _route_url(self, route: str) -> str:
- return f"{self.BASE_URL}{self.API_PATH}/{route}"
+ return f"{self.BASE_URL}{self.API_PATH}{route}"
 
- def _get(self, route: str, params: Optional[Any] = None) -> Any:
+ def _get(
+ self,
+ route: str,
+ params: Optional[Any] = None,
+ raw: bool = False,
+ ) -> Any:
  url = self._route_url(route)
  self.logger.debug(f"GET received input url={url}")
  response = requests.get(
- url,
- headers={"x-dune-api-key": self.token},
+ url=url,
+ headers=self.default_headers(),
  timeout=self.DEFAULT_TIMEOUT,
  params=params,
  )
+ if raw:
+ return response
  return self._handle_response(response)
 
  def _post(self, route: str, params: Any) -> Any:
@@ -67,7 +71,7 @@ def _post(self, route: str, params: Any) -> Any:
  response = requests.post(
  url=url,
  json=params,
- headers={"x-dune-api-key": self.token},
+ headers=self.default_headers(),
  timeout=self.DEFAULT_TIMEOUT,
  )
  return self._handle_response(response)
@@ -76,17 +80,15 @@ def execute(
  self, query: Query, performance: Optional[str] = None
  ) -> ExecutionResponse:
  """Post's to Dune API for execute `query`"""
+ params = query.request_format()
+ params["performance"] = performance or self.performance
+
  self.logger.info(
  f"executing {query.query_id} on {performance or self.performance} cluster"
  )
  response_json = self._post(
- route=f"query/{query.query_id}/execute",
- params={
- "query_parameters": {
- p.key: p.to_dict()["value"] for p in query.parameters()
- },
- "performance": performance or self.performance,
- },
+ route=f"/query/{query.query_id}/execute",
+ params=params,
  )
  try:
  return ExecutionResponse.from_dict(response_json)
@@ -95,17 +97,15 @@ def execute(
 
  def get_status(self, job_id: str) -> ExecutionStatusResponse:
  """GET status from Dune API for `job_id` (aka `execution_id`)"""
- response_json = self._get(
- route=f"execution/{job_id}/status",
- )
+ response_json = self._get(route=f"/execution/{job_id}/status")
  try:
  return ExecutionStatusResponse.from_dict(response_json)
  except KeyError as err:
  raise DuneError(response_json, "ExecutionStatusResponse", err) from err
 
  def get_result(self, job_id: str) -> ResultsResponse:
  """GET results from Dune API for `job_id` (aka `execution_id`)"""
- response_json = self._get(route=f"execution/{job_id}/results")
+ response_json = self._get(route=f"/execution/{job_id}/results")
  try:
  return ResultsResponse.from_dict(response_json)
  except KeyError as err:
@@ -119,13 +119,10 @@ def get_result_csv(self, job_id: str) -> ExecutionResultCSV:
  use this method for large results where you want lower CPU and memory overhead
  if you need metadata information use get_results() or get_status()
  """
- url = self._route_url(f"execution/{job_id}/results/csv")
+ route = f"/execution/{job_id}/results/csv"
+ url = self._route_url(f"/execution/{job_id}/results/csv")
  self.logger.debug(f"GET CSV received input url={url}")
- response = requests.get(
- url,
- headers={"x-dune-api-key": self.token},
- timeout=self.DEFAULT_TIMEOUT,
- )
+ response = self._get(route=route, raw=True)
  response.raise_for_status()
  return ExecutionResultCSV(data=BytesIO(response.content))
 
@@ -147,7 +144,7 @@ def get_latest_result(self, query: Union[Query, str, int]) -> ResultsResponse:
  query_id = int(query)
 
  response_json = self._get(
- route=f"query/{query_id}/results",
+ route=f"/query/{query_id}/results",
  params=params,
  )
  try:
@@ -157,7 +154,10 @@ def get_latest_result(self, query: Union[Query, str, int]) -> ResultsResponse:
 
  def cancel_execution(self, job_id: str) -> bool:
  """POST Execution Cancellation to Dune API for `job_id` (aka `execution_id`)"""
- response_json = self._post(route=f"execution/{job_id}/cancel", params=None)
+ response_json = self._post(
+ route=f"/execution/{job_id}/cancel",
+ params=None,
+ )
  try:
  # No need to make a dataclass for this since it's just a boolean.
  success: bool = response_json["success"]
@@ -171,6 +171,11 @@ def _refresh(
  ping_frequency: int = 5,
  performance: Optional[str] = None,
  ) -> str:
+ """
+ Executes a Dune `query`, waits until execution completes,
+ fetches and returns the results.
+ Sleeps `ping_frequency` seconds between each status request.
+ """
  job_id = self.execute(query=query, performance=performance).execution_id
  status = self.get_status(job_id)
  while status.state not in ExecutionState.terminal_states():
@@ -186,38 +191,28 @@ def _refresh(
  return job_id
 
  def refresh(
- self,
- query: Query,
- ping_frequency: int = 5,
- performance: Optional[str] = None,
+ self, query: Query, ping_frequency: int = 5, performance: Optional[str] = None
  ) -> ResultsResponse:
  """
  Executes a Dune `query`, waits until execution completes,
  fetches and returns the results.
  Sleeps `ping_frequency` seconds between each status request.
  """
  job_id = self._refresh(
- query,
- ping_frequency=ping_frequency,
- performance=performance,
+ query, ping_frequency=ping_frequency, performance=performance
  )
  return self.get_result(job_id)
 
  def refresh_csv(
- self,
- query: Query,
- ping_frequency: int = 5,
- performance: Optional[str] = None,
+ self, query: Query, ping_frequency: int = 5, performance: Optional[str] = None
  ) -> ExecutionResultCSV:
  """
  Executes a Dune query, waits till execution completes,
  fetches and the results in CSV format
  (use it load the data directly in pandas.from_csv() or similar frameworks)
  """
  job_id = self._refresh(
- query,
- ping_frequency=ping_frequency,
- performance=performance,
+ query, ping_frequency=ping_frequency, performance=performance
  )
  return self.get_result_csv(job_id)