Don't raise exception when closing a stale Thrift session (databricks#159)

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 2.6.x (Unreleased)
 
+- Fix: connector raised exception when calling close() on a closed Thrift session
 - Improve e2e test development ergonomics
 - Redact logged thrift responses by default
 - Add support for OAuth on Databricks Azure

CONTRIBUTING.md

@@ -118,6 +118,8 @@ access_token="dapi***"
 staging_ingestion_user="***@example.com"
 ```
 
+To see logging output from pytest while running tests, set `log_cli = "true"` under `tool.pytest.ini_options` in `pyproject.toml`. You can also set `log_cli_level` to any of the default Python log levels: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL`
+
 There are several e2e test suites available:
 - `PySQLCoreTestSuite`
 - `PySQLLargeQueriesSuite`

pyproject.toml

@@ -54,6 +54,8 @@ exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck
 
 [tool.pytest.ini_options]
 minversion = "6.0"
+log_cli = "false"
+log_cli_level = "INFO"
 testpaths = [
     "tests"
 ]

src/databricks/sql/client.py

@@ -190,7 +190,7 @@ def read(self) -> Optional[OAuthToken]:
             session_configuration, catalog, schema
         )
         self.open = True
-        logger.info("Successfully opened session " + str(self.get_session_id()))
+        logger.info("Successfully opened session " + str(self.get_session_id_hex()))
         self._cursors = []  # type: List[Cursor]
 
     def __enter__(self):
@@ -214,6 +214,9 @@ def __del__(self):
     def get_session_id(self):
         return self.thrift_backend.handle_to_id(self._session_handle)
 
+    def get_session_id_hex(self):
+        return self.thrift_backend.handle_to_hex_id(self._session_handle)
+
     def cursor(
         self,
         arraysize: int = DEFAULT_ARRAY_SIZE,
@@ -244,7 +247,25 @@ def _close(self, close_cursors=True) -> None:
         if close_cursors:
             for cursor in self._cursors:
                 cursor.close()
-        self.thrift_backend.close_session(self._session_handle)
+
+        logger.info(f"Closing session {self.get_session_id_hex()}")
+        if not self.open:
+            logger.debug("Session appears to have been closed already")
+
+        try:
+            self.thrift_backend.close_session(self._session_handle)
+        except DatabaseError as e:
+            if "Invalid SessionHandle" in str(e):
+                logger.warning(
+                    f"Attempted to close session that was already closed: {e}"
+                )
+            else:
+                logger.warning(
+                    f"Attempt to close session raised an exception at the server: {e}"
+                )
+        except Exception as e:
+            logger.error(f"Attempt to close session raised a local exception: {e}")
+
         self.open = False
 
     def commit(self):

src/databricks/sql/thrift_backend.py

@@ -3,6 +3,7 @@
 import logging
 import math
 import time
+import uuid
 import threading
 import lz4.frame
 from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
@@ -1021,3 +1022,8 @@ def cancel_command(self, active_op_handle):
     @staticmethod
     def handle_to_id(session_handle):
         return session_handle.sessionId.guid
+
+    @staticmethod
+    def handle_to_hex_id(session_handle: TCLIService.TSessionHandle):
+        this_uuid = uuid.UUID(bytes=session_handle.sessionId.guid)
+        return str(this_uuid)

test.env.example

@@ -4,4 +4,8 @@ http_path=""
 access_token=""
 
 # Only required to run the PySQLStagingIngestionTestSuite
-staging_ingestion_user=""
+staging_ingestion_user=""
+
+# Only required to run SQLAlchemy tests
+catalog=""
+schema=""

tests/e2e/test_driver.py

@@ -616,6 +616,22 @@ def test_close_connection_closes_cursors(self):
                     assert "RESOURCE_DOES_NOT_EXIST" in cm.exception.message
 
 
+    def test_closing_a_closed_connection_doesnt_fail(self):
+
+        with self.assertLogs("databricks.sql", level="DEBUG",) as cm:
+            # Second .close() call is when this context manager exits
+            with self.connection() as conn:
+                # First .close() call is explicit here
+                conn.close()
+
+            expected_message_was_found = False
+            for log in cm.output:
+                if expected_message_was_found:
+                    break
+                target = "Session appears to have been closed already"
+                expected_message_was_found = target in log
+            
+            self.assertTrue(expected_message_was_found, "Did not find expected log messages")
 
 
 # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep