feature: support multi-select params / array (#170)

Author: grkhr

dune_client/api/base.py

@@ -9,14 +9,17 @@
 import logging.config
 import os
 from json import JSONDecodeError
-from typing import IO, Any
+from typing import IO, TYPE_CHECKING, Any
 
 from deprecated import deprecated
 from requests import Response, Session
 from requests.adapters import HTTPAdapter, Retry
 
 from dune_client.util import get_package_version
 
+if TYPE_CHECKING:
+    from dune_client.types import QueryParameters
+
 # Headers used for pagination in CSV results
 DUNE_CSV_NEXT_URI_HEADER = "x-dune-next-uri"
 DUNE_CSV_NEXT_OFFSET_HEADER = "x-dune-next-offset"
@@ -94,15 +97,15 @@ def default_headers(self) -> dict[str, str]:
 
     def _build_parameters(
         self,
-        params: dict[str, str | int] | None = None,
+        params: QueryParameters | None = None,
         columns: list[str] | None = None,
         sample_count: int | None = None,
         filters: str | None = None,
         sort_by: list[str] | None = None,
         limit: int | None = None,
         offset: int | None = None,
         allow_partial_results: str = "true",
-    ) -> dict[str, str | int]:
+    ) -> QueryParameters:
         """
         Utility function that builds a dictionary of parameters to be used
         when retrieving advanced results (filters, pagination, sorting, etc.).
@@ -114,24 +117,24 @@ def _build_parameters(
             sample_count is None
             # We are sampling and don't use filters or pagination
             or (limit is None and offset is None and filters is None)
-        ), "sampling cannot be combined with filters or pagination"
+        ), "sapling cannot be combined with filters or pagination"
 
-        params = params or {}
-        params["allow_partial_results"] = allow_partial_results
-        if columns is not None and len(columns) > 0:
-            params["columns"] = ",".join(columns)
+        result: QueryParameters = dict(params) if params else {}
+        result["allow_partial_results"] = allow_partial_results
+        if columns:
+            result["columns"] = ",".join(columns)
         if sample_count is not None:
-            params["sample_count"] = sample_count
+            result["sample_count"] = sample_count
         if filters is not None:
-            params["filters"] = filters
-        if sort_by is not None and len(sort_by) > 0:
-            params["sort_by"] = ",".join(sort_by)
+            result["filters"] = filters
+        if sort_by:
+            result["sort_by"] = ",".join(sort_by)
         if limit is not None:
-            params["limit"] = limit
+            result["limit"] = limit
         if offset is not None:
-            params["offset"] = offset
+            result["offset"] = offset
 
-        return params
+        return result
 
 
 class BaseRouter(BaseDuneClient):

dune_client/query.py

@@ -4,27 +4,28 @@
 
 from __future__ import annotations
 
+import json
 import urllib.parse
 from dataclasses import dataclass
 from typing import Any
 
-from dune_client.types import QueryParameter
+from dune_client.types import QueryParameter, QueryParameters
 
 
 def parse_query_object_or_id(
     query: QueryBase | str | int,
-) -> tuple[dict[str, Any] | None, int]:
+) -> tuple[QueryParameters | None, int]:
     """
     Users are allowed to pass QueryBase or ID into some functions.
     This method handles both scenarios, returning a pair of the form (params, query_id)
     """
     if isinstance(query, QueryBase):
-        params = {f"params.{p.key}": p.to_dict()["value"] for p in query.parameters()}
-        query_id = query.query_id
-    else:
-        params = None
-        query_id = int(query)
-    return params, query_id
+        params: QueryParameters = {
+            f"params.{p.key}": p.to_dict()["value"] for p in query.parameters()
+        }
+        return params, query.query_id
+
+    return None, int(query)
 
 
 @dataclass
@@ -46,9 +47,17 @@ def parameters(self) -> list[QueryParameter]:
     def url(self) -> str:
         """Returns a parameterized link to the query"""
         # Include variable parameters in the URL so they are set
-        params = "&".join([f"{p.key}={p.value}" for p in self.parameters()])
-        if params:
-            return "?".join([self.base_url(), urllib.parse.quote_plus(params, safe="=&?")])
+        params = []
+        for parameter in self.parameters():
+            value = parameter.to_dict()["value"]
+            if isinstance(value, list):
+                serialized_value = json.dumps(value, separators=(",", ":"))
+            else:
+                serialized_value = value
+            params.append(f"{parameter.key}={serialized_value}")
+        param_string = "&".join(params)
+        if param_string:
+            return "?".join([self.base_url(), urllib.parse.quote_plus(param_string, safe="=&?")])
         return self.base_url()
 
     def __hash__(self) -> int:
@@ -58,7 +67,7 @@ def __hash__(self) -> int:
         """
         return self.url().__hash__()
 
