Fix test for closeSession

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

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

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(

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()