TRA-4675: restore return df param

.gitignore

@@ -14,6 +14,8 @@ build
 transpose_data_test*
 transpose_data_test*/*
 *.egg-info
+transpose_data.egg-info/*
+.eggs/*
 setup.py
 
 # keys

setup.py

@@ -8,7 +8,7 @@
 
     # version compliant with PEP440
     # https://peps.python.org/pep-0440/
-    version='4.2.0',
+    version='4.2.2',
 
     # project meta
     long_description = long_description,
@@ -45,4 +45,5 @@
         'pandas',
         'python-dotenv',
     ],
+    setup_requires=['wheel']
 )

tests/test_analytical.py

@@ -80,3 +80,16 @@ def test_query_df():
 
     except Exception:
         assert False
+
+
+def test_return_df():
+    try:
+        api = Transpose(api_key)
+
+        response = api.analytical.query("SELECT * FROM cross_chain.transaction_flows LIMIT 1;", return_df=True)
+
+        assert type(response) is DataFrame
+        assert len(response) == 1
+
+    except Exception:
+        assert False

tests/test_sql.py

@@ -60,6 +60,20 @@ def test_query_df():
         assert False
 
 
+def test_return_df():
+    try:
+        api = Transpose(api_key)
+
+        query = "SELECT * FROM ethereum.logs LIMIT 100;"
+        response = api.sql.query(query, return_df=True)
+
+        assert type(response) is DataFrame
+        assert len(response) == 100
+
+    except Exception:
+        assert False
+
+
 def test_schema():
     try:
         api = Transpose(api_key)

transpose/src/api/analytical/base.py

@@ -1,3 +1,7 @@
+from typing import Union
+
+import pandas as pd
+
 from ...util.models import QueryResult
 from ....src.util.client import post_api_request
 
@@ -12,7 +16,8 @@ def query(
         self,
         sql_query: str,
         parameters: dict = None,
-    ) -> QueryResult:
+        return_df: bool = False
+    ) -> Union[QueryResult, pd.DataFrame]:
 
         parameters = {} if parameters is None else parameters
 
@@ -26,7 +31,8 @@ def query(
             url=url,
             api_key=self.super.api_key,
             body=body,
+            return_df=return_df,
             verbose=self.super.verbose
         )
 
-        return QueryResult(result)
+        return result if return_df else QueryResult(result)

transpose/src/api/sql/base.py

@@ -1,3 +1,7 @@
+from typing import Union
+
+import pandas as pd
+
 from ...util.models import QueryResult
 from ....src.util.client import get_api_request, post_api_request
 
@@ -11,8 +15,12 @@ def __init__(self, base_class) -> None:
     def query(
         self,
         sql_query: str,
-        parameters: dict = {},
-    ) -> QueryResult:
+        parameters: dict = None,
+        return_df: bool = False
+    ) -> Union[QueryResult, pd.DataFrame]:
+
+        if parameters is None:
+            parameters = {}
 
         url = "https://api.transpose.io/sql"
         body = {
@@ -24,10 +32,11 @@ def query(
             url=url,
             api_key=self.super.api_key,
             body=body,
+            return_df=return_df,
             verbose=self.super.verbose
         )
 
-        return QueryResult(result)
+        return result if return_df else QueryResult(result)
 
     # Gets the schema from the Transpose API
     def schema(self) -> dict:

transpose/src/util/client.py

@@ -1,4 +1,7 @@
 import json
+from typing import Union
+
+import pandas as pd
 import requests
 
 from transpose.src.util.errors import raise_custom_error
@@ -12,9 +15,19 @@ def build_headers(api_key: str) -> dict:
     }
 
 
-def handle_response(request: requests.Response) -> dict:
+def handle_response(request: requests.Response, return_df: bool = False) -> Union[dict, pd.DataFrame]:
 
     if request.status_code == 200:
+
+        # return the response as a DataFrame
+        if return_df:
+
+            # check if pandas is installed
+            if not pd:
+                raise ImportError("Pandas is not installed. Please install pandas to use this feature.")
+
+            return pd.DataFrame(request.json()['results'])
+
         # return the response as a dictionary
         return request.json()
 
@@ -28,8 +41,9 @@ def get_api_request(
     api_key: str,
     body: dict = None,
     params=None,
+    return_df: bool = False,
     verbose: bool = False
-) -> dict:
+) -> Union[dict, pd.DataFrame]:
 
     # set body/parameters to an empty dictionary if not provided
     body = {} if body is None else body
@@ -47,16 +61,17 @@ def get_api_request(
         params=params
     )
 
-    return handle_response(request)
+    return handle_response(request, return_df)
 
 
 def post_api_request(
     url: str,
     api_key: str,
     body: dict,
     params=None,
+    return_df: bool = False,
     verbose: bool = False
-) -> dict:
+) -> Union[dict, pd.DataFrame]:
 
     # set body/parameters to an empty dictionary if not provided
     body = {} if body is None else body
@@ -74,4 +89,4 @@ def post_api_request(
         params=params
     )
 
-    return handle_response(request)
+    return handle_response(request, return_df)