Cursor.close() leaks server-side handle for async commands that were never fetched (#791)

Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>

Author: peco-engineer-bot[bot]

src/databricks/sql/client.py

@@ -1757,9 +1757,20 @@ 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()
+        elif self.active_command_id is not None:
+            # Async submission that never had get_async_execution_result called:
+            # the result-set close path didn't fire, so the server-side handle
+            # would leak until the session closes. Issue an explicit
+            # close_command to free it. Safe across all backends (Thrift sends
+            # TCloseOperation, SEA sends a close/cancel, kernel no-ops on
+            # unknown ids).
+            try:
+                self.backend.close_command(self.active_command_id)
+            except Exception as exc:
+                logger.warning("close_command on cursor close failed: %s", exc)
+        self.active_command_id = None
 
     @property
     def query_id(self) -> Optional[str]:

tests/e2e/test_driver.py

@@ -342,6 +342,59 @@ def test_execute_async__large_result(self, extra_params):
 
             assert len(result) == x_dimension * y_dimension
 
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+        ],
+    )
+    def test_close_frees_async_handle_without_fetch(self, extra_params):
+        """Regression test for #791.
+
+        Cursor.close() must free the server-side statement handle for an
+        async-submitted command even when get_async_execution_result() was
+        never called. Before the fix, close() only nulled active_command_id
+        and (for the never-fetched case) never issued close_command, so the
+        handle leaked until the session closed.
+        """
+        long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'"
+        with self.cursor(extra_params) as cursor:
+            cursor.execute_async(long_running_query)
+
+            command_id = cursor.active_command_id
+            backend = cursor.backend
+            assert command_id is not None
+
+            # The handle is live server-side right after async submission.
+            state = backend.get_query_state(command_id)
+            assert state in (
+                CommandState.PENDING,
+                CommandState.RUNNING,
+                CommandState.SUCCEEDED,
+            )
+
+            # Close without ever fetching the result.
+            cursor.close()
+
+            # The server must no longer report the handle as live. Polling the
+            # same command id should either report a terminal CLOSED/CANCELLED
+            # state, or raise signalling the command was closed / no longer
+            # exists server side (DatabaseError "closed server side" or a
+            # RESOURCE_DOES_NOT_EXIST RequestError). Any of these proves the
+            # server-side handle was freed by close().
+            try:
+                post_close_state = backend.get_query_state(command_id)
+                assert post_close_state in (
+                    CommandState.CLOSED,
+                    CommandState.CANCELLED,
+                ), f"Handle still live after close(): {post_close_state}"
+            except (DatabaseError, RequestError) as e:
+                assert (
+                    "closed server side" in str(e)
+                    or "does not exist" in str(e)
+                    or "RESOURCE_DOES_NOT_EXIST" in str(e)
+                )
+
 
 # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core
 # tests

tests/unit/test_client.py

@@ -330,6 +330,40 @@ def test_context_manager_closes_cursor(self):
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
+    def test_close_frees_async_command_without_result_set(self):
+        """Regression test for #791.
+
+        When a command was submitted asynchronously and never fetched,
+        active_result_set is None but active_command_id is set. close() must
+        still issue backend.close_command() to free the server-side handle.
+        """
+        mock_backend = Mock()
+        cursor = client.Cursor(Mock(), mock_backend)
+        mock_command_id = Mock()
+        cursor.active_command_id = mock_command_id
+        cursor.active_result_set = None
+
+        cursor.close()
+
+        mock_backend.close_command.assert_called_once_with(mock_command_id)
+        self.assertIsNone(cursor.active_command_id)
+        self.assertFalse(cursor.open)
+
+    def test_close_swallows_close_command_errors(self):
+        """close() must not raise if backend.close_command fails for an
+        unfetched async command (best-effort cleanup)."""
+        mock_backend = Mock()
+        mock_backend.close_command.side_effect = Exception("boom")
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.active_command_id = Mock()
+        cursor.active_result_set = None
+
+        # Should not raise.
+        cursor.close()
+
+        self.assertIsNone(cursor.active_command_id)
+        self.assertFalse(cursor.open)
+
     def dict_product(self, dicts):
         """
         Generate cartesion product of values in input dictionary, outputting a dictionary