E2E POC for python telemetry for connect logs

src/databricks/sql/client.py

@@ -49,6 +49,15 @@
     TSparkParameter,
     TOperationState,
 )
+from databricks.sql.telemetry.telemetry_client import (
+    TelemetryClientFactory,
+    TelemetryHelper,
+)
+from databricks.sql.telemetry.models.enums import DatabricksClientType
+from databricks.sql.telemetry.models.event import (
+    DriverConnectionParameters,
+    HostDetails,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -294,6 +303,31 @@ def read(self) -> Optional[OAuthToken]:
             kwargs.get("use_inline_params", False)
         )
 
+        TelemetryClientFactory.initialize_telemetry_client(
+            telemetry_enabled=self.telemetry_enabled,
+            connection_uuid=self.get_session_id_hex(),
+            auth_provider=auth_provider,
+            host_url=self.host,
+        )
+
+        self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
+            connection_uuid=self.get_session_id_hex()
+        )
+
+        driver_connection_params = DriverConnectionParameters(
+            http_path=http_path,
+            mode=DatabricksClientType.THRIFT,
+            host_info=HostDetails(host_url=server_hostname, port=self.port),
+            auth_mech=TelemetryHelper.get_auth_mechanism(auth_provider),
+            auth_flow=TelemetryHelper.get_auth_flow(auth_provider),
+            socket_timeout=kwargs.get("_socket_timeout", None),
+        )
+
+        self._telemetry_client.export_initial_telemetry_log(
+            driver_connection_params=driver_connection_params,
+            user_agent=useragent_header,
+        )
+
     def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
         """Valid values are True, False, and "silent"
 
@@ -430,6 +464,8 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
+        self._telemetry_client.close()
+
     def commit(self):
         """No-op because Databricks does not support transactions"""
         pass

src/databricks/sql/telemetry/models/enums.py

@@ -3,15 +3,14 @@
 
 class AuthFlow(Enum):
     TOKEN_PASSTHROUGH = "token_passthrough"
-    CLIENT_CREDENTIALS = "client_credentials"
     BROWSER_BASED_AUTHENTICATION = "browser_based_authentication"
-    AZURE_MANAGED_IDENTITIES = "azure_managed_identities"
 
 
 class AuthMech(Enum):
-    OTHER = "other"
-    PAT = "pat"
-    OAUTH = "oauth"
+    CLIENT_CERT = "CLIENT_CERT"  # ssl certificate authentication
+    PAT = "PAT"  # Personal Access Token authentication
+    DATABRICKS_OAUTH = "DATABRICKS_OAUTH"  # Databricks-managed OAuth flow
+    EXTERNAL_AUTH = "EXTERNAL_AUTH"  # External identity provider (AWS, Azure, etc.)
 
 
 class DatabricksClientType(Enum):

src/databricks/sql/telemetry/models/event.py

@@ -9,6 +9,7 @@
     ExecutionResultFormat,
 )
 from typing import Optional
+from databricks.sql.telemetry.utils import EnumEncoder
 
 
 @dataclass
@@ -40,26 +41,18 @@ class DriverConnectionParameters:
         host_info (HostDetails): Details about the host connection
         auth_mech (AuthMech): The authentication mechanism used
         auth_flow (AuthFlow): The authentication flow type
-        auth_scope (str): The scope of authentication
-        discovery_url (str): URL for service discovery
-        allowed_volume_ingestion_paths (str): JSON string of allowed paths for volume operations
-        azure_tenant_id (str): Azure tenant ID for Azure authentication
         socket_timeout (int): Connection timeout in milliseconds
     """
 
     http_path: str
     mode: DatabricksClientType
     host_info: HostDetails
-    auth_mech: AuthMech
-    auth_flow: AuthFlow
-    auth_scope: str
-    discovery_url: str
-    allowed_volume_ingestion_paths: str
-    azure_tenant_id: str
-    socket_timeout: int
+    auth_mech: Optional[AuthMech] = None
+    auth_flow: Optional[AuthFlow] = None
+    socket_timeout: Optional[int] = None
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -89,13 +82,13 @@ class DriverSystemConfiguration:
     runtime_name: str
     runtime_version: str
     runtime_vendor: str
-    client_app_name: str
-    locale_name: str
     driver_name: str
     char_set_encoding: str
+    client_app_name: Optional[str] = None
+    locale_name: Optional[str] = None
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -113,7 +106,7 @@ class DriverVolumeOperation:
     volume_path: str
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -131,7 +124,7 @@ class DriverErrorInfo:
     stack_trace: str
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -153,7 +146,7 @@ class SqlExecutionEvent:
     retry_count: int
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -186,4 +179,4 @@ class TelemetryEvent:
     operation_latency_ms: Optional[int] = None
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)

src/databricks/sql/telemetry/models/frontend_logs.py

@@ -1,6 +1,7 @@
 import json
 from dataclasses import dataclass, asdict
 from databricks.sql.telemetry.models.event import TelemetryEvent
+from databricks.sql.telemetry.utils import EnumEncoder
 from typing import Optional
 
 
@@ -19,7 +20,7 @@ class TelemetryClientContext:
     user_agent: str
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -35,7 +36,7 @@ class FrontendLogContext:
     client_context: TelemetryClientContext
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -51,7 +52,7 @@ class FrontendLogEntry:
     sql_driver_log: TelemetryEvent
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)
 
 
 @dataclass
@@ -74,4 +75,4 @@ class TelemetryFrontendLog:
     workspace_id: Optional[int] = None
 
     def to_json(self):
-        return json.dumps(asdict(self))
+        return json.dumps(asdict(self), cls=EnumEncoder)