remove service specific state from ExecuteResponse

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/client.py

@@ -349,10 +349,8 @@ def _results_message_to_execute_response(
             command_id=CommandId.from_sea_statement_id(response.statement_id),
             status=response.status.state,
             description=description,
-            has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=response.manifest.is_volume_operation,
-            arrow_schema_bytes=None,
             result_format=response.manifest.format,
         )
 

src/databricks/sql/backend/sea/result_set.py

@@ -15,10 +15,10 @@
 
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
-from databricks.sql.exc import ProgrammingError
+from databricks.sql.exc import CursorAlreadyClosedError, ProgrammingError, RequestError
 from databricks.sql.types import Row
 from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
-from databricks.sql.backend.types import ExecuteResponse
+from databricks.sql.backend.types import CommandState, ExecuteResponse
 from databricks.sql.result_set import ResultSet
 
 logger = logging.getLogger(__name__)
@@ -61,11 +61,9 @@ def __init__(
             buffer_size_bytes=buffer_size_bytes,
             command_id=execute_response.command_id,
             status=execute_response.status,
-            has_been_closed_server_side=execute_response.has_been_closed_server_side,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
             lz4_compressed=execute_response.lz4_compressed,
-            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
         # Assert that the backend is of the correct type
@@ -274,3 +272,21 @@ def fetchall(self) -> List[Row]:
             return self._create_json_table(self.fetchall_json())
         else:
             raise NotImplementedError("fetchall only supported for JSON data")
+
+    def close(self) -> None:
+        """
+        Close the result set.
+
+        If the connection has not been closed, and the result set has not already
+        been closed on the server for some other reason, issue a request to the server to close it.
+        """
+        try:
+            if self.results is not None:
+                self.results.close()
+            if self.status != CommandState.CLOSED and self.connection.open:
+                self.backend.close_command(self.command_id)
+        except RequestError as e:
+            if isinstance(e.args[1], CursorAlreadyClosedError):
+                logger.info("Operation was canceled by a prior request")
+        finally:
+            self.status = CommandState.CLOSED

src/databricks/sql/backend/sea/utils/filters.py

@@ -60,9 +60,7 @@ def _filter_sea_result_set(
             command_id=command_id,
             status=result_set.status,
             description=result_set.description,
-            has_been_closed_server_side=result_set.has_been_closed_server_side,
             lz4_compressed=result_set.lz4_compressed,
-            arrow_schema_bytes=result_set._arrow_schema_bytes,
             is_staging_operation=False,
         )
 

src/databricks/sql/backend/thrift_backend.py

@@ -821,14 +821,17 @@ def _results_message_to_execute_response(self, resp, operation_state):
             command_id=command_id,
             status=status,
             description=description,
-            has_been_closed_server_side=has_been_closed_server_side,
             lz4_compressed=lz4_compressed,
             is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
-            arrow_schema_bytes=schema_bytes,
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
-        return execute_response, is_direct_results
+        return (
+            execute_response,
+            is_direct_results,
+            has_been_closed_server_side,
+            schema_bytes,
+        )
 
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
@@ -881,10 +884,8 @@ def get_execution_result(
             command_id=command_id,
             status=status,
             description=description,
-            has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
-            arrow_schema_bytes=schema_bytes,
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
@@ -898,6 +899,8 @@ def get_execution_result(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+            arrow_schema_bytes=schema_bytes,
+            has_been_closed_server_side=False,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -1016,9 +1019,12 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results = self._handle_execute_response(
-                resp, cursor
-            )
+            (
+                execute_response,
+                is_direct_results,
+                has_been_closed_server_side,
+                schema_bytes,
+            ) = self._handle_execute_response(resp, cursor)
 
             t_row_set = None
             if resp.directResults and resp.directResults.resultSet:
@@ -1034,6 +1040,8 @@ def execute_command(
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
                 is_direct_results=is_direct_results,
+                has_been_closed_server_side=has_been_closed_server_side,
+                arrow_schema_bytes=schema_bytes,
             )
 
     def get_catalogs(
@@ -1055,9 +1063,12 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        (
+            execute_response,
+            is_direct_results,
+            has_been_closed_server_side,
+            schema_bytes,
+        ) = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1073,6 +1084,8 @@ def get_catalogs(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+            has_been_closed_server_side=has_been_closed_server_side,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def get_schemas(
@@ -1100,9 +1113,12 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        (
+            execute_response,
+            is_direct_results,
+            has_been_closed_server_side,
+            schema_bytes,
+        ) = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1118,6 +1134,8 @@ def get_schemas(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+            has_been_closed_server_side=has_been_closed_server_side,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def get_tables(
@@ -1149,9 +1167,12 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        (
+            execute_response,
+            is_direct_results,
+            has_been_closed_server_side,
+            schema_bytes,
+        ) = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1167,6 +1188,8 @@ def get_tables(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+            has_been_closed_server_side=has_been_closed_server_side,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def get_columns(
@@ -1198,9 +1221,12 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        (
+            execute_response,
+            is_direct_results,
+            has_been_closed_server_side,
+            schema_bytes,
+        ) = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1216,6 +1242,8 @@ def get_columns(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
+            has_been_closed_server_side=has_been_closed_server_side,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/backend/types.py

@@ -419,8 +419,6 @@ class ExecuteResponse:
     command_id: CommandId
     status: CommandState
     description: List[Tuple]
-    has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
     is_staging_operation: bool = False
-    arrow_schema_bytes: Optional[bytes] = None
     result_format: Optional[Any] = None

src/databricks/sql/result_set.py

@@ -41,12 +41,10 @@ def __init__(
         buffer_size_bytes: int,
         command_id: CommandId,
         status: CommandState,
-        has_been_closed_server_side: bool = False,
         is_direct_results: bool = False,
         description: List[Tuple] = [],
         is_staging_operation: bool = False,
         lz4_compressed: bool = False,
-        arrow_schema_bytes: Optional[bytes] = None,
     ):
         """
         A ResultSet manages the results of a single command.
@@ -57,7 +55,6 @@ def __init__(
             :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch
             :param command_id: The command ID
             :param status: The command status
-            :param has_been_closed_server_side: Whether the command has been closed on the server
             :param is_direct_results: Whether the command has more rows
             :param description: column description of the results
             :param is_staging_operation: Whether the command is a staging operation
@@ -71,12 +68,10 @@ def __init__(
         self.description: List[Tuple] = description
         self.command_id: CommandId = command_id
         self.status: CommandState = status
-        self.has_been_closed_server_side: bool = has_been_closed_server_side
         self.is_direct_results: bool = is_direct_results
         self.results: Optional[ResultSetQueue] = None
         self._is_staging_operation: bool = is_staging_operation
         self.lz4_compressed: bool = lz4_compressed
-        self._arrow_schema_bytes: Optional[bytes] = arrow_schema_bytes
 
     def __iter__(self):
         while True:
@@ -158,28 +153,12 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all remaining rows as an Arrow table."""
         pass
 
+    @abstractmethod
     def close(self) -> None:
         """
         Close the result set.
-
-        If the connection has not been closed, and the result set has not already
-        been closed on the server for some other reason, issue a request to the server to close it.
         """
-        try:
-            if self.results:
-                self.results.close()
-            if (
-                self.status != CommandState.CLOSED
-                and not self.has_been_closed_server_side
-                and self.connection.open
-            ):
-                self.backend.close_command(self.command_id)
-        except RequestError as e:
-            if isinstance(e.args[1], CursorAlreadyClosedError):
-                logger.info("Operation was canceled by a prior request")
-        finally:
-            self.has_been_closed_server_side = True
-            self.status = CommandState.CLOSED
+        pass
 
 
 class ThriftResultSet(ResultSet):
@@ -196,6 +175,8 @@ def __init__(
         max_download_threads: int = 10,
         ssl_options=None,
         is_direct_results: bool = True,
+        has_been_closed_server_side: bool = False,
+        arrow_schema_bytes: Optional[bytes] = None,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -210,11 +191,15 @@ def __init__(
             :param max_download_threads: Maximum number of download threads for cloud fetch
             :param ssl_options: SSL options for cloud fetch
             :param is_direct_results: Whether there are more rows to fetch
+            :param has_been_closed_server_side: Whether the command has been closed on the server
+            :param arrow_schema_bytes: The schema of the result set
         """
 
         # Initialize ThriftResultSet-specific attributes
         self._use_cloud_fetch = use_cloud_fetch
         self.is_direct_results = is_direct_results
+        self.has_been_closed_server_side = has_been_closed_server_side
+        self._arrow_schema_bytes = arrow_schema_bytes
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -225,7 +210,7 @@ def __init__(
             results_queue = ThriftResultSetQueueFactory.build_queue(
                 row_set_type=execute_response.result_format,
                 t_row_set=t_row_set,
-                arrow_schema_bytes=execute_response.arrow_schema_bytes or b"",
+                arrow_schema_bytes=self._arrow_schema_bytes or b"",
                 max_download_threads=max_download_threads,
                 lz4_compressed=execute_response.lz4_compressed,
                 description=execute_response.description,
@@ -239,12 +224,10 @@ def __init__(
             buffer_size_bytes=buffer_size_bytes,
             command_id=execute_response.command_id,
             status=execute_response.status,
-            has_been_closed_server_side=execute_response.has_been_closed_server_side,
             is_direct_results=is_direct_results,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
             lz4_compressed=execute_response.lz4_compressed,
-            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
         # Assert that the backend is of the correct type
@@ -460,3 +443,29 @@ 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 close(self) -> None:
+        """
+        Close the result set.
+
+        If the connection has not been closed, and the result set has not already
+        been closed on the server for some other reason, issue a request to the server to close it.
+        """
+        try:
+            if self.results:
+                self.results.close()
+            print(f"has_been_closed_server_side: {self.has_been_closed_server_side}")
+            print(f"status: {self.status}")
+            print(f"connection.open: {self.connection.open}")
+            if (
+                self.status != CommandState.CLOSED
+                and not self.has_been_closed_server_side
+                and self.connection.open
+            ):
+                self.backend.close_command(self.command_id)
+        except RequestError as e:
+            if isinstance(e.args[1], CursorAlreadyClosedError):
+                logger.info("Operation was canceled by a prior request")
+        finally:
+            self.has_been_closed_server_side = True
+            self.status = CommandState.CLOSED