[471] - Close underlying HTTP Client on closing Connection

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

@@ -477,8 +477,8 @@ def attempt_request(attempt):
                     )
                 )
             finally:
-                # Calling `close()` here releases the active HTTP connection back to the pool
-                self._transport.close()
+                # Calling `release_connection()` here releases the active HTTP connection back to the pool
+                self._transport.release_connection()
 
             return RequestErrorInfo(
                 error=error,

src/databricks/sql/telemetry/telemetry_client.py

@@ -4,6 +4,7 @@
 import json
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import Future
+from concurrent.futures import wait
 from datetime import datetime, timezone
 from typing import List, Dict, Any, Optional, TYPE_CHECKING
 from databricks.sql.telemetry.models.event import (
@@ -182,6 +183,7 @@ def __init__(
         self._user_agent = None
         self._events_batch = []
         self._lock = threading.RLock()
+        self._pending_futures = set()
         self._driver_connection_params = None
         self._host_url = host_url
         self._executor = executor
@@ -245,6 +247,9 @@ def _send_telemetry(self, events):
                 timeout=900,
             )
 
+            with self._lock:
+                self._pending_futures.add(future)
+
             future.add_done_callback(
                 lambda fut: self._telemetry_request_callback(fut, sent_count=sent_count)
             )
@@ -303,6 +308,9 @@ def _telemetry_request_callback(self, future, sent_count: int):
 
         except Exception as e:
             logger.debug("Telemetry request failed with exception: %s", e)
+        finally:
+            with self._lock:
+                self._pending_futures.discard(future)
 
     def _export_telemetry_log(self, **telemetry_event_kwargs):
         """
@@ -356,10 +364,29 @@ def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
         )
 
     def close(self):
-        """Flush remaining events before closing"""
-        logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
+        """Schedules the client to be closed in the background."""
+        logger.debug(
+            "Scheduling background closure for TelemetryClient of connection %s",
+            self._session_id_hex,
+        )
+        self._executor.submit(self._close_and_wait)
+
+    def _close_and_wait(self):
+        """Flush remaining events and wait for them to complete before closing."""
         self._flush()
 
+        with self._lock:
+            futures_to_wait_on = list(self._pending_futures)
+
+        if futures_to_wait_on:
+            logger.debug(
+                "Waiting for %s pending telemetry requests to complete.",
+                len(futures_to_wait_on),
+            )
+            wait(futures_to_wait_on)
+
+        self._http_client.close()
+
 
 class TelemetryClientFactory:
     """
@@ -460,7 +487,6 @@ def initialize_telemetry_client(
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
-
             with TelemetryClientFactory._lock:
                 TelemetryClientFactory._initialize()
 

tests/unit/test_sea_backend.py

@@ -220,6 +220,7 @@ def test_session_management(self, sea_client, mock_http_client, thrift_session_i
             path=sea_client.SESSION_PATH_WITH_ID.format("test-session-789"),
             data={"session_id": "test-session-789", "warehouse_id": "abc123"},
         )
+        mock_http_client.close.assert_called_once()
 
         # Test close_session with invalid ID type
         with pytest.raises(ValueError) as excinfo:

tests/unit/test_thrift_backend.py

@@ -1436,8 +1436,12 @@ def test_op_handle_respected_in_close_command(self, tcli_service_class):
         )
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
-    def test_session_handle_respected_in_close_session(self, tcli_service_class):
+    @patch("databricks.sql.auth.thrift_http_client.THttpClient", autospec=True)
+    def test_session_handle_respected_in_close_session(
+        self, mock_http_client_class, tcli_service_class
+    ):
         tcli_service_instance = tcli_service_class.return_value
+        mock_http_client_instance = mock_http_client_class.return_value
         thrift_backend = ThriftDatabricksClient(
             "foobar",
             443,
@@ -1447,12 +1451,15 @@ def test_session_handle_respected_in_close_session(self, tcli_service_class):
             ssl_options=SSLOptions(),
             http_client=MagicMock(),
         )
+        thrift_backend._transport = mock_http_client_instance
+
         session_id = SessionId.from_thrift_handle(self.session_handle)
         thrift_backend.close_session(session_id)
         self.assertEqual(
             tcli_service_instance.CloseSession.call_args[0][0].sessionHandle,
             self.session_handle,
         )
+        mock_http_client_instance.close.assert_called_once()
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_non_arrow_non_column_based_set_triggers_exception(