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