diff --git a/.gitignore b/.gitignore index 79eff2b..c24fbba 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ build transpose_data_test* transpose_data_test*/* *.egg-info +transpose_data.egg-info/* +.eggs/* setup.py # keys diff --git a/setup.py b/setup.py index b6cb016..c2f4766 100644 --- a/setup.py +++ b/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'] ) diff --git a/tests/test_analytical.py b/tests/test_analytical.py index 27ae245..66df130 100644 --- a/tests/test_analytical.py +++ b/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 diff --git a/tests/test_sql.py b/tests/test_sql.py index 7e6818c..d1da237 100644 --- a/tests/test_sql.py +++ b/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) diff --git a/transpose/src/api/analytical/base.py b/transpose/src/api/analytical/base.py index fc62b25..dbf5a71 100644 --- a/transpose/src/api/analytical/base.py +++ b/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) diff --git a/transpose/src/api/sql/base.py b/transpose/src/api/sql/base.py index 6d79578..a394ff5 100644 --- a/transpose/src/api/sql/base.py +++ b/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: diff --git a/transpose/src/util/client.py b/transpose/src/util/client.py index cfc139c..24b7ae5 100644 --- a/transpose/src/util/client.py +++ b/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,7 +61,7 @@ def get_api_request( params=params ) - return handle_response(request) + return handle_response(request, return_df) def post_api_request( @@ -55,8 +69,9 @@ def post_api_request( 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)