removed debugging, added TelemetryClientFactory

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -54,9 +54,7 @@
 )
 from databricks.sql.telemetry.telemetry_client import (
     TelemetryHelper,
-    initialize_telemetry_client,
-    get_telemetry_client,
-    close_telemetry_client,
+    TelemetryClientFactory,
 )
 from databricks.sql.telemetry.models.enums import DatabricksClientType
 from databricks.sql.telemetry.models.event import (
@@ -308,14 +306,14 @@ def read(self) -> Optional[OAuthToken]:
             kwargs.get("use_inline_params", False)
         )
 
-        initialize_telemetry_client(
+        TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
             session_id_hex=self.get_session_id_hex(),
             auth_provider=auth_provider,
             host_url=self.host,
         )
 
-        self._telemetry_client = get_telemetry_client(
+        self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
             session_id_hex=self.get_session_id_hex()
         )
 
@@ -472,7 +470,7 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
-        close_telemetry_client(self.get_session_id_hex())
+        TelemetryClientFactory.close(self.get_session_id_hex())
 
     def commit(self):
         """No-op because Databricks does not support transactions"""

src/databricks/sql/exc.py

@@ -1,7 +1,7 @@
 import json
 import logging
 
-from databricks.sql.telemetry.telemetry_client import get_telemetry_client
+from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +22,9 @@ def __init__(
 
         error_name = self.__class__.__name__
         if session_id_hex:
-            telemetry_client = get_telemetry_client(session_id_hex)
+            telemetry_client = TelemetryClientFactory.get_telemetry_client(
+                session_id_hex
+            )
             telemetry_client.export_failure_log(error_name, self.message)
 
     def __str__(self):

src/databricks/sql/telemetry/telemetry_client.py

@@ -305,111 +305,131 @@ def close(self):
         self._flush()
 
 
-# Module-level state
-_clients: Dict[str, BaseTelemetryClient] = {}
-_executor: Optional[ThreadPoolExecutor] = None
-_initialized: bool = False
-_lock = threading.Lock()
-_original_excepthook = None
-_excepthook_installed = False
-
-
-def _initialize():
-    """Initialize the telemetry system if not already initialized"""
-    global _initialized, _executor
-    if not _initialized:
-        _clients.clear()
-        _executor = ThreadPoolExecutor(max_workers=10)
-        _install_exception_hook()
-        _initialized = True
-        logger.debug("Telemetry system initialized with thread pool (max_workers=10)")
-
-
-def _install_exception_hook():
-    """Install global exception handler for unhandled exceptions"""
-    global _excepthook_installed, _original_excepthook
-    if not _excepthook_installed:
-        _original_excepthook = sys.excepthook
-        sys.excepthook = _handle_unhandled_exception
-        _excepthook_installed = True
-        logger.debug("Global exception handler installed for telemetry")
-
-
-def _handle_unhandled_exception(exc_type, exc_value, exc_traceback):
-    """Handle unhandled exceptions by sending telemetry and flushing thread pool"""
-    logger.debug("Handling unhandled exception: %s", exc_type.__name__)
-
-    clients_to_close = list(_clients.values())
-    for client in clients_to_close:
-        client.close()
-
-    # Call the original exception handler to maintain normal behavior
-    if _original_excepthook:
-        _original_excepthook(exc_type, exc_value, exc_traceback)
-
-
-def initialize_telemetry_client(
-    telemetry_enabled, session_id_hex, auth_provider, host_url
-):
-    """Initialize a telemetry client for a specific connection if telemetry is enabled"""
-    try:
-        with _lock:
-            _initialize()
-            if session_id_hex not in _clients:
-                logger.debug(
-                    "Creating new TelemetryClient for connection %s", session_id_hex
-                )
-                if telemetry_enabled:
-                    _clients[session_id_hex] = TelemetryClient(
-                        telemetry_enabled=telemetry_enabled,
-                        session_id_hex=session_id_hex,
-                        auth_provider=auth_provider,
-                        host_url=host_url,
-                        executor=_executor,
-                    )
-                    print("i have initialized the telemetry client yes")
-                else:
-                    _clients[session_id_hex] = NoopTelemetryClient()
-                    print("i have initialized the noop client yes")
-    except Exception as e:
-        logger.debug("Failed to initialize telemetry client: %s", e)
-        # Fallback to NoopTelemetryClient to ensure connection doesn't fail
-        _clients[session_id_hex] = NoopTelemetryClient()
-
-
-def get_telemetry_client(session_id_hex):
-    """Get the telemetry client for a specific connection"""
-    try:
-        if session_id_hex in _clients:
-            return _clients[session_id_hex]
-        else:
-            logger.error(
-                "Telemetry client not initialized for connection %s", session_id_hex
+class TelemetryClientFactory:
+    """
+    Static factory class for creating and managing telemetry clients.
+    It uses a thread pool to handle asynchronous operations.
+    """
+
+    _clients: Dict[
+        str, BaseTelemetryClient
+    ] = {}  # Map of session_id_hex -> BaseTelemetryClient
+    _executor: Optional[ThreadPoolExecutor] = None
+    _initialized: bool = False
+    _lock = threading.Lock()  # Thread safety for factory operations
+    _original_excepthook = None
+    _excepthook_installed = False
+
+    @classmethod
+    def _initialize(cls):
+        """Initialize the factory if not already initialized"""
+
+        if not cls._initialized:
+            cls._clients = {}
+            cls._executor = ThreadPoolExecutor(
+                max_workers=10
+            )  # Thread pool for async operations TODO: Decide on max workers
+            cls._install_exception_hook()
+            cls._initialized = True
+            logger.debug(
+                "TelemetryClientFactory initialized with thread pool (max_workers=10)"
             )
-            return NoopTelemetryClient()
-    except Exception as e:
-        logger.debug("Failed to get telemetry client: %s", e)
-        return NoopTelemetryClient()
-
-
-def close_telemetry_client(session_id_hex):
-    """Remove the telemetry client for a specific connection"""
-    global _initialized, _executor
-    with _lock:
-        # if (telemetry_client := _clients.pop(session_id_hex, None)) is not None:
-        if session_id_hex in _clients:
-            telemetry_client = _clients.pop(session_id_hex)
-            logger.debug("Removing telemetry client for connection %s", session_id_hex)
-            telemetry_client.close()
-
-        # Shutdown executor if no more clients
+
+    @classmethod
+    def _install_exception_hook(cls):
+        """Install global exception handler for unhandled exceptions"""
+        if not cls._excepthook_installed:
+            cls._original_excepthook = sys.excepthook
+            sys.excepthook = cls._handle_unhandled_exception
+            cls._excepthook_installed = True
+            logger.debug("Global exception handler installed for telemetry")
+
+    @classmethod
+    def _handle_unhandled_exception(cls, exc_type, exc_value, exc_traceback):
+        """Handle unhandled exceptions by sending telemetry and flushing thread pool"""
+        logger.debug("Handling unhandled exception: %s", exc_type.__name__)
+
+        clients_to_close = list(cls._clients.values())
+        for client in clients_to_close:
+            client.close()
+
+        # Call the original exception handler to maintain normal behavior
+        if cls._original_excepthook:
+            cls._original_excepthook(exc_type, exc_value, exc_traceback)
+
+    @staticmethod
+    def initialize_telemetry_client(
+        telemetry_enabled,
+        session_id_hex,
+        auth_provider,
+        host_url,
+    ):
+        """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
-            if not _clients and _executor:
+
+            with TelemetryClientFactory._lock:
+                TelemetryClientFactory._initialize()
+
+                if session_id_hex not in TelemetryClientFactory._clients:
+                    logger.debug(
+                        "Creating new TelemetryClient for connection %s",
+                        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,
+                        )
+                    else:
+                        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
+            TelemetryClientFactory._clients[session_id_hex] = NoopTelemetryClient()
+
+    @staticmethod
+    def get_telemetry_client(session_id_hex):
+        """Get the telemetry client for a specific connection"""
+        try:
+            if session_id_hex in TelemetryClientFactory._clients:
+                return TelemetryClientFactory._clients[session_id_hex]
+            else:
+                logger.error(
+                    "Telemetry client not initialized for connection %s",
+                    session_id_hex,
+                )
+                return NoopTelemetryClient()
+        except Exception as e:
+            logger.debug("Failed to get telemetry client: %s", e)
+            return NoopTelemetryClient()
+
+    @staticmethod
+    def close(session_id_hex):
+        """Close and remove the telemetry client for a specific connection"""
+
+        with TelemetryClientFactory._lock:
+            if (
+                telemetry_client := TelemetryClientFactory._clients.pop(
+                    session_id_hex, None
+                )
+            ) is not None:
+                logger.debug(
+                    "Removing telemetry client for connection %s", session_id_hex
+                )
+                telemetry_client.close()
+
+            # Shutdown executor if no more clients
+            if not TelemetryClientFactory._clients and TelemetryClientFactory._executor:
                 logger.debug(
                     "No more telemetry clients, shutting down thread pool executor"
                 )
-                _executor.shutdown(wait=True)
-                _executor = None
-                _initialized = False
-        except Exception as e:
-            logger.debug("Failed to shutdown thread pool executor: %s", e)
+                TelemetryClientFactory._executor.shutdown(wait=True)
+                TelemetryClientFactory._executor = None
+                TelemetryClientFactory._initialized = False

tests/unit/test_client.py

@@ -336,7 +336,6 @@ def test_negative_fetch_throws_exception(self):
             result_set.fetchmany(-1)
 
     def test_context_manager_closes_cursor(self):
-        print("hellow")
         mock_close = Mock()
         with client.Cursor(Mock(), Mock()) as cursor:
             cursor.close = mock_close
@@ -352,8 +351,7 @@ def test_context_manager_closes_cursor(self):
             cursor.close.assert_called()
 
     @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
-    def test_a_context_manager_closes_connection(self, mock_client_class):
-        print("hellow1")
+    def test_context_manager_closes_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
         mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
@@ -792,7 +790,6 @@ def test_cursor_context_manager_handles_exit_exception(self):
 
     def test_connection_close_handles_cursor_close_exception(self):
         """Test that _close handles exceptions from cursor.close() properly."""
-        print("banana")
         cursors_closed = []
 
         def mock_close_with_exception():