Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
27 changes: 27 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down