From 54f7e6121468f042ee513ee2cf3c1f8dee6ba462 Mon Sep 17 00:00:00 2001
From: Nanook <nanookclaw@users.noreply.github.com>
Date: Wed, 20 May 2026 00:12:33 +0000
Subject: [PATCH] fix: close unfetched async commands

Signed-off-by: Nanook <nanookclaw@users.noreply.github.com>
---
 src/databricks/sql/client.py | 10 +++++++---
 tests/unit/test_client.py    | 27 +++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index 7fc815cd8..aaff15e8e 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1713,9 +1713,13 @@ def cancel(self) -> None:
     def close(self) -> None:
         """Close cursor"""
         self.open = False
-        self.active_command_id = None
-        if self.active_result_set:
-            self._close_and_clear_active_result_set()
+        try:
+            if self.active_result_set:
+                self._close_and_clear_active_result_set()
+            elif self.active_command_id is not None:
+                self.backend.close_command(self.active_command_id)
+        finally:
+            self.active_command_id = None
 
     @property
     def query_id(self) -> Optional[str]:
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 4a8cb0b68..26f10eb7f 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -280,6 +280,33 @@ def test_context_manager_closes_cursor(self):
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
+    def test_close_closes_active_command_without_result_set(self):
+        mock_backend = Mock()
+        mock_command_id = Mock(spec=CommandId)
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.active_command_id = mock_command_id
+
+        cursor.close()
+
+        mock_backend.close_command.assert_called_once_with(mock_command_id)
+        self.assertIsNone(cursor.active_command_id)
+        self.assertIsNone(cursor.active_result_set)
+
+    def test_close_delegates_to_active_result_set_without_double_closing_command(self):
+        mock_backend = Mock()
+        mock_result_set = Mock()
+        mock_command_id = Mock(spec=CommandId)
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.active_result_set = mock_result_set
+        cursor.active_command_id = mock_command_id
+
+        cursor.close()
+
+        mock_result_set.close.assert_called_once_with()
+        mock_backend.close_command.assert_not_called()
+        self.assertIsNone(cursor.active_command_id)
+        self.assertIsNone(cursor.active_result_set)
+
     def dict_product(self, dicts):
         """
         Generate cartesion product of values in input dictionary, outputting a dictionary