From 5ae9486b50a07e848e36a6df9a15d58a80dd0e14 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 22 Aug 2022 20:36:58 -0500 Subject: [PATCH 1/2] Add test: cursors are closed when connection closes Signed-off-by: Jesse Whitehouse --- tests/e2e/driver_tests.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/e2e/driver_tests.py b/tests/e2e/driver_tests.py index 358f0b263..9e400770d 100644 --- a/tests/e2e/driver_tests.py +++ b/tests/e2e/driver_tests.py @@ -544,6 +544,32 @@ def test_decimal_not_returned_as_strings_arrow(self): decimal_type = arrow_df.field(0).type self.assertTrue(pyarrow.types.is_decimal(decimal_type)) + def test_close_connection_closes_cursors(self): + + from databricks.sql.thrift_api.TCLIService import ttypes + + with self.connection() as conn: + cursor = conn.cursor() + cursor.execute('SELECT id, id `id2`, id `id3` FROM RANGE(1000000) order by RANDOM()') + ars = cursor.active_result_set + + # We must manually run this check because thrift_backend always forces `has_been_closed_server_side` to True + + # Cursor op state should be open before connection is closed + status_request = ttypes.TGetOperationStatusReq(operationHandle=ars.command_id, getProgressUpdate=False) + op_status_at_server = ars.thrift_backend._client.GetOperationStatus(status_request) + assert op_status_at_server.operationState != ttypes.TOperationState.CLOSED_STATE + + conn.close() + + # When connection closes, any cursor operations should no longer exist at the server + with self.assertRaises(thrift.Thrift.TApplicationException) as cm: + op_status_at_server = ars.thrift_backend._client.GetOperationStatus(status_request) + if hasattr(cm, "exception"): + assert "RESOURCE_DOES_NOT_EXIST" in cm.exception.message + + + # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep # the 429/503 subsuites separate since they execute under different circumstances. From b80cb7120c4bd9433da9c9641a3113f0a3a5c7bb Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 22 Aug 2022 20:37:08 -0500 Subject: [PATCH 2/2] Close cursors before closing connection Signed-off-by: Jesse Whitehouse --- src/databricks/sql/client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 717980af4..2286c89bd 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -181,12 +181,11 @@ def close(self) -> None: self._close() def _close(self, close_cursors=True) -> None: - self.thrift_backend.close_session(self._session_handle) - self.open = False - if close_cursors: for cursor in self._cursors: cursor.close() + self.thrift_backend.close_session(self._session_handle) + self.open = False def commit(self): """No-op because Databricks does not support transactions"""