[471] - Close underlying HTTP Client on closing Connection

src/databricks/sql/auth/thrift_http_client.py

@@ -147,11 +147,14 @@ 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)
 

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

@@ -287,9 +287,9 @@ def close_session(self, session_id: SessionId) -> None:
 
         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()
+        if sea_session_id is None:
+            raise ValueError("Not a valid SEA session ID")
 
         request_data = DeleteSessionRequest(
             warehouse_id=self.warehouse_id,
@@ -302,6 +302,9 @@ def close_session(self, session_id: SessionId) -> None:
             data=request_data.to_dict(),
         )
 
+        # close the HTTP client
+        self._http_client.close()
+
     def _extract_description_from_manifest(
         self, manifest: ResultManifest
     ) -> List[Tuple]:

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

@@ -183,7 +183,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

@@ -483,8 +483,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

@@ -2,6 +2,7 @@
 import time
 import logging
 import json
+from concurrent.futures import ThreadPoolExecutor, wait
 from queue import Queue, Full
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import Future
@@ -189,6 +190,7 @@ def __init__(
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
+        self._pending_futures = set()
 
         # OPTIMIZATION: Use lock-free Queue instead of list + lock
         # Queue is thread-safe internally and has better performance under concurrency
@@ -293,6 +295,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)
             )
@@ -351,6 +356,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):
         """
@@ -404,20 +412,40 @@ 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)
+        """Schedule client closure."""
+        logger.debug(
+            "Scheduling 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:
+            pending_events = list(self._pending_futures)
+
+        if pending_events:
+            logger.debug(
+                "Waiting for %s pending telemetry requests to complete.",
+                len(pending_events),
+            )
+            wait(pending_events)
+
+        logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
+        self._http_client.close()
+
 
 class TelemetryClientFactory:
     """
     Static factory class for creating and managing telemetry clients.
     It uses a thread pool to handle asynchronous operations and a single flush thread for all clients.
     """
 
-    _clients: Dict[
-        str, BaseTelemetryClient
-    ] = {}  # Map of session_id_hex -> BaseTelemetryClient
+    _clients: Dict[str, BaseTelemetryClient] = (
+        {}
+    )  # Map of session_id_hex -> BaseTelemetryClient
     _executor: Optional[ThreadPoolExecutor] = None
     _initialized: bool = False
     _lock = threading.RLock()  # Thread safety for factory operations
@@ -518,21 +546,21 @@ def initialize_telemetry_client(
                         session_id_hex,
                     )
                     if telemetry_enabled:
-                        TelemetryClientFactory._clients[
-                            session_id_hex
-                        ] = TelemetryClient(
-                            telemetry_enabled=telemetry_enabled,
-                            session_id_hex=session_id_hex,
-                            auth_provider=auth_provider,
-                            host_url=host_url,
-                            executor=TelemetryClientFactory._executor,
-                            batch_size=batch_size,
-                            client_context=client_context,
+                        TelemetryClientFactory._clients[session_id_hex] = (
+                            TelemetryClient(
+                                telemetry_enabled=telemetry_enabled,
+                                session_id_hex=session_id_hex,
+                                auth_provider=auth_provider,
+                                host_url=host_url,
+                                executor=TelemetryClientFactory._executor,
+                                batch_size=batch_size,
+                                client_context=client_context,
+                            )
                         )
                     else:
-                        TelemetryClientFactory._clients[
-                            session_id_hex
-                        ] = NoopTelemetryClient()
+                        TelemetryClientFactory._clients[session_id_hex] = (
+                            NoopTelemetryClient()
+                        )
         except Exception as e:
             logger.debug("Failed to initialize telemetry client: %s", e)
             # Fallback to NoopTelemetryClient to ensure connection doesn't fail

tests/unit/test_sea_backend.py

@@ -222,6 +222,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(