diff --git a/dune_client/api/table.py b/dune_client/api/table.py index ed1f9a3..ceca4e7 100644 --- a/dune_client/api/table.py +++ b/dune_client/api/table.py @@ -4,10 +4,10 @@ """ from __future__ import annotations -from typing import List, Dict, Any, IO +from typing import List, Dict, IO from dune_client.api.base import BaseRouter -from dune_client.models import DuneError +from dune_client.models import DuneError, InsertTableResult, CreateTableResult class TableAPI(BaseRouter): @@ -54,7 +54,7 @@ def create_table( schema: List[Dict[str, str]], description: str = "", is_private: bool = False, - ) -> Any: + ) -> CreateTableResult: """ https://docs.dune.com/api-reference/tables/endpoint/create The create table endpoint allows you to create an empty table @@ -65,7 +65,7 @@ def create_table( - Column names in the table can’t start with a special character or a digit. """ - return self._post( + result_json = self._post( route="/table/create", params={ "namespace": namespace, @@ -75,6 +75,7 @@ def create_table( "is_private": is_private, }, ) + return CreateTableResult.from_dict(result_json) def insert_table( self, @@ -82,7 +83,7 @@ def insert_table( table_name: str, data: IO[bytes], content_type: str, - ) -> Any: + ) -> InsertTableResult: """ https://docs.dune.com/api-reference/tables/endpoint/insert The insert table endpoint allows you to insert data into an existing table in Dune. @@ -92,8 +93,9 @@ def insert_table( - The file has to have the same schema as the table """ - return self._post( + result_json = self._post( route=f"/table/{namespace}/{table_name}/insert", headers={"Content-Type": content_type}, data=data, ) + return InsertTableResult.from_dict(result_json) diff --git a/dune_client/models.py b/dune_client/models.py index c19a174..3c05ae8 100644 --- a/dune_client/models.py +++ b/dune_client/models.py @@ -11,7 +11,7 @@ from io import BytesIO from os import SEEK_END from typing import Optional, Any, Union, List, Dict - +from dataclasses_json import DataClassJsonMixin from dateutil.parser import parse from dune_client.types import DuneRecord @@ -354,3 +354,24 @@ def __add__(self, other: ResultsResponse) -> ResultsResponse: self.next_uri = other.next_uri self.next_offset = other.next_offset return self + + +@dataclass +class CreateTableResult(DataClassJsonMixin): + """ + Data type returned by table/create operation + """ + + example_query: str + full_name: str + namespace: str + table_name: str + + +@dataclass +class InsertTableResult(DataClassJsonMixin): + """ + Data type returned by table/insert operation + """ + + rows_written: int diff --git a/requirements/prod.txt b/requirements/prod.txt index ebd6c5c..a25f63b 100644 --- a/requirements/prod.txt +++ b/requirements/prod.txt @@ -1,4 +1,5 @@ aiohttp>=3.8.5 +dataclasses-json==0.6.4 types-python-dateutil>=2.8.19.14 types-PyYAML>=6.0.12.11 types-requests>=2.28.0 diff --git a/tests/e2e/test_client.py b/tests/e2e/test_client.py index 4fdc7a5..3c3dcde 100644 --- a/tests/e2e/test_client.py +++ b/tests/e2e/test_client.py @@ -11,6 +11,8 @@ ExecutionResponse, ExecutionStatusResponse, DuneError, + InsertTableResult, + CreateTableResult, ) from dune_client.types import QueryParameter from dune_client.client import DuneClient @@ -125,7 +127,7 @@ def test_parameters_recognized(self): { "text_field": "different word", "number_field": 22, - "date_field": "1991-01-01T00:00:00Z", + "date_field": "1991-01-01 00:00:00.000", "list_field": "Option 2", } ], @@ -256,12 +258,14 @@ def test_create_table_success(self): ], is_private=False, ), - { - "namespace": namespace, - "table_name": table_name, - "full_name": f"dune.{namespace}.{table_name}", - "example_query": f"select * from dune.{namespace}.{table_name} limit 10", - }, + CreateTableResult.from_dict( + { + "namespace": namespace, + "table_name": table_name, + "full_name": f"dune.{namespace}.{table_name}", + "example_query": f"select * from dune.{namespace}.{table_name} limit 10", + } + ), ) @unittest.skip("This is a plus subscription endpoint.") @@ -277,7 +281,7 @@ def test_insert_table_csv_success(self): data=data, content_type="text/csv", ), - {"rows_written": 1}, + InsertTableResult(rows_written=1), ) @unittest.skip("This is a plus subscription endpoint.") @@ -293,7 +297,7 @@ def test_insert_table_json_success(self): data=data, content_type="application/x-ndjson", ), - {"rows_written": 1}, + InsertTableResult(rows_written=1), ) def test_download_csv_with_pagination(self): @@ -350,7 +354,7 @@ def test_download_csv_success_by_id(self): { "text_field": "different word", "number_field": 22, - "date_field": "1991-01-01T00:00:00Z", + "date_field": "1991-01-01 00:00:00.000", "list_field": "Option 2", } ], @@ -372,7 +376,7 @@ def test_download_csv_success_with_params(self): pandas.read_csv(result_csv.data).to_dict(orient="records"), [ { - "date_field": "2022-05-04T00:00:00Z", + "date_field": "2022-05-04 00:00:00.000", "list_field": "Option 1", "number_field": 3.1415926535, "text_field": "Plain Text",