telemetry client changes

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -23,7 +23,6 @@
 from abc import ABC, abstractmethod
 
 
-<<<<<<< HEAD
 class BaseTelemetryClient(ABC):
     @abstractmethod
     def export_initial_telemetry_log(self, **kwargs):
@@ -53,29 +52,13 @@ 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.auth_provider = auth_provider
-        self.is_authenticated = is_authenticated
-        self.batch_size = batch_size
-        self.user_agent = user_agent
+        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.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):
@@ -92,67 +75,18 @@ 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 _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, **kwargs):
+        http_path = kwargs.get("http_path", None)
+        port = kwargs.get("port", None)
+        socket_timeout = kwargs.get("socket_timeout", None)
 
-    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"
@@ -189,26 +123,10 @@ def export_initial_telemetry_log(self, http_path, port, socket_timeout):
 
         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:
@@ -227,7 +145,6 @@ 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
@@ -269,59 +186,6 @@ 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:
@@ -343,7 +207,6 @@ 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
@@ -353,14 +216,3 @@ 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)