Revert "shifted thread pool executor to telemetry manager"

This reverts commit 419dac3302f67edd00c1b1a8254947f6e55f74b8.

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -23,6 +23,7 @@
 from abc import ABC, abstractmethod
 
 
+<<<<<<< HEAD
 class BaseTelemetryClient(ABC):
     @abstractmethod
     def export_initial_telemetry_log(self, **kwargs):
@@ -52,13 +53,29 @@ class TelemetryClient(BaseTelemetryClient):
     def __init__(self, telemetry_enabled, batch_size, connection_uuid, **kwargs):
         self.telemetry_enabled = telemetry_enabled
         self.batch_size = batch_size
+=======
+class TelemetryClient:
+    def __init__(
+        self,
+        host,
+        connection_uuid,
+        batch_size,
+        auth_provider=None,
+        is_authenticated=False,
+        user_agent=None,
+    ):
+        self.host_url = host
+>>>>>>> parent of 419dac3 (shifted thread pool executor to telemetry manager)
         self.connection_uuid = connection_uuid
-        self.host_url = kwargs.get("host_url", None)
-        self.auth_provider = kwargs.get("auth_provider", None)
-        self.is_authenticated = kwargs.get("is_authenticated", False)
-        self.user_agent = kwargs.get("user_agent", None)
+        self.auth_provider = auth_provider
+        self.is_authenticated = is_authenticated
+        self.batch_size = batch_size
+        self.user_agent = user_agent
         self.events_batch = []
         self.lock = threading.Lock()
+        self.executor = ThreadPoolExecutor(
+            max_workers=10  # TODO: Decide on max workers
+        )  # Thread pool for async operations
         self.DriverConnectionParameters = None
 
     def export_event(self, event):
@@ -75,18 +92,67 @@ def flush(self):
             self.events_batch = []
 
         if events_to_flush:
+<<<<<<< HEAD
             telemetry_client_factory._send_telemetry(
                 events_to_flush,
                 self.host_url,
                 self.is_authenticated,
                 self.auth_provider,
             )
+=======
+            self.executor.submit(self._send_telemetry, events_to_flush)
 
-    def export_initial_telemetry_log(self, **kwargs):
-        http_path = kwargs.get("http_path", None)
-        port = kwargs.get("port", None)
-        socket_timeout = kwargs.get("socket_timeout", None)
+    def _send_telemetry(self, events):
+        """Send telemetry events to the server"""
+        request = {
+            "uploadTime": int(time.time() * 1000),
+            "items": [],
+            "protoLogs": [event.to_json() for event in events],
+        }
+
+        path = "/telemetry-ext" if self.is_authenticated else "/telemetry-unauth"
+        url = f"https://{self.host_url}{path}"
+
+        headers = {"Accept": "application/json", "Content-Type": "application/json"}
 
