[PECO-1263] Separate get_status and poll_for_status methods (#313)

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

Author: Jesse

src/databricks/sql/ae.py

@@ -15,6 +15,10 @@
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 
+class AsyncExecutionException(Exception):
+    pass
+
+
 @dataclass
 class FakeCursor:
     active_op_handle: Optional[ttypes.TOperationHandle]
@@ -50,8 +54,7 @@ def _toperationstate_to_ae_status(
 
 class AsyncExecution:
     """
-    A class that represents an async execution of a query. Exposes just two methods:
-    get_result_or_status and cancel
+    A class that represents an async execution of a query.
 
     AsyncExecutions are effectively connectionless. But because thrift_backend is entangled
     with client.py, the AsyncExecution needs access to both a Connection and a ThriftBackend
@@ -83,18 +86,20 @@ def __init__(
     status: AsyncExecutionStatus
     query_id: UUID
 
-    def get_result_or_status(self) -> Union["ResultSet", AsyncExecutionStatus]:
-        """Get the result of the async execution. If execution has not completed, return False."""
+    def get_result(self) -> "ResultSet":
+        """Get a result set for this async execution
+
+        Raises an exception if the query is still running or has been canceled.
+        """
 
         if self.status == AsyncExecutionStatus.CANCELED:
-            return self.status
+            raise AsyncExecutionException("Query was canceled: %s" % self.query_id)
+        if self.is_running:
+            raise AsyncExecutionException("Query is still running: %s" % self.query_id)
         if self.status == AsyncExecutionStatus.FINISHED:
             self._thrift_fetch_result()
         if self.status == AsyncExecutionStatus.FETCHED:
             return self._result_set
-        else:
-            self._thrift_get_operation_status()
-            return self.status
 
     def cancel(self) -> None:
         """Cancel the query"""
@@ -106,11 +111,13 @@ def _thrift_cancel_operation(self) -> None:
         _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"""
+    def poll_for_status(self) -> None:
+        """Check the thrift server for the status of this operation and set self.status
+
+        This will result in an error if the operation has been canceled or aborted at the server"""
 
         _output = self._thrift_backend._poll_for_status(self.t_operation_handle)
-        self.status = _toperationstate_to_ae_status(_output)
+        self.status = _toperationstate_to_ae_status(_output.operationState)
 
     def _thrift_fetch_result(self) -> None:
         """Execute TFetchResultReq and store the result"""

src/databricks/sql/thrift_backend.py

@@ -775,10 +775,7 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
                 num_rows,
             ) = convert_column_based_set_to_arrow_table(t_row_set.columns, description)
         elif t_row_set.arrowBatches is not None:
-            (
-                arrow_table,
-                num_rows,
-            ) = convert_arrow_based_set_to_arrow_table(
+            (arrow_table, num_rows,) = convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, schema_bytes
             )
         else:

tests/e2e/test_execute_async.py

@@ -1,50 +1,62 @@
 from tests.e2e.test_driver import PySQLPytestTestCase
 
-from databricks.sql.ae import AsyncExecutionStatus as AsyncExecutionStatus
-
+from databricks.sql.ae import (
+    AsyncExecutionStatus,
+    AsyncExecutionException,
+    AsyncExecution,
+)
+import pytest
 import time
 
-LONG_RUNNING_QUERY =  """
+LONG_RUNNING_QUERY = """
 SELECT SUM(A.id - B.id)
 FROM range(1000000000) A CROSS JOIN range(100000000) B 
 GROUP BY (A.id - B.id)
 """
 
+
 class TestExecuteAsync(PySQLPytestTestCase):
+    @pytest.fixture
+    def long_running_ae(self, scope="function") -> AsyncExecution:
+        """Start a long-running query so we can make assertions about it."""
+        with self.connection() as conn:
+            ae = conn.execute_async(LONG_RUNNING_QUERY)
+            yield ae
+
+            # cancellation is idempotent
+            ae.cancel()
 
     def test_basic_api(self):
-        """This is a WIP test of the basic API defined in PECO-1263
-        """
+        """This is a WIP test of the basic API defined in PECO-1263"""
         # This is taken directly from the design doc
 
         with self.connection() as conn:
             ae = conn.execute_async("select :param `col`", {"param": 1})
             while ae.is_running:
-                ae.get_result_or_status()
+                ae.poll_for_status()
                 time.sleep(1)
-            
-            result = ae.get_result_or_status().fetchone()
 
-        assert result.col == 1
+            result = ae.get_result().fetchone()
 
-    def test_cancel_running_query(self):
-        """Start a long-running query and cancel it
-        """
+        assert result.col == 1
 
-        with self.connection() as conn:
-            ae = conn.execute_async(LONG_RUNNING_QUERY)
-            time.sleep(2)
-            ae.cancel()
+    def test_cancel_running_query(self, long_running_ae: AsyncExecution):
+        long_running_ae.cancel()
+        assert long_running_ae.status == AsyncExecutionStatus.CANCELED
 
-            status = ae.get_result_or_status()
+    def test_cant_get_results_while_running(self, long_running_ae: AsyncExecution):
+        with pytest.raises(AsyncExecutionException, match="Query is still running"):
+            long_running_ae.get_result()
 
-        assert ae.status == AsyncExecutionStatus.CANCELED
+    def test_cant_get_results_after_cancel(self, long_running_ae: AsyncExecution):
+        long_running_ae.cancel()
+        with pytest.raises(AsyncExecutionException, match="Query was canceled"):
+            long_running_ae.get_result()
 
-    
 
     def test_staging_operation(self):
         """We need to test what happens with a staging operation since this query won't have a result set
         that user needs. It could be sufficient to notify users that they shouldn't use this API for staging/volumes
         queries...
         """
-        assert False
+        assert False