ai: apply changes for #873 (2 review threads)

Addresses:
  - #3635951484 at tests/e2e/test_driver.py:379
  - #3635951489 at src/databricks/sql/client.py:1775

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

Author: peco-engineer-bot[bot]

src/databricks/sql/client.py

@@ -1772,7 +1772,7 @@ def close(self) -> None:
             try:
                 self.backend.close_command(self.active_command_id)
             except RequestError as e:
-                if isinstance(e.args[1], CursorAlreadyClosedError):
+                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.

tests/e2e/test_driver.py

@@ -371,11 +371,20 @@ def test_execute_async__close_without_fetch_frees_handle(self, extra_params):
             # without ever calling get_async_execution_result().
             cursor.close()
 
-            # After close, the server-side handle must have been freed, so a
-            # re-poll of the saved command id should raise a server error.
-            # Pre-fix (leak): this poll succeeds. Post-fix: it raises.
-            with pytest.raises((RequestError, OperationalError, DatabaseError)):
-                backend.get_query_state(command_id)
+            # 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/unit/test_client.py

@@ -222,6 +222,32 @@ def test_cursor_close_does_not_double_close_when_result_set_present(self):
         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)