unauth telemetry client

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -250,13 +250,14 @@ def read(self) -> Optional[OAuthToken]:
                 server_hostname, **kwargs
             )
         except Exception as e:
-            raise AuthenticationError(
-                message=f"Failed to create authentication provider: {str(e)}",
-                host_url=server_hostname,
+            TelemetryClientFactory.connection_failure_log(
+                error_name=e.__class__.__name__,
+                error_message=str(e),
+                host_url=self.host,
                 http_path=http_path,
                 port=self.port,
-                original_exception=e,
-            ) from e
+            )
+            raise e
 
         self.server_telemetry_enabled = True
         self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
@@ -316,14 +317,15 @@ def read(self) -> Optional[OAuthToken]:
                 session_configuration, catalog, schema
             )
         except Exception as e:
-            raise ConnectionError(
-                message=f"Failed to establish connection: {str(e)}",
+            TelemetryClientFactory.connection_failure_log(
+                error_name=e.__class__.__name__,
+                error_message=str(e),
                 host_url=self.host,
                 http_path=http_path,
                 port=self.port,
                 user_agent=useragent_header,
-                original_exception=e,
-            ) from e
+            )
+            raise e
 
         self._session_handle = self._open_session_resp.sessionHandle
         self.protocol_version = self.get_protocol_version(self._open_session_resp)
@@ -1418,7 +1420,6 @@ def _fill_results_buffer(self):
         self.results = results
         self.has_more_rows = has_more_rows
 
-    @log_latency()
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
         ResultRow = Row(*column_names)
@@ -1431,7 +1432,6 @@ def _convert_columnar_table(self, table):
 
         return result
 
-    @log_latency()
     def _convert_arrow_table(self, table):
         column_names = [c[0] for c in self.description]
         ResultRow = Row(*column_names)
@@ -1474,7 +1474,6 @@ def _convert_arrow_table(self, table):
     def rownumber(self):
         return self._next_row_index
 
-    @log_latency()
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
@@ -1517,7 +1516,6 @@ def merge_columnar(self, result1, result2):
         ]
         return ColumnTable(merged_result, result1.column_names)
 
-    @log_latency()
     def fetchmany_columnar(self, size: int):
         """
         Fetch the next set of rows of a query result, returning a Columnar Table.
@@ -1543,7 +1541,6 @@ def fetchmany_columnar(self, size: int):
 
         return results
 
-    @log_latency()
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         results = self.results.remaining_rows()
@@ -1570,7 +1567,6 @@ def fetchall_arrow(self) -> "pyarrow.Table":
             return pyarrow.Table.from_pydict(data)
         return results
 
-    @log_latency()
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
         results = self.results.remaining_rows()

src/databricks/sql/exc.py

@@ -1,9 +1,9 @@
 import json
 import logging
 
+logger = logging.getLogger(__name__)
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
-logger = logging.getLogger(__name__)
 
 ### PEP-249 Mandated ###
 # https://peps.python.org/pep-0249/#exceptions
@@ -22,29 +22,10 @@ def __init__(
 
         error_name = self.__class__.__name__
         if session_id_hex:
-            # Normal case: we have a session, send to regular telemetry client
             telemetry_client = TelemetryClientFactory.get_telemetry_client(
                 session_id_hex
             )
             telemetry_client.export_failure_log(error_name, self.message)
-        elif (
-            isinstance(self, (ConnectionError, AuthenticationError))
-            and "host_url" in self.context
-        ):
-            # Connection error case: no session but we should still send telemetry
-            self._send_connection_error_telemetry(error_name)
-
-    def _send_connection_error_telemetry(self, error_name):
-        """Send connection error telemetry to unauthenticated endpoint"""
-
-        TelemetryClientFactory.send_connection_error_telemetry(
-            error_name=error_name,
-            error_message=self.message or str(self),
-            host_url=self.context["host_url"],
-            http_path=self.context.get("http_path", ""),
-            port=self.context.get("port", 443),
-            user_agent=self.context.get("user_agent"),
-        )
 
     def __str__(self):
         return self.message

src/databricks/sql/telemetry/telemetry_client.py

@@ -440,74 +440,33 @@ def close(session_id_hex):
                 TelemetryClientFactory._initialized = False
 
     @staticmethod
-    def send_connection_error_telemetry(
+    def connection_failure_log(
         error_name: str,
         error_message: str,
         host_url: str,
         http_path: str,
-        port: int = 443,
+        port: int,
         user_agent: Optional[str] = None,
     ):
         """Send error telemetry when connection creation fails, without requiring a session"""
-        try:
-            logger.debug("Sending connection error telemetry for host: %s", host_url)
-
-            # Initialize factory if needed (with proper locking)
-            with TelemetryClientFactory._lock:
-                TelemetryClientFactory._initialize()
 
-            # Create driver connection params for the failed connection
-            driver_connection_params = DriverConnectionParameters(
-                http_path=http_path,
-                mode=DatabricksClientType.THRIFT,
-                host_info=HostDetails(host_url=host_url, port=port),
-            )
-
-            error_info = DriverErrorInfo(
-                error_name=error_name, stack_trace=error_message
-            )
+        UNAUTH_SESSION_ID = "Unauth_session_id"
 
-            telemetry_frontend_log = TelemetryFrontendLog(
-                frontend_log_event_id=str(uuid.uuid4()),
-                context=FrontendLogContext(
-                    client_context=TelemetryClientContext(
-                        timestamp_millis=int(time.time() * 1000),
-                        user_agent=user_agent or "PyDatabricksSqlConnector",
-                    )
-                ),
-                entry=FrontendLogEntry(
-                    sql_driver_log=TelemetryEvent(
-                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
-                        driver_connection_params=driver_connection_params,
-                        error_info=error_info,
-                    )
-                ),
-            )
-
-            # Send to unauthenticated endpoint since we don't have working auth
-            request = {
-                "uploadTime": int(time.time() * 1000),
-                "items": [],
-                "protoLogs": [telemetry_frontend_log.to_json()],
-            }
-
-            url = f"https://{host_url}/telemetry-unauth"
-            headers = {"Accept": "application/json", "Content-Type": "application/json"}
+        TelemetryClientFactory.initialize_telemetry_client(
+            telemetry_enabled=True,
+            session_id_hex=UNAUTH_SESSION_ID,
+            auth_provider=None,
+            host_url=host_url,
+        )
 
-            # Send synchronously for connection errors since we're probably about to exit
-            response = requests.post(
-                url,
-                data=json.dumps(request),
-                headers=headers,
-                timeout=5,
-            )
-            if response.status_code == 200:
-                logger.debug("Connection error telemetry sent successfully")
-            else:
-                logger.debug(
-                    "Connection error telemetry failed with status: %s",
-                    response.status_code,
-                )
+        telemetry_client = TelemetryClientFactory.get_telemetry_client(
+            UNAUTH_SESSION_ID
+        )
+        telemetry_client._driver_connection_params = DriverConnectionParameters(
+            http_path=http_path,
+            mode=DatabricksClientType.THRIFT,
+            host_info=HostDetails(host_url=host_url, port=port),
+        )
+        telemetry_client._user_agent = user_agent
 
-        except Exception as e:
-            logger.debug("Failed to send connection error telemetry: %s", e)
+        telemetry_client.export_failure_log(error_name, error_message)