convert TelemetryClientFactory to module-level functions, replace NoopTelemetryClient class with NOOP_TELEMETRY_CLIENT singleton, updated tests accordingly

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -53,8 +53,9 @@
     TOperationState,
 )
 from databricks.sql.telemetry.telemetry_client import (
-    TelemetryClientFactory,
     TelemetryHelper,
+    initialize_telemetry_client,
+    get_telemetry_client,
 )
 from databricks.sql.telemetry.models.enums import DatabricksClientType
 from databricks.sql.telemetry.models.event import (
@@ -306,14 +307,14 @@ def read(self) -> Optional[OAuthToken]:
             kwargs.get("use_inline_params", False)
         )
 
-        TelemetryClientFactory.initialize_telemetry_client(
+        initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
             connection_uuid=self.get_session_id_hex(),
             auth_provider=auth_provider,
             host_url=self.host,
         )
 
-        self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
+        self._telemetry_client = get_telemetry_client(
             connection_uuid=self.get_session_id_hex()
         )
 

src/databricks/sql/exc.py

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

src/databricks/sql/telemetry/telemetry_client.py

@@ -115,27 +115,16 @@ def close(self):
         pass
 
 
-class NoopTelemetryClient(BaseTelemetryClient):
-    """
-    NoopTelemetryClient is a telemetry client that does not send any events to the server.
-    It is used when telemetry is disabled.
-    """
-
-    _instance = None
-
-    def __new__(cls):
-        if cls._instance is None:
-            cls._instance = super(NoopTelemetryClient, cls).__new__(cls)
-        return cls._instance
-
-    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
-        pass
-
-    def export_failure_log(self, error_name, error_message):
-        pass
-
-    def close(self):
-        pass
+# A single instance of the no-op client that can be reused
+NOOP_TELEMETRY_CLIENT = type(
+    "NoopTelemetryClient",
+    (BaseTelemetryClient,),
+    {
+        "export_initial_telemetry_log": lambda self, *args, **kwargs: None,
+        "export_failure_log": lambda self, *args, **kwargs: None,
+        "close": lambda self: None,
+    },
+)()
 
 
 class TelemetryClient(BaseTelemetryClient):
@@ -301,129 +290,111 @@ def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._connection_uuid)
         self._flush()
-        TelemetryClientFactory.close(self._connection_uuid)
-
-
-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 connection_uuid -> 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"""
-        with cls._lock:
-            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)"
-                )
-
-    @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")
+        _remove_telemetry_client(self._connection_uuid)
+
+
+# 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
+    with _lock:
+        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)"
+            )
 
-    @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()
+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")
 
-        # 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,
-        connection_uuid,
-        auth_provider,
-        host_url,
-    ):
-        """Initialize a telemetry client for a specific connection if telemetry is enabled"""
-        try:
-            TelemetryClientFactory._initialize()
+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__)
 
-            with TelemetryClientFactory._lock:
-                if connection_uuid not in TelemetryClientFactory._clients:
-                    logger.debug(
-                        "Creating new TelemetryClient for connection %s",
-                        connection_uuid,
-                    )
-                    if telemetry_enabled:
-                        TelemetryClientFactory._clients[
-                            connection_uuid
-                        ] = TelemetryClient(
-                            telemetry_enabled=telemetry_enabled,
-                            connection_uuid=connection_uuid,
-                            auth_provider=auth_provider,
-                            host_url=host_url,
-                            executor=TelemetryClientFactory._executor,
-                        )
-                    else:
-                        TelemetryClientFactory._clients[
-                            connection_uuid
-                        ] = 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[connection_uuid] = NoopTelemetryClient()
+    clients_to_close = list(_clients.values())
+    for client in clients_to_close:
+        client.close()
 
-    @staticmethod
-    def get_telemetry_client(connection_uuid):
-        """Get the telemetry client for a specific connection"""
-        try:
-            if connection_uuid in TelemetryClientFactory._clients:
-                return TelemetryClientFactory._clients[connection_uuid]
-            else:
-                logger.error(
-                    "Telemetry client not initialized for connection %s",
-                    connection_uuid,
-                )
-                return NoopTelemetryClient()
-        except Exception as e:
-            logger.debug("Failed to get telemetry client: %s", e)
-            return NoopTelemetryClient()
+    # Call the original exception handler to maintain normal behavior
+    if _original_excepthook:
+        _original_excepthook(exc_type, exc_value, exc_traceback)
 
-    @staticmethod
-    def close(connection_uuid):
-        """Close and remove the telemetry client for a specific connection"""
 
-        with TelemetryClientFactory._lock:
-            if connection_uuid in TelemetryClientFactory._clients:
-                logger.debug(
-                    "Removing telemetry client for connection %s", connection_uuid
-                )
-                TelemetryClientFactory._clients.pop(connection_uuid, None)
+def initialize_telemetry_client(
+    telemetry_enabled, connection_uuid, auth_provider, host_url
+):
+    """Initialize a telemetry client for a specific connection if telemetry is enabled"""
+    try:
+        _initialize()
 
-            # Shutdown executor if no more clients
-            if not TelemetryClientFactory._clients and TelemetryClientFactory._executor:
+        with _lock:
+            if connection_uuid not in _clients:
                 logger.debug(
-                    "No more telemetry clients, shutting down thread pool executor"
+                    "Creating new TelemetryClient for connection %s", connection_uuid
                 )
-                TelemetryClientFactory._executor.shutdown(wait=True)
-                TelemetryClientFactory._executor = None
-                TelemetryClientFactory._initialized = False
+                if telemetry_enabled:
+                    _clients[connection_uuid] = TelemetryClient(
+                        telemetry_enabled=telemetry_enabled,
+                        connection_uuid=connection_uuid,
+                        auth_provider=auth_provider,
+                        host_url=host_url,
+                        executor=_executor,
+                    )
+                else:
+                    _clients[connection_uuid] = NOOP_TELEMETRY_CLIENT
+    except Exception as e:
+        logger.debug("Failed to initialize telemetry client: %s", e)
+        # Fallback to NoopTelemetryClient to ensure connection doesn't fail
+        _clients[connection_uuid] = NOOP_TELEMETRY_CLIENT
+
+
+def get_telemetry_client(connection_uuid):
+    """Get the telemetry client for a specific connection"""
+    try:
+        if connection_uuid in _clients:
+            return _clients[connection_uuid]
+        else:
+            logger.error(
+                "Telemetry client not initialized for connection %s", connection_uuid
+            )
+            return NOOP_TELEMETRY_CLIENT
+    except Exception as e:
+        logger.debug("Failed to get telemetry client: %s", e)
+        return NOOP_TELEMETRY_CLIENT
+
+
+def _remove_telemetry_client(connection_uuid):
+    """Remove the telemetry client for a specific connection"""
+    global _initialized, _executor
+    with _lock:
+        if connection_uuid in _clients:
+            logger.debug("Removing telemetry client for connection %s", connection_uuid)
+            _clients.pop(connection_uuid, None)
+
+        # Shutdown executor if no more clients
+        if not _clients and _executor:
+            logger.debug(
+                "No more telemetry clients, shutting down thread pool executor"
+            )
+            _executor.shutdown(wait=True)
+            _executor = None
+            _initialized = False