+        if self.is_authenticated and self.auth_provider:
+            self.auth_provider.add_headers(headers)
+
+        # print("\n=== Request Details ===", flush=True)
+        # print(f"URL: {url}", flush=True)
+        # print("\nHeaders:", flush=True)
+        # for key, value in headers.items():
+        #     print(f"  {key}: {value}", flush=True)
+
+        # print("\nRequest Body:", flush=True)
+        # print(json.dumps(request, indent=2), flush=True)
+        # sys.stdout.flush()
+
+        response = requests.post(
+            url, data=json.dumps(request), headers=headers, timeout=10
+        )
+
+        # print("\n=== Response Details ===", flush=True)
+        # print(f"Status Code: {response.status_code}", flush=True)
+        # print("\nResponse Headers:", flush=True)
+        # for key, value in response.headers.items():
+        #     print(f"  {key}: {value}", flush=True)
+
+        # print("\nResponse Body:", flush=True)
+        # try:
+        #     response_json = response.json()
+        #     print(json.dumps(response_json, indent=2), flush=True)
+        # except json.JSONDecodeError:
+        #     print(response.text, flush=True)
+        # sys.stdout.flush()
+
+    def close(self):
+        """Flush remaining events and shut down executor"""
+        self.flush()
+        self.executor.shutdown(wait=True)
+>>>>>>> parent of 419dac3 (shifted thread pool executor to telemetry manager)
+
+    def export_initial_telemetry_log(self, http_path, port, socket_timeout):
         discovery_url = None
         if hasattr(self.auth_provider, "oauth_manager") and hasattr(
             self.auth_provider.oauth_manager, "idp_endpoint"
@@ -123,10 +189,26 @@ def export_initial_telemetry_log(self, **kwargs):
 
         self.export_event(telemetry_frontend_log)
 
+<<<<<<< HEAD
     def close(self):
         """Flush remaining events before closing"""
         self.flush()
         telemetry_client_factory.close(self.connection_uuid)
+=======
+    def export_failure_log(self, errorName, errorMessage):
+        pass
+
+    def export_sql_latency_log(
+        self, latency_ms, sql_execution_event, sql_statement_id=None
+    ):
+        """Export telemetry for sql execution"""
+        pass
+
+    def export_volume_latency_log(self, latency_ms, volume_operation):
+        """Export telemetry for volume operation"""
+        pass
+
+>>>>>>> parent of 419dac3 (shifted thread pool executor to telemetry manager)
 
 
 class TelemetryClientFactory:
@@ -145,6 +227,7 @@ def __init__(self):
             return
 
         self._clients = {}  # Map of connection_uuid -> TelemetryClient
+<<<<<<< HEAD
         self.executor = ThreadPoolExecutor(
             max_workers=10
         )  # Thread pool for async operations TODO: Decide on max workers
@@ -186,6 +269,59 @@ def _send_telemetry(self, events, host_url, is_authenticated, auth_provider):
             requests.post, url, data=json.dumps(request), headers=headers, timeout=10
         )
 
+=======
+        self._initialized = True
+
+    def initialize(
+        self,
+        host,
+        connection_uuid,
+        batch_size,
+        auth_provider=None,
+        is_authenticated=False,
+        user_agent=None,
+    ):
+        """Initialize a telemetry client for a specific connection"""
+        if connection_uuid not in self._clients:
+            self._clients[connection_uuid] = TelemetryClient(
+                host=host,
+                connection_uuid=connection_uuid,
+                batch_size=batch_size,
+                auth_provider=auth_provider,
+                is_authenticated=is_authenticated,
+                user_agent=user_agent,
+            )
+
+    def export_failure_log(self, error_name, error_message, connection_uuid):
+        """Export error logs for a specific connection or all connections if connection_uuid is None"""
+        pass
+
+    def export_initial_telemetry_log(
+        self, http_path, port, socket_timeout, connection_uuid
+    ):
+        """Export initial telemetry for a specific connection"""
+        if connection_uuid in self._clients:
+            self._clients[connection_uuid].export_initial_telemetry_log(
+                http_path, port, socket_timeout
+            )
+
+    def export_sql_latency_log(
+        self,
+        latency_ms,
+        sql_execution_event,
+        sql_statement_id=None,
+        connection_uuid=None,
+    ):
+        """Export latency logs for sql execution for a specific connection"""
+        pass
+
+    def export_volume_latency_log(
+        self, latency_ms, volume_operation, connection_uuid=None
+    ):
+        """Export latency logs for volume operation for a specific connection"""
+        pass
+
+>>>>>>> parent of 419dac3 (shifted thread pool executor to telemetry manager)
     @classmethod
     def getDriverSystemConfiguration(cls) -> DriverSystemConfiguration:
         if cls._DRIVER_SYSTEM_CONFIGURATION is None:
@@ -207,6 +343,7 @@ def getDriverSystemConfiguration(cls) -> DriverSystemConfiguration:
         return cls._DRIVER_SYSTEM_CONFIGURATION
 
     def close(self, connection_uuid):
+<<<<<<< HEAD
         del self._clients[connection_uuid]
 
         # Shutdown executor if no more clients
@@ -216,3 +353,14 @@ def close(self, connection_uuid):
 
 # Create a global instance
 telemetry_client_factory = TelemetryClientFactory()
+=======
+        """Close telemetry client(s)"""
+        if connection_uuid:
+            if connection_uuid in self._clients:
+                self._clients[connection_uuid].close()
+                del self._clients[connection_uuid]
+
+
+# Create a global instance
+telemetry_client = TelemetryManager()
+>>>>>>> parent of 419dac3 (shifted thread pool executor to telemetry manager)