diff --git a/dune_client/api/base.py b/dune_client/api/base.py index ddb3ed4..c00b016 100644 --- a/dune_client/api/base.py +++ b/dune_client/api/base.py @@ -9,7 +9,7 @@ import logging.config import os from json import JSONDecodeError -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union, IO from requests import Response, Session from requests.adapters import HTTPAdapter, Retry @@ -179,15 +179,22 @@ def _get( return response return self._handle_response(response) - def _post(self, route: str, params: Optional[Any] = None) -> Any: + def _post( + self, + route: str, + params: Optional[Any] = None, + data: Optional[IO[bytes]] = None, + headers: Optional[Dict[str, str]] = None, + ) -> Any: """Generic interface for the POST method of a Dune API request""" url = self._route_url(route) self.logger.debug(f"POST received input url={url}, params={params}") response = self.http.post( url=url, json=params, - headers=self.default_headers(), + headers=dict(self.default_headers(), **headers if headers else {}), timeout=self.request_timeout, + data=data, ) return self._handle_response(response) diff --git a/dune_client/api/extensions.py b/dune_client/api/extensions.py index c5f3047..e4c074c 100644 --- a/dune_client/api/extensions.py +++ b/dune_client/api/extensions.py @@ -19,6 +19,7 @@ ) from dune_client.api.execution import ExecutionAPI from dune_client.api.query import QueryAPI +from dune_client.api.table import TableAPI from dune_client.models import ( ResultsResponse, DuneError, @@ -36,7 +37,7 @@ POLL_FREQUENCY_SECONDS = 1 -class ExtendedAPI(ExecutionAPI, QueryAPI): +class ExtendedAPI(ExecutionAPI, QueryAPI, TableAPI): """ Provides higher level helper methods for faster and easier development on top of the base ExecutionAPI. @@ -316,40 +317,6 @@ def download_csv( ), ) - ############################ - # Plus Subscription Features - ############################ - def upload_csv( - self, - table_name: str, - data: str, - description: str = "", - is_private: bool = False, - ) -> bool: - """ - https://docs.dune.com/api-reference/tables/endpoint/upload - The write API allows you to upload any .csv file into Dune. The only limitations are: - - - File has to be < 200 MB - - Column names in the table can't start with a special character or digits. - - Private uploads require a Plus subscription. - - Below are the specifics of how to work with the API. - """ - response_json = self._post( - route="/table/upload/csv", - params={ - "table_name": table_name, - "description": description, - "data": data, - "is_private": is_private, - }, - ) - try: - return bool(response_json["success"]) - except KeyError as err: - raise DuneError(response_json, "UploadCsvResponse", err) from err - ############################################################################################## # Plus Features: these features use APIs that are only available on paid subscription plans ############################################################################################## diff --git a/dune_client/api/table.py b/dune_client/api/table.py new file mode 100644 index 0000000..ed1f9a3 --- /dev/null +++ b/dune_client/api/table.py @@ -0,0 +1,99 @@ +""" +Table API endpoints enables users to +create and insert data into Dune. +""" + +from __future__ import annotations +from typing import List, Dict, Any, IO + +from dune_client.api.base import BaseRouter +from dune_client.models import DuneError + + +class TableAPI(BaseRouter): + """ + Implementation of Table endpoints - Plus subscription only + https://docs.dune.com/api-reference/tables/ + """ + + def upload_csv( + self, + table_name: str, + data: str, + description: str = "", + is_private: bool = False, + ) -> bool: + """ + https://docs.dune.com/api-reference/tables/endpoint/upload + This endpoint allows you to upload any .csv file into Dune. The only limitations are: + + - File has to be < 200 MB + - Column names in the table can't start with a special character or digits. + - Private uploads require a Plus subscription. + + Below are the specifics of how to work with the API. + """ + response_json = self._post( + route="/table/upload/csv", + params={ + "table_name": table_name, + "description": description, + "data": data, + "is_private": is_private, + }, + ) + try: + return bool(response_json["success"]) + except KeyError as err: + raise DuneError(response_json, "UploadCsvResponse", err) from err + + def create_table( + self, + namespace: str, + table_name: str, + schema: List[Dict[str, str]], + description: str = "", + is_private: bool = False, + ) -> Any: + """ + https://docs.dune.com/api-reference/tables/endpoint/create + The create table endpoint allows you to create an empty table + with a specific schema in Dune. + + The only limitations are: + - If a table already exists with the same name, the request will fail. + - Column names in the table can’t start with a special character or a digit. + """ + + return self._post( + route="/table/create", + params={ + "namespace": namespace, + "table_name": table_name, + "schema": schema, + "description": description, + "is_private": is_private, + }, + ) + + def insert_table( + self, + namespace: str, + table_name: str, + data: IO[bytes], + content_type: str, + ) -> Any: + """ + https://docs.dune.com/api-reference/tables/endpoint/insert + The insert table endpoint allows you to insert data into an existing table in Dune. + + The only limitations are: + - The file has to be in json or csv format + - The file has to have the same schema as the table + """ + + return self._post( + route=f"/table/{namespace}/{table_name}/insert", + headers={"Content-Type": content_type}, + data=data, + ) diff --git a/tests/e2e/test_async_client.py b/tests/e2e/test_async_client.py index 20ac35a..1d9b97f 100644 --- a/tests/e2e/test_async_client.py +++ b/tests/e2e/test_async_client.py @@ -41,7 +41,6 @@ async def test_refresh_context_manager(self): async def test_refresh_with_pagination(self): # Arrange async with AsyncDuneClient(self.valid_api_key) as cl: - # Act results = (await cl.refresh(self.multi_rows_query, batch_size=1)).get_rows() @@ -60,7 +59,6 @@ async def test_refresh_with_pagination(self): async def test_refresh_with_filters(self): # Arrange async with AsyncDuneClient(self.valid_api_key) as cl: - # Act results = ( await cl.refresh(self.multi_rows_query, filters="number < 3") @@ -78,7 +76,6 @@ async def test_refresh_with_filters(self): async def test_refresh_csv_with_pagination(self): # Arrange async with AsyncDuneClient(self.valid_api_key) as cl: - # Act result_csv = await cl.refresh_csv(self.multi_rows_query, batch_size=1) @@ -97,7 +94,6 @@ async def test_refresh_csv_with_pagination(self): async def test_refresh_csv_with_filters(self): # Arrange async with AsyncDuneClient(self.valid_api_key) as cl: - # Act result_csv = await cl.refresh_csv( self.multi_rows_query, filters="number < 3" diff --git a/tests/e2e/test_client.py b/tests/e2e/test_client.py index 07882ac..426c9d4 100644 --- a/tests/e2e/test_client.py +++ b/tests/e2e/test_client.py @@ -236,6 +236,66 @@ def test_upload_csv_success(self): True, ) + @unittest.skip("This is a plus subscription endpoint.") + def test_create_table_success(self): + # Make sure the table doesn't already exist. + # You will need to change the namespace to your own. + client = DuneClient(self.valid_api_key) + + namespace = "test" + table_name = "dataset_e2e_test" + + self.assertEqual( + client.create_table( + namespace=namespace, + table_name=table_name, + description="e2e test table", + schema=[ + {"name": "date", "type": "timestamp"}, + {"name": "dgs10", "type": "double"}, + ], + 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", + }, + ) + + @unittest.skip("This is a plus subscription endpoint.") + def test_insert_table_csv_success(self): + # Make sure the table already exists and csv matches table schema. + # You will need to change the namespace to your own. + client = DuneClient(self.valid_api_key) + with open("./tests/fixtures/sample_table_insert.csv", "rb") as data: + self.assertEqual( + client.insert_table( + namespace="test", + table_name="dataset_e2e_test", + data=data, + content_type="text/csv", + ), + None, + ) + + @unittest.skip("This is a plus subscription endpoint.") + def test_insert_table_json_success(self): + # Make sure the table already exists and json matches table schema. + # You will need to change the namespace to your own. + client = DuneClient(self.valid_api_key) + with open("./tests/fixtures/sample_table_insert.json", "rb") as data: + self.assertEqual( + client.insert_table( + namespace="test", + table_name="dataset_e2e_test", + data=data, + content_type="application/x-ndjson", + ), + None, + ) + def test_download_csv_with_pagination(self): # Arrange client = DuneClient(self.valid_api_key) diff --git a/tests/fixtures/sample_table_insert.csv b/tests/fixtures/sample_table_insert.csv new file mode 100644 index 0000000..70b2f9e --- /dev/null +++ b/tests/fixtures/sample_table_insert.csv @@ -0,0 +1,2 @@ +date,dgs10 +2020-12-01T23:33:00,10 \ No newline at end of file diff --git a/tests/fixtures/sample_table_insert.json b/tests/fixtures/sample_table_insert.json new file mode 100644 index 0000000..f2b93f0 --- /dev/null +++ b/tests/fixtures/sample_table_insert.json @@ -0,0 +1 @@ +{"date":"2020-12-01T23:33:00","dgs10":10} \ No newline at end of file