[PECO-1263] Extract asyncexecution from thrift backend (#312)

Teach AsyncExecution to cope with results from ThriftBackend

This way, thrift_backend.py doesn't need to even know about AsyncExecution

This de-complicates the dependency between these two modules.

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

src/databricks/sql/ae.py

@@ -2,6 +2,7 @@
 from typing import Optional, Union, TYPE_CHECKING
 from databricks.sql.results import ResultSet
 
+from dataclasses import dataclass
 
 if TYPE_CHECKING:
     from databricks.sql.thrift_backend import ThriftBackend
@@ -14,6 +15,11 @@
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 
+@dataclass
+class FakeCursor:
+    active_op_handle: Optional[ttypes.TOperationHandle]
+
+
 class AsyncExecutionStatus(Enum):
     """An enum that represents the status of an async execution"""
 
@@ -22,6 +28,8 @@ class AsyncExecutionStatus(Enum):
     FINISHED = 2
     CANCELED = 3
     FETCHED = 4
+
+    # todo: when is this ever evaluated?
     ABORTED = 5
 
 
@@ -44,21 +52,29 @@ class AsyncExecution:
     """
     A class that represents an async execution of a query. Exposes just two methods:
     get_result_or_status and cancel
+
+    AsyncExecutions are effectively connectionless. But because thrift_backend is entangled
+    with client.py, the AsyncExecution needs access to both a Connection and a ThriftBackend
+
+    This will need to be refactored for cleanliness in the future.
     """
 
     _connection: "Connection"
+    _thrift_backend: "ThriftBackend"
     _result_set: Optional["ResultSet"]
     _execute_statement_response: Optional[ttypes.TExecuteStatementResp]
 
     def __init__(
         self,
+        thrift_backend: "ThriftBackend",
         connection: "Connection",
         query_id: UUID,
         query_secret: UUID,
         status: AsyncExecutionStatus,
         execute_statement_response: Optional[ttypes.TExecuteStatementResp] = None,
     ):
         self._connection = connection
+        self._thrift_backend = thrift_backend
         self._execute_statement_response = execute_statement_response
         self.query_id = query_id
         self.query_secret = query_secret
@@ -87,13 +103,13 @@ def cancel(self) -> None:
     def _thrift_cancel_operation(self) -> None:
         """Execute TCancelOperation"""
 
-        _output = self._connection.thrift_backend.async_cancel_command(self.t_operation_handle)
+        _output = self._thrift_backend.async_cancel_command(self.t_operation_handle)
         self.status = AsyncExecutionStatus.CANCELED
 
     def _thrift_get_operation_status(self) -> None:
         """Execute GetOperationStatusReq and map thrift execution status to DbsqlAsyncExecutionStatus"""
 
-        _output = self._connection.thrift_backend._poll_for_status(self.t_operation_handle)
+        _output = self._thrift_backend._poll_for_status(self.t_operation_handle)
         self.status = _toperationstate_to_ae_status(_output)
 
     def _thrift_fetch_result(self) -> None:
@@ -104,10 +120,10 @@ def _thrift_fetch_result(self) -> None:
         # support JSON and Thrift binary result formats in addition to arrow.
 
         # in the case of direct results this creates a second cursor...how can I avoid that?
-        with self._connection.cursor() as cursor:
-            er = self._connection.thrift_backend._handle_execute_response(
-                self._execute_statement_response, cursor
-            )
+
+        er = self._thrift_backend._handle_execute_response(
+            self._execute_statement_response, FakeCursor(None)
+        )
 
         self._result_set = ResultSet(
             connection=self._connection,
@@ -123,18 +139,37 @@ def is_running(self) -> bool:
             AsyncExecutionStatus.RUNNING,
             AsyncExecutionStatus.PENDING,
         ]
-    
+
     @property
     def t_operation_handle(self) -> ttypes.TOperationHandle:
-        """Return the current AsyncExecution as a Thrift TOperationHandle
-        """
+        """Return the current AsyncExecution as a Thrift TOperationHandle"""
 
         handle = ttypes.TOperationHandle(
             operationId=ttypes.THandleIdentifier(
                 guid=self.query_id.bytes, secret=self.query_secret.bytes
             ),
             operationType=ttypes.TOperationType.EXECUTE_STATEMENT,
-            hasResultSet=True
+            hasResultSet=True,
         )
 
         return handle
+
+    @classmethod
+    def from_thrift_response(
+        cls,
+        connection: "Connection",
+        thrift_backend: "ThriftBackend",
+        resp: ttypes.TExecuteStatementResp,
+    ) -> "AsyncExecution":
+        """This method is meant to be consumed by `client.py`"""
+
+        return cls(
+            connection=connection,
+            thrift_backend=thrift_backend,
+            query_id=UUID(bytes=resp.operationHandle.operationId.guid),
+            query_secret=UUID(bytes=resp.operationHandle.operationId.secret),
+            status=_toperationstate_to_ae_status(
+                resp.directResults.operationStatus.operationState
+            ),
+            execute_statement_response=resp,
+        )

src/databricks/sql/client.py

@@ -25,8 +25,8 @@
 
 from databricks.sql.results import ResultSet
 
-if TYPE_CHECKING:
-    from databricks.sql.ae import AsyncExecution, AsyncExecutionStatus
+
+from databricks.sql.ae import AsyncExecution, AsyncExecutionStatus
 
 logger = logging.getLogger(__name__)
 
@@ -386,7 +386,7 @@ def execute_async(
         )
 
         with self.cursor() as cursor:
-            ae: "AsyncExecution" = self.thrift_backend.async_execute_statement(
+            execute_statement_resp = self.thrift_backend.async_execute_statement(
                 statement=prepared_operation,
                 session_handle=self._session_handle,
                 max_rows=cursor.arraysize,
@@ -397,8 +397,11 @@ def execute_async(
                 parameters=prepared_params,
             )
 
-        # should we log this?
-        return ae
+        return AsyncExecution.from_thrift_response(
+            connection=self,
+            thrift_backend=self.thrift_backend,
+            resp=execute_statement_resp,
+        )
 
 
 class Cursor:

src/databricks/sql/thrift_backend.py

@@ -9,12 +9,10 @@
 from typing import List, Union, Callable, TYPE_CHECKING, Optional
 
 
-from databricks.sql.ae import _toperationstate_to_ae_status, AsyncExecution
-
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.parameters.native import TSparkParameter
-    from databricks.sql.ae import AsyncExecution
+
 
 import pyarrow
 import thrift.transport.THttpClient
@@ -430,7 +428,7 @@ def async_execute_statement(
         cursor: "Cursor",
         use_cloud_fetch: bool = True,
         parameters: Optional[List["TSparkParameter"]] = [],
-    ) -> "AsyncExecution":
+    ) -> ttypes.TExecuteStatementResp:
         """Send an ExecuteStatement command to the server, and return an AsyncExecution object.
 
         Args:
@@ -479,28 +477,7 @@ def async_execute_statement(
             self._client.ExecuteStatement, req
         )
 
-        query_id = guid = uuid.UUID(
-            hex=self.guid_to_hex_id(resp.operationHandle.operationId.guid)
-        )
-
-        secret = uuid.UUID(
-            hex=self.guid_to_hex_id(resp.operationHandle.operationId.secret)
-        )
-
-        # operationStatus -> TOperationstate
-        status = _toperationstate_to_ae_status(
-            resp.directResults.operationStatus.operationState
-        )
-
-        ae = AsyncExecution(
-            connection=cursor.connection,
-            query_id=query_id,
-            query_secret=secret,
-            status=status,
-            execute_statement_response=resp,
-        )
-
-        return ae
+        return resp
 
     # FUTURE: Consider moving to https://github.com/litl/backoff or
     # https://github.com/jd/tenacity for retry logic.
@@ -1187,9 +1164,7 @@ def cancel_command(self, active_op_handle):
         req = ttypes.TCancelOperationReq(active_op_handle)
         self.make_request(self._client.CancelOperation, req)
 
-    def async_cancel_command(
-        self, op_handle: ttypes.TOperationHandle
-    ) -> None:
+    def async_cancel_command(self, op_handle: ttypes.TOperationHandle) -> None:
         """Cancel a query using the thrift operation handle
 
         Args: