removed ABC from mixin, added try/catch block around executor shutdown

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

Author: saishreeeee

.github/workflows/code-quality-checks.yml

@@ -61,7 +61,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       - name: Run tests
-        run: poetry run python -m pytest tests/unit -v -s --log-cli-level=DEBUG
+        run: poetry run python -m pytest tests/unit 
   run-unit-tests-with-arrow:
     runs-on: ubuntu-latest
     strategy:
@@ -112,7 +112,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       -   name: Run tests
-          run: poetry run python -m pytest tests/unit -v -s --log-cli-level=DEBUG
+          run: poetry run python -m pytest tests/unit
   check-linting:
     runs-on: ubuntu-latest
     strategy:

src/databricks/sql/telemetry/telemetry_client.py

@@ -123,6 +123,13 @@ class NoopTelemetryClient(BaseTelemetryClient):
     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
 
@@ -133,10 +140,6 @@ def close(self):
         pass
 
 
-# A single instance of the no-op client that can be reused
-NOOP_TELEMETRY_CLIENT = NoopTelemetryClient()
-
-
 class TelemetryClient(BaseTelemetryClient):
     """
     Telemetry client class that handles sending telemetry events in batches to the server.
@@ -369,11 +372,11 @@ def initialize_telemetry_client(
                         executor=_executor,
                     )
                 else:
-                    _clients[session_id_hex] = NOOP_TELEMETRY_CLIENT
+                    _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
-        _clients[session_id_hex] = NOOP_TELEMETRY_CLIENT
+        _clients[session_id_hex] = NoopTelemetryClient()
 
 
 def get_telemetry_client(session_id_hex):
@@ -385,10 +388,10 @@ def get_telemetry_client(session_id_hex):
             logger.error(
                 "Telemetry client not initialized for connection %s", session_id_hex
             )
-            return NOOP_TELEMETRY_CLIENT
+            return NoopTelemetryClient()
     except Exception as e:
         logger.debug("Failed to get telemetry client: %s", e)
-        return NOOP_TELEMETRY_CLIENT
+        return NoopTelemetryClient()
 
 
 def close_telemetry_client(session_id_hex):
@@ -401,10 +404,13 @@ def close_telemetry_client(session_id_hex):
             telemetry_client.close()
 
         # 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
+        try:
+            if not _clients and _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)

src/databricks/sql/telemetry/utils.py

@@ -1,10 +1,9 @@
 import json
 from enum import Enum
 from dataclasses import asdict, is_dataclass
-from abc import ABC
 
 
-class JsonSerializableMixin(ABC):
+class JsonSerializableMixin:
     """Mixin class to provide JSON serialization capabilities to dataclasses."""
 
     def to_json(self) -> str:

tests/unit/test_telemetry.py

@@ -5,7 +5,7 @@
 
 from databricks.sql.telemetry.telemetry_client import (
     TelemetryClient,
-    NOOP_TELEMETRY_CLIENT,
+    NoopTelemetryClient,
     initialize_telemetry_client,
     get_telemetry_client,
     close_telemetry_client,
@@ -30,8 +30,8 @@
 
 @pytest.fixture
 def noop_telemetry_client():
-    """Fixture for NOOP_TELEMETRY_CLIENT."""
-    return NOOP_TELEMETRY_CLIENT
+    """Fixture for NoopTelemetryClient."""
+    return NoopTelemetryClient()
 
 
 @pytest.fixture
@@ -79,7 +79,13 @@ def telemetry_system_reset():
 
 
 class TestNoopTelemetryClient:
-    """Tests for the NOOP_TELEMETRY_CLIENT."""
+    """Tests for the NoopTelemetryClient."""
+
+    def test_singleton(self):
+        """Test that NoopTelemetryClient is a singleton."""
+        client1 = NoopTelemetryClient()
+        client2 = NoopTelemetryClient()
+        assert client1 is client2
 
     def test_export_initial_telemetry_log(self, noop_telemetry_client):
         """Test that export_initial_telemetry_log does nothing."""
@@ -362,12 +368,12 @@ def test_initialize_telemetry_client_disabled(self, telemetry_system_reset):
         )
 
         client = get_telemetry_client(session_id_hex)
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert isinstance(client, NoopTelemetryClient)
 
     def test_get_telemetry_client_nonexistent(self, telemetry_system_reset):
         """Test getting a non-existent telemetry client."""
         client = get_telemetry_client("nonexistent-uuid")
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert isinstance(client, NoopTelemetryClient)
 
     def test_close_telemetry_client(self, telemetry_system_reset):
         """Test closing a telemetry client."""
@@ -392,7 +398,7 @@ def test_close_telemetry_client(self, telemetry_system_reset):
         client.close.assert_called_once()
 
         client = get_telemetry_client(session_id_hex)
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert isinstance(client, NoopTelemetryClient)
 
     def test_close_telemetry_client_noop(self, telemetry_system_reset):
         """Test closing a no-op telemetry client."""
@@ -405,7 +411,7 @@ def test_close_telemetry_client_noop(self, telemetry_system_reset):
         )
 
         client = get_telemetry_client(session_id_hex)
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert isinstance(client, NoopTelemetryClient)
 
         client.close = MagicMock()
 
@@ -414,7 +420,7 @@ def test_close_telemetry_client_noop(self, telemetry_system_reset):
         client.close.assert_called_once()
 
         client = get_telemetry_client(session_id_hex)
-        assert client is NOOP_TELEMETRY_CLIENT
+        assert isinstance(client, NoopTelemetryClient)
 
     @patch("databricks.sql.telemetry.telemetry_client._handle_unhandled_exception")
     def test_global_exception_hook(self, mock_handle_exception, telemetry_system_reset):