diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d3f8831..35a8e0aa 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 53ec9735..6ab8b45f 100644 --- a/CONTRIBUTING.md +++ b/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` diff --git a/pyproject.toml b/pyproject.toml index f522d1ce..f651421e 100644 --- a/pyproject.toml +++ b/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" ] diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 722ed778..14e59df6 100644 --- a/src/databricks/sql/client.py +++ b/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): diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index b4afeaa3..7756c56a 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/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) diff --git a/test.env.example b/test.env.example index 5ce7eca5..94aed419 100644 --- a/test.env.example +++ b/test.env.example @@ -4,4 +4,8 @@ http_path="" access_token="" # Only required to run the PySQLStagingIngestionTestSuite -staging_ingestion_user="" \ No newline at end of file +staging_ingestion_user="" + +# Only required to run SQLAlchemy tests +catalog="" +schema="" \ No newline at end of file diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 37428116..26b7d186 100644 --- a/tests/e2e/test_driver.py +++ b/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