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

src/databricks/sql/client.py

@@ -16,6 +16,7 @@
 from databricks.sql import *
 from databricks.sql.exc import (
     OperationalError,
+    RequestError,
     SessionAlreadyClosedError,
     CursorAlreadyClosedError,
     InterfaceError,
@@ -1759,9 +1760,37 @@ 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 and self.connection.open:
+            # Async submission whose result was never fetched (no
+            # get_execution_result call), so the result-set close path never
+            # fired. Issue an explicit close_command to free the server-side
+            # statement handle instead of leaking it until session close.
+            # Gate on connection.open (mirroring ResultSet.close) so we don't
+            # attempt a network call on an already-torn-down session.
+            try:
+                self.backend.close_command(self.active_command_id)
+            except RequestError as e:
+                if len(e.args) > 1 and isinstance(e.args[1], CursorAlreadyClosedError):
+                    # Already-closed handle (e.g. a prior cancel() or concurrent
+                    # session teardown) is an expected, benign case — mirror
+                    # ResultSet.close and log at info, not warning.
+                    logger.info("Operation was canceled by a prior request")
+                else:
+                    logger.warning("close_command on cursor close failed: %s", e)
+            except Exception as exc:
+                # Unlike the RequestError branch above (an expected, often
+                # transient network failure), reaching here means something we
+                # did not anticipate — surface the full traceback so it is
+                # diagnosable rather than indistinguishable from a benign
+                # network blip.
+                logger.warning(
+                    "close_command on cursor close failed unexpectedly: %s",
+                    exc,
+                    exc_info=True,
+                )
+        self.active_command_id = None
 
     @property
     def query_id(self) -> Optional[str]:

tests/e2e/test_driver.py

@@ -342,6 +342,50 @@ def test_execute_async__large_result(self, extra_params):
 
             assert len(result) == x_dimension * y_dimension
 
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_execute_async__close_without_fetch_frees_handle(self, extra_params):
+        """Closing a cursor whose async command result was never fetched must free
+        the server-side statement handle (issue #791). Otherwise the handle leaks
+        until the session closes."""
+        with self.cursor(extra_params) as cursor:
+            cursor.execute_async("SELECT 1")
+
+            # Capture the server-side command id before we close the cursor.
+            command_id = cursor.active_command_id
+            assert command_id is not None
+
+            backend = cursor.backend
+
+            # Sanity: the handle is live and pollable before close.
+            backend.get_query_state(command_id)
+
+            # User decides not to wait for the result and closes the cursor
+            # without ever calling get_async_execution_result().
+            cursor.close()
+
+            # After close, the server-side handle must have been freed. How a
+            # re-poll of the saved command id surfaces that is backend-specific:
+            #   - Thrift raises a server error on the closed handle.
+            #   - SEA's get_query_state() does a plain GET and returns the
+            #     status.state, so a freed statement may come back as a terminal
+            #     CLOSED/CANCELLED state instead of raising.
+            # Accept either signal as proof the handle was freed. Pre-fix (leak),
+            # the poll instead returns a live/terminal-success state.
+            try:
+                state = backend.get_query_state(command_id)
+            except (RequestError, OperationalError, DatabaseError):
+                pass
+            else:
+                assert state in (CommandState.CLOSED, CommandState.CANCELLED)
+
 
 # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core
 # tests

tests/unit/test_client.py

@@ -161,6 +161,93 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                     # Should NOT have called backend.close_command (already closed)
                     mock_backend.close_command.assert_not_called()
 
+    def test_cursor_close_frees_command_when_no_result_set(self):
+        """Closing a cursor with an unfetched async command frees the server handle.
+
+        When active_result_set is None but active_command_id is set (an async
+        submission whose result was never fetched), close() must issue an
+        explicit backend.close_command so the server-side statement handle is
+        not leaked. This branch is backend-agnostic (Thrift/SEA/kernel).
+        """
+        mock_backend = Mock(spec=ThriftDatabricksClient)
+        mock_connection = Mock()
+        cursor = client.Cursor(connection=mock_connection, backend=mock_backend)
+
+        command_id = Mock(spec=CommandId)
+        cursor.active_command_id = command_id
+        cursor.active_result_set = None
+
+        cursor.close()
+
+        mock_backend.close_command.assert_called_once_with(command_id)
+        self.assertIsNone(cursor.active_command_id)
+
+    def test_cursor_close_skips_command_when_connection_closed(self):
+        """Closing a cursor after its session is gone must not call close_command.
+
+        Mirrors ResultSet.close, which gates backend.close_command on
+        connection.open, to avoid a network call on a dead session (and a
+        spurious warning) during shutdown ordering.
+        """
+        mock_backend = Mock(spec=ThriftDatabricksClient)
+        mock_connection = Mock()
+        mock_connection.open = False
+        cursor = client.Cursor(connection=mock_connection, backend=mock_backend)
+
+        cursor.active_command_id = Mock(spec=CommandId)
+        cursor.active_result_set = None
+
+        cursor.close()
+
+        mock_backend.close_command.assert_not_called()
+        self.assertIsNone(cursor.active_command_id)
+
+    def test_cursor_close_does_not_double_close_when_result_set_present(self):
+        """Closing a cursor with an active result set must NOT call close_command.
+
+        The result-set close path is responsible for freeing the handle in that
+        case, so the elif branch must not fire (no double-close).
+        """
+        mock_backend = Mock(spec=ThriftDatabricksClient)
+        mock_connection = Mock()
+        cursor = client.Cursor(connection=mock_connection, backend=mock_backend)
+
+        cursor.active_command_id = Mock(spec=CommandId)
+        result_set = Mock(spec=ResultSet)
+        cursor.active_result_set = result_set
+
+        cursor.close()
+
+        result_set.close.assert_called_once()
+        mock_backend.close_command.assert_not_called()
+        self.assertIsNone(cursor.active_command_id)
+
+    def test_cursor_close_tolerates_single_arg_request_error(self):
+        """A RequestError with only one positional arg must not escape close().
+
+        Some RequestErrors (e.g. from unified_http_client during shutdown) are
+        built with a single message arg, so e.args[1] would IndexError. close()
+        must guard the length check and swallow the error like any other
+        close_command failure.
+        """
+        from databricks.sql.exc import RequestError
+
+        mock_backend = Mock(spec=ThriftDatabricksClient)
+        mock_backend.close_command.side_effect = RequestError(
+            "HTTP client is closing or has been closed"
+        )
+        mock_connection = Mock()
+        cursor = client.Cursor(connection=mock_connection, backend=mock_backend)
+
+        cursor.active_command_id = Mock(spec=CommandId)
+        cursor.active_result_set = None
+
+        # Should not raise despite the single-arg RequestError.
+        cursor.close()
+
+        mock_backend.close_command.assert_called_once()
+        self.assertIsNone(cursor.active_command_id)
+
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_cant_open_cursor_on_closed_connection(self, mock_client_class):
         connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)