diff --git a/python/requirements.txt b/python/requirements.txt index 1060483..5441f68 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -1,3 +1,3 @@ pydantic==1.9.1 requests==2.28.1 -urllib3==1.26.11 +urllib3==1.26.15 diff --git a/python/shroomdk/errors/__init__.py b/python/shroomdk/errors/__init__.py index fbd6d5a..1ad6245 100644 --- a/python/shroomdk/errors/__init__.py +++ b/python/shroomdk/errors/__init__.py @@ -2,6 +2,7 @@ from .not_found_error import NotFoundError # noqa: F401 from .query_run_errors import ( # noqa: F401 QueryRunCancelledError, QueryRunExecutionError, + QueryRunInvalidStateToCancel, QueryRunRateLimitError, QueryRunTimeoutError, ) diff --git a/python/shroomdk/errors/query_run_errors.py b/python/shroomdk/errors/query_run_errors.py index a334fb8..64cdd48 100644 --- a/python/shroomdk/errors/query_run_errors.py +++ b/python/shroomdk/errors/query_run_errors.py @@ -51,3 +51,16 @@ class QueryRunCancelledError(BaseError): ): self.message = f"QUERY_RUN_CANCELLED_ERROR: your query has been cancelled. errorName={error_name}, errorMessage={error_message}, errorData={error_data}" super().__init__(self.message) + + +class QueryRunInvalidStateToCancel(BaseError): + """ + Base class for all QueryRunInvalidStateToCancel errors. + """ + + def __init__( + self, + msg: Optional[str] = None, + ): + self.message = msg + super().__init__(self.message) diff --git a/python/shroomdk/integrations/query_integration/compass_query_integration.py b/python/shroomdk/integrations/query_integration/compass_query_integration.py index e5af70b..2e21a57 100644 --- a/python/shroomdk/integrations/query_integration/compass_query_integration.py +++ b/python/shroomdk/integrations/query_integration/compass_query_integration.py @@ -9,6 +9,7 @@ from shroomdk.errors import ( SDKError, UserError, ) +from shroomdk.errors.query_run_errors import QueryRunInvalidStateToCancel from shroomdk.models import ( Query, QueryDefaults, @@ -16,6 +17,7 @@ from shroomdk.models import ( QueryStatus, SleepConfig, ) +from shroomdk.models.compass.cancel_query_run import CancelQueryRunRpcRequestParams from shroomdk.models.compass.core.page import Page from shroomdk.models.compass.core.query_run import QueryRun from shroomdk.models.compass.core.result_format import ResultFormat @@ -112,6 +114,19 @@ class CompassQueryIntegration(object): return response.result.queryRun + def cancel_query_run(self, query_run_id: str) -> QueryRun: + response = self.rpc.cancel_query_run( + CancelQueryRunRpcRequestParams(queryRunId=query_run_id) + ) + + if response.error or not response.result: + if response.error and response.error.code == -32165: + raise QueryRunInvalidStateToCancel(response.error.message) + + raise NotFoundError(f"QueryRun<{query_run_id}> not found") + + return response.result.queryRun + def get_query_results( self, query_run_id: str, @@ -146,16 +161,6 @@ class CompassQueryIntegration(object): filters: Optional[Union[List[Filter], None]] = [], sort_by: Optional[Union[List[SortBy], None]] = [], ) -> GetQueryRunResultsRpcResult: - # f2 = [] - # if filters: - # for f in filters: - # d = f.dict() - # d2 = {} - # for k, v in d.items(): - # if v is not None: - # d2[k] = v - # f2.append(Filter(**d2)) - query_results_resp = self.rpc.get_query_result( GetQueryRunResultsRpcParams( queryRunId=query_run_id, diff --git a/python/shroomdk/models/compass/cancel_query_run.py b/python/shroomdk/models/compass/cancel_query_run.py new file mode 100644 index 0000000..cf91ce6 --- /dev/null +++ b/python/shroomdk/models/compass/cancel_query_run.py @@ -0,0 +1,26 @@ +from typing import List, Union + +from pydantic import BaseModel + +from .core.query_run import QueryRun +from .core.rpc_request import RpcRequest +from .core.rpc_response import RpcResponse + + +# Request +class CancelQueryRunRpcRequestParams(BaseModel): + queryRunId: str + + +class CancelQueryRunRpcRequest(RpcRequest): + method: str = "cancelQueryRun" + params: List[CancelQueryRunRpcRequestParams] + + +# Response +class CancelQueryRunRpcResult(BaseModel): + queryRun: QueryRun + + +class CancelQueryRunRpcResponse(RpcResponse): + result: Union[CancelQueryRunRpcResult, None] diff --git a/python/shroomdk/rpc.py b/python/shroomdk/rpc.py index 4ad5fcd..ea9ded4 100644 --- a/python/shroomdk/rpc.py +++ b/python/shroomdk/rpc.py @@ -4,6 +4,11 @@ from typing import List import requests from requests.adapters import HTTPAdapter, Retry +from shroomdk.models.compass.cancel_query_run import ( + CancelQueryRunRpcRequest, + CancelQueryRunRpcRequestParams, + CancelQueryRunRpcResponse, +) from shroomdk.models.compass.get_sql_statement import ( GetSqlStatementParams, GetSqlStatementRequest, @@ -101,6 +106,21 @@ class RPC(object): return get_sql_statement_resp + def cancel_query_run( + self, params: CancelQueryRunRpcRequestParams + ) -> CancelQueryRunRpcResponse: + result = self._session.post( + self.url, + data=json.dumps(CancelQueryRunRpcRequest(params=[params]).dict()), + headers=self._headers, + ) + + data = self._handle_response(result, "cancelQueryRun") + + cancel_query_run_resp = CancelQueryRunRpcResponse(**data) + + return cancel_query_run_resp + def get_query_result( self, params: GetQueryRunResultsRpcParams ) -> GetQueryRunResultsRpcResponse: @@ -161,10 +181,10 @@ class RPC(object): allowed_methods=self._METHOD_ALLOWLIST, ) - adapter = HTTPAdapter(max_retries=retry_strategy) # type: ignore - http = requests.Session() - http.mount("https://", adapter) - http.mount("http://", adapter) + adapter = HTTPAdapter(max_retries=retry_strategy, pool_connections=1, pool_maxsize=1) # type: ignore + session = requests.Session() + session.mount("https://", adapter) + session.mount("http://", adapter) - self.__session = http + self.__session = session return self.__session diff --git a/python/shroomdk/shroomdk.py b/python/shroomdk/shroomdk.py index 1051ad7..ea0039f 100644 --- a/python/shroomdk/shroomdk.py +++ b/python/shroomdk/shroomdk.py @@ -3,16 +3,14 @@ from typing import List, Optional, Union from shroomdk.integrations.query_integration.compass_query_integration import ( CompassQueryIntegration, ) -from shroomdk.models import Query +from shroomdk.models import Filter, Query, SortBy from shroomdk.models.compass.core.query_run import QueryRun from shroomdk.models.compass.core.sql_statement import SqlStatement from shroomdk.models.compass.get_sql_statement import GetSqlStatementParams from shroomdk.models.query_result_set import QueryResultSet from shroomdk.rpc import RPC -from .models import Filter, SortBy - -API_BASE_URL = "https://rpc-prod.flompass.pizza" +API_BASE_URL = "https://rpc.flipsidecrypto.com" SDK_VERSION = "2.0.0" SDK_PACKAGE = "python" @@ -76,3 +74,6 @@ class ShroomDK(object): def get_sql_statement(self, sql_statement_id: str) -> SqlStatement: return self.query_integration.get_sql_statement(sql_statement_id) + + def cancel_query_run(self, query_run_id: str) -> QueryRun: + return self.query_integration.cancel_query_run(query_run_id)