[PECO-1263] Implement a .returned_as_direct_result property for AsyncExecution status

src/databricks/sql/ae.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Optional, Union, TYPE_CHECKING
 from databricks.sql.exc import RequestError
-from databricks.sql.results import ResultSet
+from databricks.sql.results import ResultSet, execute_response_contains_direct_results
 
 from datetime import datetime
 
@@ -81,6 +81,7 @@ class AsyncExecution:
     ]
     _last_sync_timestamp: Optional[datetime] = None
     _result_set: Optional["ResultSet"] = None
+    _is_available: bool = True
 
     def __init__(
         self,
@@ -101,6 +102,8 @@ def __init__(
 
         if execute_statement_response:
             self._execute_statement_response = execute_statement_response
+            if execute_response_contains_direct_results(execute_statement_response):
+                self._is_available = False
         else:
             self._execute_statement_response = FakeExecuteStatementResponse(
                 directResults=False, operationHandle=self.t_operation_handle
@@ -225,6 +228,15 @@ def last_sync_timestamp(self) -> Optional[datetime]:
         """The timestamp of the last time self.status was synced with the server"""
         return self._last_sync_timestamp
 
+    @property
+    def is_available(self) -> bool:
+        """Indicates whether the result of this query can be fetched from a separate thread.
+
+        Only returns False if the query returned its results directly when `execute_async` was called.
+        """
+
+        return self._is_available
+
     @classmethod
     def from_thrift_response(
         cls,

src/databricks/sql/results.py

@@ -9,6 +9,7 @@
 from databricks.sql.exc import (
     CursorAlreadyClosedError,
 )
+from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import Row
 from databricks.sql.utils import ExecuteResponse
 
@@ -223,3 +224,14 @@ def map_col_type(type_):
             (column.name, map_col_type(column.datatype), None, None, None, None, None)
             for column in table_schema_message.columns
         ]
+
+
+def execute_response_contains_direct_results(
+    execute_response: ttypes.TExecuteStatementResp,
+) -> bool:
+    """
+    Returns True if the thrift TExecuteStatementResp contains metadata
+    This indicates the statement has finished executing at the server.
+    """
+
+    return bool(execute_response.directResults.resultSetMetadata)

tests/e2e/test_execute_async.py

@@ -166,6 +166,9 @@ def test_get_async_execution_no_results_when_direct_results_were_sent(self):
 
         with self.connection() as conn:
             ae = conn.execute_async(DIRECT_RESULTS_QUERY, {"param": 1})
+            assert (
+                not ae.is_available
+            ), "Queries that return direct results should not be available"
             query_id, query_secret = ae.serialize().split(":")
             ae.get_result()
 
@@ -193,9 +196,13 @@ def test_get_async_execution_twice(self):
         """
         with self.connection() as conn_1, self.connection() as conn_2:
             ae_1 = conn_1.execute_async(LONG_ISH_QUERY)
+            assert (
+                ae_1.is_available
+            ), "A long query does not return direct results so is_available should be True"
 
             query_id, query_secret = ae_1.serialize().split(":")
             ae_2 = conn_2.get_async_execution(query_id, query_secret)
+            assert ae_2.is_available
 
             while ae_1.is_running:
                 time.sleep(1)