[Premium Feature] run_sql (#78)

* add premium feature run_sql

dune_client/api/extensions.py

@@ -1,14 +1,15 @@
 """
 Extended functionality for the ExecutionAPI
 """
-
+from __future__ import annotations
 import time
 from io import BytesIO
 from typing import Union, Optional, Any
 
 from deprecated import deprecated
 
 from dune_client.api.execution import ExecutionAPI
+from dune_client.api.query import QueryAPI
 from dune_client.models import (
     ResultsResponse,
     DuneError,
@@ -17,9 +18,10 @@
     ExecutionResultCSV,
 )
 from dune_client.query import QueryBase, parse_query_object_or_id
+from dune_client.types import QueryParameter
 
 
-class ExtendedAPI(ExecutionAPI):
+class ExtendedAPI(ExecutionAPI, QueryAPI):
     """
     Provides higher level helper methods for faster
     and easier development on top of the base ExecutionAPI.
@@ -131,6 +133,29 @@ def upload_csv(self, table_name: str, data: str, description: str = "") -> bool:
         except KeyError as err:
             raise DuneError(response_json, "UploadCsvResponse", err) from err
 
+    ##############################################################################################
+    # Premium Features: these features use APIs that are only available on paid subscription plans
+    ##############################################################################################
+
+    def run_sql(
+        self,
+        query_sql: str,
+        params: Optional[list[QueryParameter]] = None,
+        is_private: bool = True,
+        archive_after: bool = True,
+    ) -> ResultsResponse:
+        """
+        Allows user to provide execute raw_sql via the CRUD interface
+        - create, run, get results with optional archive/delete.
+        - Query is by default made private and archived after execution.
+        Requires premium subscription!
+        """
+        query = self.create_query("API Query", query_sql, params, is_private)
+        results = self.run_query(query.base)
+        if archive_after:
+            self.archive_query(query.base.query_id)
+        return results
+
     ######################
     # Deprecated Functions
     ######################

dune_client/client.py

@@ -4,10 +4,9 @@
 https://duneanalytics.notion.site/API-Documentation-1b93d16e0fa941398e15047f643e003a
 """
 from dune_client.api.extensions import ExtendedAPI
-from dune_client.api.query import QueryAPI
 
 
-class DuneClient(QueryAPI, ExtendedAPI):
+class DuneClient(ExtendedAPI):
     """
     An interface for Dune API with a few convenience methods
     combining the use of endpoints (e.g. run_query)
@@ -16,13 +15,14 @@ class DuneClient(QueryAPI, ExtendedAPI):
 
         DuneClient
         |
-        |--- QueryAPI(BaseRouter)
-        |       - Contains CRUD Operations on Queries
-        |
         |--- ExtendedAPI
-                | - Contains higher level methods for higher productivity
-                |       (things like `run_query`, `run_query_csv`, etc..)
+                |   - Contains compositions of execution methods like `run_query`
+                |               (things like `run_query`, `run_query_csv`, etc..)
+                |   - make use of both Execution and Query APIs
                 |
                 |--- ExecutionAPI(BaseRouter)
-                        - Contains query execution and result related methods.
+                |        - Contains query execution methods.
+                |
+                |--- QueryAPI(BaseRouter)
+                |       - Contains CRUD Operations on Queries
     """

tests/e2e/test_client.py

@@ -278,6 +278,17 @@ def test_archive(self):
         self.assertEqual(self.client.archive_query(self.existing_query_id), True)
         self.assertEqual(self.client.unarchive_query(self.existing_query_id), False)
 
+    @unittest.skip("Works fine, but creates too many queries!")
+    def test_run_sql(self):
+        query_sql = "select 85"
+        results = self.client.run_sql(query_sql)
+        self.assertEqual(results.get_rows(), [{"_col0": 85}])
+
+        # The default functionality is meant to create a private query and then archive it.
+        query = self.client.get_query(results.query_id)
+        self.assertTrue(query.meta.is_archived)
+        self.assertTrue(query.meta.is_private)
+
 
 if __name__ == "__main__":
     unittest.main()