Missed adding some files in previous commit

Author: jprakash-db

src/databricks/sql/client.py

@@ -733,7 +733,7 @@ def execute(
         self,
         operation: str,
         parameters: Optional[TParameterCollection] = None,
-        perform_async = True
+        perform_async = False
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -814,10 +814,37 @@ def execute(
 
         return self
 
-    def executeAsync(self,
+    def execute_async(self,
                      operation: str,
                      parameters: Optional[TParameterCollection] = None,):
-        return execute(operation, parameters, True)
+        return self.execute(operation, parameters, True)
+
+    def get_query_status(self):
+        self._check_not_closed()
+        return self.thrift_backend.get_query_status(self.active_op_handle)
+
+    def get_execution_result(self):
+        self._check_not_closed()
+
+        operation_state = self.get_query_status()
+        if operation_state.statusCode == ttypes.TStatusCode.SUCCESS_STATUS or operation_state.statusCode == ttypes.TStatusCode.SUCCESS_WITH_INFO_STATUS:
+            execute_response=self.thrift_backend.get_execution_result(self.active_op_handle)
+            self.active_result_set = ResultSet(
+                self.connection,
+                execute_response,
+                self.thrift_backend,
+                self.buffer_size_bytes,
+                self.arraysize,
+            )
+
+            if execute_response.is_staging_operation:
+                self._handle_staging_operation(
+                    staging_allowed_local_path=self.thrift_backend.staging_allowed_local_path
+                )
+
+            return self
+        else:
+            raise Error(f"get_execution_result failed with status code {operation_state.statusCode}")
 
     def executemany(self, operation, seq_of_parameters):
         """
@@ -1126,7 +1153,7 @@ def __init__(
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._next_row_index = 0
 
-        if execute_response.arrow_queue:
+        if execute_response.arrow_queue or True:
             # In this case the server has taken the fast path and returned an initial batch of
             # results
             self.results = execute_response.arrow_queue

src/databricks/sql/constants.py

@@ -0,0 +1,12 @@
+from databricks.sql.thrift_api.TCLIService import ttypes
+
+class QueryExecutionStatus:
+	INITIALIZED_STATE=ttypes.TOperationState.INITIALIZED_STATE
+	RUNNING_STATE = ttypes.TOperationState.RUNNING_STATE
+	FINISHED_STATE = ttypes.TOperationState.FINISHED_STATE
+	CANCELED_STATE = ttypes.TOperationState.CANCELED_STATE
+	CLOSED_STATE = ttypes.TOperationState.CLOSED_STATE
+	ERROR_STATE = ttypes.TOperationState.ERROR_STATE
+	UKNOWN_STATE = ttypes.TOperationState.UKNOWN_STATE
+	PENDING_STATE = ttypes.TOperationState.PENDING_STATE
+	TIMEDOUT_STATE = ttypes.TOperationState.TIMEDOUT_STATE

src/databricks/sql/thrift_backend.py

@@ -769,6 +769,66 @@ def _results_message_to_execute_response(self, resp, operation_state):
             arrow_schema_bytes=schema_bytes,
         )
 
+    def get_execution_result(self, op_handle):
+
+        assert op_handle is not None
+
+        req = ttypes.TFetchResultsReq(
+            operationHandle=ttypes.TOperationHandle(
+                op_handle.operationId,
+                op_handle.operationType,
+                False,
+                op_handle.modifiedRowCount,
+            ),
+            maxRows=max_rows,
+            maxBytes=max_bytes,
+            orientation=ttypes.TFetchOrientation.FETCH_NEXT,
+            includeResultSetMetadata=True,
+        )
+
+        resp = self.make_request(self._client.FetchResults, req)
+
+        t_result_set_metadata_resp = resp.resultSetMetaData
+
+        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
+        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
+        has_more_rows = resp.hasMoreRows
+        description = self._hive_schema_to_description(
+            t_result_set_metadata_resp.schema
+        )
+
+        if pyarrow:
+            schema_bytes = (
+                t_result_set_metadata_resp.arrowSchema
+                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
+                .serialize()
+                .to_pybytes()
+            )
+        else:
+            schema_bytes = None
+
+        queue = ResultSetQueueFactory.build_queue(
+            row_set_type=resp.resultSetMetadata.resultFormat,
+            t_row_set=resp.results,
+            arrow_schema_bytes=schema_bytes,
+            max_download_threads=self.max_download_threads,
+            lz4_compressed=lz4_compressed,
+            description=description,
+            ssl_options=self._ssl_options,
+        )
+
+        return ExecuteResponse(
+            arrow_queue=queue,
+            status=resp.status,
+            has_been_closed_server_side=has_been_closed_server_side,
+            has_more_rows=has_more_rows,
+            lz4_compressed=lz4_compressed,
+            is_staging_operation=is_staging_operation,
+            command_handle=resp.operationHandle,
+            description=description,
+            arrow_schema_bytes=schema_bytes,
+        )
+
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
         if initial_operation_status_resp:
             self._check_command_not_in_error_or_closed_state(
@@ -787,6 +847,12 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
             self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)
         return operation_state
 
+    def get_query_status(self, op_handle):
+        poll_resp = self._poll_for_status(op_handle)
+        operation_state = poll_resp.status
+        self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)
+        return operation_state
+
     @staticmethod
     def _check_direct_results_for_error(t_spark_direct_results):
         if t_spark_direct_results:
@@ -848,7 +914,10 @@ def execute_command(
         )
         resp = self.make_request(self._client.ExecuteStatement, req)
 
-        return self._handle_execute_response(resp, cursor, perform_async)
+        if perform_async:
+            return self._handle_execute_response_async(resp, cursor)
+        else:
+            return self._handle_execute_response(resp, cursor)
 
     def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
         assert session_handle is not None
@@ -936,19 +1005,34 @@ def get_columns(
         resp = self.make_request(self._client.GetColumns, req)
         return self._handle_execute_response(resp, cursor)
 
-    def _handle_execute_response(self, resp, cursor, perform_async=False):
+    def _handle_execute_response(self, resp, cursor):
         cursor.active_op_handle = resp.operationHandle
         self._check_direct_results_for_error(resp.directResults)
 
-        if perform_async:
-            final_operation_state=ttypes.TStatusCode.STILL_EXECUTING_STATUS
-        else:
-            final_operation_state=self._wait_until_command_done(
+        final_operation_state = self._wait_until_command_done(
             resp.operationHandle,
-            resp.directResults and resp.directResults.operationStatus)
+            resp.directResults and resp.directResults.operationStatus,
+            )
 
         return self._results_message_to_execute_response(resp, final_operation_state)
 
+    def _handle_execute_response_async(self, resp, cursor):
+        cursor.active_op_handle = resp.operationHandle
+        self._check_direct_results_for_error(resp.directResults)
+        operation_status = resp.status.statusCode
+
+        return ExecuteResponse(
+            arrow_queue=None,
+            status=operation_status,
+            has_been_closed_server_side=None,
+            has_more_rows=None,
+            lz4_compressed=None,
+            is_staging_operation=None,
+            command_handle=resp.operationHandle,
+            description=None,
+            arrow_schema_bytes=None,
+        )
+
     def fetch_results(
         self,
         op_handle,