Fix cancel query

python/requirements.txt

@@ -1,3 +1,3 @@
 pydantic==1.9.1
 requests==2.28.1
-urllib3==1.26.11
+urllib3==1.26.15

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,
 )

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)

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,

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]

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

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)