make error classes more consistent

duneanalytics · Sep 8, 2023 · 4567baa · 4567baa
1 parent 1d8b6af
commit 4567baa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/dune_client/api/extensions.py b/dune_client/api/extensions.py
@@ -122,7 +122,7 @@ def upload_csv(self, table_name: str, data: str, description: str = "") -> bool:
         try:
             return bool(response_json["success"])
         except KeyError as err:
-            raise DuneError(response_json, "upload_csv response", err) from err
+            raise DuneError(response_json, "UploadCsvResponse", err) from err
 
     ######################
     # Deprecated Functions

diff --git a/dune_client/api/query.py b/dune_client/api/query.py
@@ -43,7 +43,7 @@ def create_query(
             # Note that this requires an extra request.
             return self.get_query(query_id)
         except KeyError as err:
-            raise DuneError(response_json, "create_query Response", err) from err
+            raise DuneError(response_json, "CreateQueryResponse", err) from err
 
     def get_query(self, query_id: int) -> DuneQuery:
         """
@@ -96,7 +96,7 @@ def update_query(  # pylint: disable=too-many-arguments
             # No need to make a dataclass for this since it's just a boolean.
             return int(response_json["query_id"])
         except KeyError as err:
-            raise DuneError(response_json, "update_query Response", err) from err
+            raise DuneError(response_json, "UpdateQueryResponse", err) from err
 
     def archive_query(self, query_id: int) -> bool:
         """
@@ -108,7 +108,7 @@ def archive_query(self, query_id: int) -> bool:
             # No need to make a dataclass for this since it's just a boolean.
             return self.get_query(int(response_json["query_id"])).meta.is_archived
         except KeyError as err:
-            raise DuneError(response_json, "make_private Response", err) from err
+            raise DuneError(response_json, "ArchiveQueryResponse", err) from err
 
     def unarchive_query(self, query_id: int) -> bool:
         """
@@ -120,18 +120,24 @@ def unarchive_query(self, query_id: int) -> bool:
             # No need to make a dataclass for this since it's just a boolean.
             return self.get_query(int(response_json["query_id"])).meta.is_archived
         except KeyError as err:
-            raise DuneError(response_json, "make_private Response", err) from err
+            raise DuneError(response_json, "UnarchiveQueryResponse", err) from err
 
     def make_private(self, query_id: int) -> None:
         """
         https://dune.com/docs/api/api-reference/edit-queries/private-query
         """
         response_json = self._post(route=f"/query/{query_id}/private")
-        assert self.get_query(int(response_json["query_id"])).meta.is_private
+        try:
+            assert self.get_query(int(response_json["query_id"])).meta.is_private
+        except KeyError as err:
+            raise DuneError(response_json, "MakePrivateResponse", err) from err
 
     def make_public(self, query_id: int) -> None:
         """
         https://dune.com/docs/api/api-reference/edit-queries/private-query
         """
         response_json = self._post(route=f"/query/{query_id}/unprivate")
-        assert not self.get_query(int(response_json["query_id"])).meta.is_private
+        try:
+            assert not self.get_query(int(response_json["query_id"])).meta.is_private
+        except KeyError as err:
+            raise DuneError(response_json, "MakePublicResponse", err) from err