Merge branch 'main' into sea-migration

Author: varun-edachali-dbx

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/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
@@ -163,6 +167,9 @@ def remaining_rows(self):
         self.cur_row_index += slice.num_rows
         return slice
 
+    def close(self):
+        return
+
 
 class ArrowQueue(ResultSetQueue):
     def __init__(
@@ -200,6 +207,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__(
@@ -352,6 +362,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]