From 4c2c256c805a6e5cda3d149acc9f5512eb70ea37 Mon Sep 17 00:00:00 2001
From: Benjamin Smith <bh2smith@users.noreply.github.com>
Date: Wed, 6 Sep 2023 14:15:12 +0200
Subject: [PATCH] Add upload_csv route (#61)

* add upload/csv route
* add better response handling
---
 dune_client/base_client.py |  1 -
 dune_client/client.py      | 23 +++++++++++++++++++++++
 tests/e2e/test_client.py   | 12 +++++++++++-
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/dune_client/base_client.py b/dune_client/base_client.py
index 3c91cd2..271153d 100644
--- a/dune_client/base_client.py
+++ b/dune_client/base_client.py
@@ -9,7 +9,6 @@
 from typing import Dict
 
 
-# pylint: disable=too-few-public-methods
 class BaseDuneClient:
     """
     A Base Client for Dune which sets up default values
diff --git a/dune_client/client.py b/dune_client/client.py
index d24ae48..1d391bd 100644
--- a/dune_client/client.py
+++ b/dune_client/client.py
@@ -374,3 +374,26 @@ def make_public(self, query_id: int) -> None:
         """
         response_json = self._post(route=f"/query/{query_id}/unprivate")
         assert not self.get_query(int(response_json["query_id"])).meta.is_private
+
+    def upload_csv(self, table_name: str, data: str, description: str = "") -> bool:
+        """
+        https://dune.com/docs/api/api-reference/upload-data/?h=data+upload#endpoint
+        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.
+
+        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,
+            },
+        )
+        try:
+            return bool(response_json["success"])
+        except KeyError as err:
+            raise DuneError(response_json, "upload_csv response", err) from err
diff --git a/tests/e2e/test_client.py b/tests/e2e/test_client.py
index e07da85..bde0959 100644
--- a/tests/e2e/test_client.py
+++ b/tests/e2e/test_client.py
@@ -156,7 +156,6 @@ def test_internal_error(self):
 
     def test_invalid_job_id_error(self):
         dune = DuneClient(self.valid_api_key)
-
         with self.assertRaises(DuneError) as err:
             dune.get_status("Wonky Job ID")
         self.assertEqual(
@@ -175,6 +174,17 @@ def test_get_latest_result_with_query_id(self):
         results = dune.get_latest_result(self.query.query_id).get_rows()
         self.assertGreater(len(results), 0)
 
+    def test_upload_csv_success(self):
+        client = DuneClient(self.valid_api_key)
+        self.assertEqual(
+            client.upload_csv(
+                table_name="e2e-test",
+                description="best data",
+                data="column1,column2\nvalue1,value2\nvalue3,value4",
+            ),
+            True,
+        )
+
 
 @unittest.skip("This is an enterprise only endpoint that can no longer be tested.")
 class TestCRUDOps(unittest.TestCase):