changed export_event and flush to private functions

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -467,7 +467,7 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
-        self._telemetry_client.close()
+        TelemetryClientFactory.close(self.get_session_id_hex())
 
     def commit(self):
         """No-op because Databricks does not support transactions"""

src/databricks/sql/telemetry/telemetry_client.py

@@ -102,14 +102,6 @@ class BaseTelemetryClient(ABC):
     It is used to define the interface for telemetry clients.
     """
 
-    @abstractmethod
-    def export_event(self, event):
-        pass
-
-    @abstractmethod
-    def flush(self):
-        pass
-
     @abstractmethod
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
         pass
@@ -136,12 +128,6 @@ def __new__(cls):
             cls._instance = super(NoopTelemetryClient, cls).__new__(cls)
         return cls._instance
 
-    def export_event(self, event):
-        pass
-
-    def flush(self):
-        pass
-
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
         pass
 
@@ -182,7 +168,7 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
 
-    def export_event(self, event):
+    def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
         logger.debug("Exporting event for connection %s", self._connection_uuid)
         with self._lock:
@@ -191,9 +177,9 @@ def export_event(self, event):
             logger.debug(
                 "Batch size limit reached (%s), flushing events", self._batch_size
             )
-            self.flush()
+            self._flush()
 
-    def flush(self):
+    def _flush(self):
         """Flush the current batch of events to the server"""
         with self._lock:
             events_to_flush = self._events_batch.copy()
@@ -313,11 +299,7 @@ def export_failure_log(self, error_name, error_message):
     def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._connection_uuid)
-        try:
-            self.flush()
-            TelemetryClientFactory.close(self._connection_uuid)
-        except Exception as e:
-            logger.debug("Failed to close telemetry client: %s", e)
+        self._flush()
 
 
 class TelemetryClientFactory:
@@ -365,20 +347,8 @@ def _handle_unhandled_exception(cls, exc_type, exc_value, exc_traceback):
         logger.debug("Handling unhandled exception: %s", exc_type.__name__)
 
         # Flush existing thread pool work and wait for completion
-        logger.debug(
-            "Flushing pending telemetry and waiting for thread pool completion..."
-        )
-        for client in cls._clients.items():
-            client.flush()  # Submit any pending events
-
-        if cls._executor:
-            try:
-                cls._executor.shutdown(
-                    wait=True
-                )  # This waits for all submitted work to complete
-                logger.debug("Thread pool shutdown completed successfully")
-            except Exception as e:
-                logger.debug("Thread pool shutdown failed: %s", e)
+        for uuid, _ in cls._clients.items():
+            cls.close(uuid)
 
         # Call the original exception handler to maintain normal behavior
         if cls._original_excepthook:
@@ -445,6 +415,7 @@ def close(connection_uuid):
                 logger.debug(
                     "Removing telemetry client for connection %s", connection_uuid
                 )
+                TelemetryClientFactory.get_telemetry_client(connection_uuid).close()
                 TelemetryClientFactory._clients.pop(connection_uuid, None)
 
             # Shutdown executor if no more clients
@@ -455,3 +426,4 @@ def close(connection_uuid):
                 TelemetryClientFactory._executor.shutdown(wait=True)
                 TelemetryClientFactory._executor = None
                 TelemetryClientFactory._initialized = False
+                

tests/unit/test_telemetry.py

@@ -177,18 +177,18 @@ def test_export_failure_log(
     def test_export_event(self, telemetry_client_setup):
         """Test exporting an event."""
         client = telemetry_client_setup["client"]
-        client.flush = MagicMock()
+        client._flush = MagicMock()
 
         for i in range(5):
-            client.export_event(f"event-{i}")
+            client._export_event(f"event-{i}")
 
-        client.flush.assert_not_called()
+        client._flush.assert_not_called()
         assert len(client._events_batch) == 5
 
         for i in range(5, 10):
-            client.export_event(f"event-{i}")
+            client._export_event(f"event-{i}")
 
-        client.flush.assert_called_once()
+        client._flush.assert_called_once()
         assert len(client._events_batch) == 10
 
     @patch("requests.post")
@@ -244,7 +244,7 @@ def test_flush(self, telemetry_client_setup):
         client._events_batch = ["event1", "event2"]
         client._send_telemetry = MagicMock()
 
-        client.flush()
+        client._flush()
 
         client._send_telemetry.assert_called_once_with(["event1", "event2"])
         assert client._events_batch == []
@@ -253,13 +253,11 @@ def test_flush(self, telemetry_client_setup):
     def test_close(self, mock_factory_class, telemetry_client_setup):
         """Test closing the client."""
         client = telemetry_client_setup["client"]
-        connection_uuid = telemetry_client_setup["connection_uuid"]
-        client.flush = MagicMock()
+        client._flush = MagicMock()
 
         client.close()
 
-        client.flush.assert_called_once()
-        mock_factory_class.close.assert_called_once_with(connection_uuid)
+        client._flush.assert_called_once()
 
 
 class TestTelemetryClientFactory: