diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 00a19154a..75871d387 100755 --- a/src/databricks/sql/client.py +++ b/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]: diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 5fe3db037..d5e170e99 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -342,6 +342,62 @@ 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_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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 66a3722ec..36189e251 100644 --- a/tests/unit/test_client.py +++ b/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