made telemetry client close a module-level function

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/client.py

@@ -56,6 +56,7 @@
     TelemetryHelper,
     initialize_telemetry_client,
     get_telemetry_client,
+    close_telemetry_client,
 )
 from databricks.sql.telemetry.models.enums import DatabricksClientType
 from databricks.sql.telemetry.models.event import (
@@ -471,7 +472,7 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
-        self._telemetry_client.close()
+        close_telemetry_client(self.get_session_id_hex())
 
     def commit(self):
         """No-op because Databricks does not support transactions"""

src/databricks/sql/telemetry/telemetry_client.py

@@ -300,7 +300,6 @@ def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
         self._flush()
-        _remove_telemetry_client(self._session_id_hex)
 
 
 # Module-level state
@@ -392,13 +391,14 @@ def get_telemetry_client(session_id_hex):
         return NOOP_TELEMETRY_CLIENT
 
 
-def _remove_telemetry_client(session_id_hex):
+def close_telemetry_client(session_id_hex):
     """Remove the telemetry client for a specific connection"""
     global _initialized, _executor
     with _lock:
         if session_id_hex in _clients:
             logger.debug("Removing telemetry client for connection %s", session_id_hex)
-            _clients.pop(session_id_hex, None)
+            telemetry_client = _clients.pop(session_id_hex, None)
+            telemetry_client.close()
 
         # Shutdown executor if no more clients
         if not _clients and _executor:

tests/unit/test_telemetry.py

@@ -8,7 +8,7 @@
     NOOP_TELEMETRY_CLIENT,
     initialize_telemetry_client,
     get_telemetry_client,
-    _remove_telemetry_client,
+    close_telemetry_client,
     TelemetryHelper,
     BaseTelemetryClient
 )
@@ -385,7 +385,33 @@ def test_close_telemetry_client(self, telemetry_system_reset):
         client = get_telemetry_client(session_id_hex)
         assert isinstance(client, TelemetryClient)
 
-        _remove_telemetry_client(session_id_hex)
+        client.close = MagicMock()
+        
+        close_telemetry_client(session_id_hex)
+        
+        client.close.assert_called_once()
+        
+        client = get_telemetry_client(session_id_hex)
+        assert client is NOOP_TELEMETRY_CLIENT
+
+    def test_close_telemetry_client_noop(self, telemetry_system_reset):
+        """Test closing a no-op telemetry client."""
+        session_id_hex = "test-uuid"
+        initialize_telemetry_client(
+            telemetry_enabled=False,
+            session_id_hex=session_id_hex,
+            auth_provider=MagicMock(),
+            host_url="test-host",
+        )
+        
+        client = get_telemetry_client(session_id_hex)
+        assert client is NOOP_TELEMETRY_CLIENT
+        
+        client.close = MagicMock()
+        
+        close_telemetry_client(session_id_hex)
+        
+        client.close.assert_called_once()
 
         client = get_telemetry_client(session_id_hex)
         assert client is NOOP_TELEMETRY_CLIENT