add close() for Queue, add ResultSet invocation

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -1505,6 +1505,7 @@ def close(self) -> None:
                 and self.connection.open
             ):
                 self.thrift_backend.close_command(self.command_id)
+                self.results.close()
         except RequestError as e:
             if isinstance(e.args[1], CursorAlreadyClosedError):
                 logger.info("Operation was canceled by a prior request")

src/databricks/sql/utils.py

@@ -46,6 +46,10 @@ def next_n_rows(self, num_rows: int):
     def remaining_rows(self):
         pass
 
+    @abstractmethod
+    def close(self):
+        pass
+
 
 class ResultSetQueueFactory(ABC):
     @staticmethod
@@ -157,6 +161,9 @@ def remaining_rows(self):
         self.cur_row_index += slice.num_rows
         return slice
 
+    def close(self):
+        return
+
 
 class ArrowQueue(ResultSetQueue):
     def __init__(
@@ -192,6 +199,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__(
@@ -341,6 +351,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()
+
 
 ExecuteResponse = namedtuple(
     "ExecuteResponse",