-    def request_format(self) -> dict[str, dict[str, str] | str | None]:
+    def request_format(self) -> dict[str, str | QueryParameters]:
         """Transforms Query objects to params to pass in API"""
         return {"query_parameters": {p.key: p.to_dict()["value"] for p in self.parameters()}}
 

dune_client/types.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import re
+from collections.abc import Sequence
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
@@ -16,6 +17,7 @@
     from datetime import datetime
 
 DuneRecord = dict[str, Any]
+QueryParameters = dict[str, str | list[str] | int]
 
 
 class Address:
@@ -125,7 +127,12 @@ def __eq__(self, other: object) -> bool:
         )
 
     def __hash__(self) -> int:
-        return hash((self.key, self.value, self.type.value))
+        value = (
+            tuple(self.value)
+            if isinstance(self.value, Sequence) and not isinstance(self.value, str)
+            else self.value
+        )
+        return hash((self.key, value, self.type.value))
 
     @classmethod
     def text_type(cls, name: str, value: str) -> QueryParameter:
@@ -148,25 +155,31 @@ def date_type(cls, name: str, value: datetime | str) -> QueryParameter:
         return cls(name, ParameterType.DATE, value)
 
     @classmethod
-    def enum_type(cls, name: str, value: str) -> QueryParameter:
-        """Constructs a Query parameter of type number"""
-        return cls(name, ParameterType.ENUM, value)
+    def enum_type(cls, name: str, value: str | Sequence[str]) -> QueryParameter:
+        """Constructs a Query parameter of type enum or multi-select"""
+        if isinstance(value, str):
+            return cls(name, ParameterType.ENUM, value)
+        if isinstance(value, Sequence):
+            return cls(name, ParameterType.ENUM, tuple(value))
+        raise TypeError(f"Unsupported enum value type for parameter '{name}': {type(value)!r}")
 
-    def value_str(self) -> str:
-        """Returns string value of parameter"""
+    def serialized_value(self) -> str | list[str]:
+        """Returns JSON-ready value of parameter"""
         if self.type in (ParameterType.TEXT, ParameterType.NUMBER, ParameterType.ENUM):
+            if isinstance(self.value, Sequence) and not isinstance(self.value, str):
+                return [str(v) for v in self.value]
             return str(self.value)
         if self.type == ParameterType.DATE:
             # This is the postgres string format of timestamptz
             return str(self.value.strftime("%Y-%m-%d %H:%M:%S"))
         raise TypeError(f"Type {self.type} not recognized!")
 
-    def to_dict(self) -> dict[str, str]:
+    def to_dict(self) -> dict[str, str | list[str]]:
         """Converts QueryParameter into string json format accepted by Dune API"""
-        results: dict[str, str] = {
+        results: dict[str, str | list[str]] = {
             "key": self.key,
             "type": self.type.value,
-            "value": self.value_str(),
+            "value": self.serialized_value(),
         }
         return results
 

tests/unit/test_query.py

@@ -1,4 +1,5 @@
 import unittest
+import urllib.parse
 from datetime import datetime
 
 from dune_client.query import QueryBase, parse_query_object_or_id
@@ -13,17 +14,24 @@ def setUp(self) -> None:
             QueryParameter.text_type("Text", "plain text"),
             QueryParameter.number_type("Number", 12),
             QueryParameter.date_type("Date", "2021-01-01 12:34:56"),
+            QueryParameter.enum_type("Multi", ["a1", "a2"]),
         ]
         self.query = QueryBase(name="", query_id=0, params=self.query_params)
 
     def test_base_url(self):
         assert self.query.base_url() == "https://dune.com/queries/0"
 
     def test_url(self):
-        assert (
-            self.query.url()
-            == "https://dune.com/queries/0?Enum=option1&Text=plain+text&Number=12&Date=2021-01-01+12%3A34%3A56"
+        raw_params = (
+            'Enum=option1&Text=plain text&Number=12&Date=2021-01-01 12:34:56&Multi=["a1","a2"]'
         )
+        expected_url = "?".join(
+            [
+                "https://dune.com/queries/0",
+                urllib.parse.quote_plus(raw_params, safe="=&?"),
+            ]
+        )
+        assert self.query.url() == expected_url
         assert QueryBase(0, "", []).url() == "https://dune.com/queries/0"
 
     def test_parameters(self):
@@ -36,6 +44,7 @@ def test_request_format(self):
                 "Text": "plain text",
                 "Number": "12",
                 "Date": "2021-01-01 12:34:56",
+                "Multi": ["a1", "a2"],
             }
         }
         assert self.query.request_format() == expected_answer
@@ -62,6 +71,7 @@ def test_parse_object_or_id(self):
             "params.Enum": "option1",
             "params.Number": "12",
             "params.Text": "plain text",
+            "params.Multi": ["a1", "a2"],
         }
         expected_query_id = self.query.query_id
         # Query Object