Merge pull request #114 from duneanalytics/PRO-937-add-table-endpoint…

…s-to-sdk

PRO-937: Add `create` and `insert` table endpoints to sdk

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)
 

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
     ##############################################################################################

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,
+        )

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"

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)

tests/fixtures/sample_table_insert.csv

@@ -0,0 +1,2 @@
+date,dgs10
+2020-12-01T23:33:00,10

tests/fixtures/sample_table_insert.json

@@ -0,0 +1 @@
+{"date":"2020-12-01T23:33:00","dgs10":10}