preliminary connection closure func

Author: Varun0157

src/databricks/sql/auth/thrift_http_client.py

@@ -105,7 +105,6 @@ def startRetryTimer(self):
         self.retry_policy and self.retry_policy.start_retry_timer()
 
     def open(self):
-
         # self.__pool replaces the self.__http used by the original THttpClient
         _pool_kwargs = {"maxsize": self.max_connections}
 
@@ -140,19 +139,21 @@ def open(self):
         else:
             self.__pool = pool_class(self.host, self.port, **_pool_kwargs)
 
-    def close(self):
+    def release_connection(self):
         self.__resp and self.__resp.drain_conn()
         self.__resp and self.__resp.release_conn()
         self.__resp = None
 
+    def close(self):
+        self.__pool.close()
+
     def read(self, sz):
         return self.__resp.read(sz)
 
     def isOpen(self):
         return self.__resp is not None
 
     def flush(self):
-
         # Pull data out of buffer that will be sent in this request
         data = self.__wbuf.getvalue()
         self.__wbuf = BytesIO()

src/databricks/sql/backend/sea/backend.py

@@ -273,7 +273,7 @@ def open_session(
 
         return SessionId.from_sea_session_id(session_id)
 
-    def close_session(self, session_id: SessionId) -> None:
+    def _close_session(self, session_id: SessionId) -> None:
         """
         Closes an existing session with the Databricks SQL service.
 
@@ -285,8 +285,6 @@ def close_session(self, session_id: SessionId) -> None:
             OperationalError: If there's an error closing the session
         """
 
-        logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
-
         if session_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
@@ -302,6 +300,23 @@ def close_session(self, session_id: SessionId) -> None:
             data=request_data.to_dict(),
         )
 
+    def close_session(self, session_id: SessionId) -> None:
+        """
+        Closes the session and the underlying HTTP client.
+
+        Args:
+            session_id: The session identifier returned by open_session()
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error closing the session
+        """
+
+        logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
+
+        self._close_session(session_id)
+        self._http_client.close()
+
     def _extract_description_from_manifest(
         self, manifest: ResultManifest
     ) -> List[Tuple]:

src/databricks/sql/backend/sea/utils/http_client.py

@@ -197,7 +197,7 @@ def _open(self):
     def close(self):
         """Close the connection pool."""
         if self._pool:
-            self._pool.clear()
+            self._pool.close()
 
     def using_proxy(self) -> bool:
         """Check if proxy is being used."""

src/databricks/sql/backend/thrift_backend.py

@@ -232,7 +232,7 @@ def __init__(
         try:
             self._transport.open()
         except:
-            self._transport.close()
+            self._transport.release_connection()
             raise
 
         self._request_lock = threading.RLock()
@@ -478,7 +478,7 @@ def attempt_request(attempt):
                 )
             finally:
                 # Calling `close()` here releases the active HTTP connection back to the pool
-                self._transport.close()
+                self._transport.release_connection()
 
             return RequestErrorInfo(
                 error=error,
@@ -607,7 +607,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             self._session_id_hex = session_id.hex_guid
             return session_id
         except:
-            self._transport.close()
+            self._transport.release_connection()
             raise
 
     def close_session(self, session_id: SessionId) -> None:
@@ -619,7 +619,8 @@ def close_session(self, session_id: SessionId) -> None:
         try:
             self.make_request(self._client.CloseSession, req)
         finally:
-            self._transport.close()
+            self._transport.release_connection()
+        self._transport.close()
 
     def _check_command_not_in_error_or_closed_state(
         self, op_handle, get_operations_resp