Merge branch 'sea-migration' into chunk-latency

Author: saishreeeee

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

@@ -69,3 +69,6 @@ def remaining_rows(self) -> List[List[str]]:
         slice = self.data_array[self.cur_row_index :]
         self.cur_row_index += len(slice)
         return slice
+
+    def close(self):
+        return

src/databricks/sql/result_set.py

@@ -170,6 +170,7 @@ def close(self) -> None:
         been closed on the server for some other reason, issue a request to the server to close it.
         """
         try:
+            self.results.close()
             if (
                 self.status != CommandState.CLOSED
                 and not self.has_been_closed_server_side

src/databricks/sql/utils.py

@@ -51,6 +51,10 @@ def next_n_rows(self, num_rows: int):
     def remaining_rows(self):
         pass
 
+    @abstractmethod
+    def close(self):
+        pass
+
 
 class ThriftResultSetQueueFactory(ABC):
     @staticmethod
@@ -171,6 +175,9 @@ def remaining_rows(self):
         self.cur_row_index += slice.num_rows
         return slice
 
+    def close(self):
+        return
+
 
 class ArrowQueue(ResultSetQueue):
     def __init__(
@@ -208,6 +215,9 @@ def remaining_rows(self) -> "pyarrow.Table":
         self.cur_row_index += slice.num_rows
         return slice
 
+    def close(self):
+        return
+
 
 class CloudFetchQueue(ResultSetQueue):
     def __init__(
@@ -372,6 +382,9 @@ def _create_empty_table(self) -> "pyarrow.Table":
         # Create a 0-row table with just the schema bytes
         return create_arrow_table_from_arrow_file(self.schema_bytes, self.description)
 
+    def close(self):
+        self.download_manager._shutdown_manager()
+
 
 def _bound(min_x, max_x, x):
     """Bound x by [min_x, max_x]

tests/unit/test_client.py

@@ -199,6 +199,8 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             session_id_hex=Mock(),
             statement_type=Mock(),
         )
+        result_set.results = mock_results
+
         # Setup session mock on the mock_connection
         mock_session = Mock()
         mock_session.open = False
@@ -208,12 +210,14 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
 
         self.assertFalse(mock_backend.close_command.called)
         self.assertTrue(result_set.has_been_closed_server_side)
+        mock_results.close.assert_called_once()
 
     def test_closing_result_set_hard_closes_commands(self):
         mock_results_response = Mock()
         mock_results_response.has_been_closed_server_side = False
         mock_connection = Mock()
         mock_thrift_backend = Mock()
+        mock_results = Mock()
         # Setup session mock on the mock_connection
         mock_session = Mock()
         mock_session.open = True
@@ -228,6 +232,7 @@ def test_closing_result_set_hard_closes_commands(self):
         mock_thrift_backend.close_command.assert_called_once_with(
             mock_results_response.command_id
         )
+        mock_results.close.assert_called_once()
 
     def test_executing_multiple_commands_uses_the_most_recent_command(self):
         mock_result_sets = [Mock(), Mock()]