Skip to content

Commit ebf4ac4

Browse files
committed
[ES-402013] Check that closing the connection also closes any active cursors.
Signed-off-by: Jesse Whitehouse <[email protected]>
1 parent 9bb2323 commit ebf4ac4

File tree

2 files changed

+15
-3
lines changed

2 files changed

+15
-3
lines changed

src/databricks/sql/client.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -181,13 +181,14 @@ def close(self) -> None:
181181
self._close()
182182

183183
def _close(self, close_cursors=True) -> None:
184-
self.thrift_backend.close_session(self._session_handle)
185-
self.open = False
186-
187184
if close_cursors:
188185
for cursor in self._cursors:
189186
cursor.close()
190187

188+
self.thrift_backend.close_session(self._session_handle)
189+
self.open = False
190+
191+
191192
def commit(self):
192193
"""No-op because Databricks does not support transactions"""
193194
pass

tests/e2e/driver_tests.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,17 @@ def test_decimal_not_returned_as_strings_arrow(self):
544544
decimal_type = arrow_df.field(0).type
545545
self.assertTrue(pyarrow.types.is_decimal(decimal_type))
546546

547+
def test_close_connection_closes_cursors(self):
548+
with self.connection() as conn:
549+
cursor = conn.cursor()
550+
cursor.execute('SELECT id, id `id2`, id `id3` FROM RANGE(1000000) order by RANDOM()')
551+
ars = cursor.active_result_set
552+
assert not bool(ars.has_been_closed_server_side)
553+
conn.close()
554+
assert bool(ars.has_been_closed_server_side)
555+
cursor.close()
556+
557+
547558

548559
# use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
549560
# the 429/503 subsuites separate since they execute under different circumstances.

0 commit comments

Comments
 (0)