diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 3b162225..b1ecaf3a 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -16,7 +16,7 @@ from io import BytesIO from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager, Retry -from databricks.sql.exc import MaxRetryDurationError, NonRecoverableNetworkError, UnsafeToRetryError +from databricks.sql.exc import MaxRetryDurationError, NonRecoverableNetworkError, UnsafeToRetryError, SessionAlreadyClosedError from enum import Enum, auto @@ -188,9 +188,9 @@ def should_retry(self, method: str, status_code: int) -> bool: if ( status_code == 404 and self.command_type in [CommandType.CANCEL_OPERATION, CommandType.CLOSE_SESSION] - and len(self.history) > 1 + and len(self.history) > 0 ): - return False + raise SessionAlreadyClosedError("CloseSession received 404 code from Databricks. Session is already closed.") # Request failed, was an ExecuteStatement and the command may have reached the server diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 347b9c10..4b1b708a 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -257,8 +257,9 @@ def _close(self, close_cursors=True) -> None: try: self.thrift_backend.close_session(self._session_handle) - except SessionAlreadyClosedError as e: - logger.info("Session was closed by a prior request") + except RequestError as e: + if isinstance(e.args[1], SessionAlreadyClosedError): + logger.info("Session was closed by a prior request") except DatabaseError as e: if "Invalid SessionHandle" in str(e): logger.warning( diff --git a/tests/e2e/common/retry_test_mixins.py b/tests/e2e/common/retry_test_mixins.py index 0fa9b137..d89817bf 100644 --- a/tests/e2e/common/retry_test_mixins.py +++ b/tests/e2e/common/retry_test_mixins.py @@ -198,8 +198,6 @@ def test_retry_abort_close_session_on_404(self): ] with self.connection(extra_params={**self._retry_policy}) as conn: - conn.close() - conn.open = True with mock_sequential_server_responses(responses): with self.assertLogs("databricks.sql", level="INFO",) as cm: conn.close()