added classes required for telemetry

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

Author: saishreeeee

examples/query_execute.py

@@ -8,7 +8,7 @@
 ) as connection:
 
     with connection.cursor() as cursor:
-        cursor.execute("SELECT * FROM default.diamonds LIMIT 2")
+        cursor.execute("SELECT * FROM main.eng_lumberjack.staging_frontend_log_sql_driver_log limit 1")
         result = cursor.fetchall()
 
         for row in result:

examples/test_telemetry.py

@@ -0,0 +1,20 @@
+import os
+import databricks.sql as sql
+
+# Create connection with telemetry enabled
+conn = sql.connect(
+    server_hostname=os.environ["DATABRICKS_SERVER_HOSTNAME"],
+    http_path=os.environ["DATABRICKS_HTTP_PATH"],
+    access_token=os.environ["DATABRICKS_TOKEN"],
+    enable_telemetry=True,  # Enable telemetry
+    telemetry_batch_size=1  # Set batch size to 1
+)
+
+# Execute a simple query to generate telemetry
+cursor = conn.cursor()
+cursor.execute("SELECT * FROM main.eng_lumberjack.staging_frontend_log_sql_driver_log limit 1")
+cursor.fetchall()
+
+# Close the connection
+cursor.close()
+conn.close() 

src/databricks/sql/client.py

@@ -1,5 +1,21 @@
 import time
 from typing import Dict, Tuple, List, Optional, Any, Union, Sequence
+import uuid
+from databricks.sql.telemetry.TelemetryEvent import TelemetryEvent
+from databricks.sql.telemetry.DriverSystemConfiguration import DriverSystemConfiguration
+from databricks.sql.telemetry.TelemetryClient import TelemetryClient
+from databricks.sql.telemetry.NoopTelemetryClient import NoopTelemetryClient
+from databricks.sql.telemetry.TelemetryFrontendLog import TelemetryFrontendLog
+from databricks.sql.telemetry.FrontendLogContext import FrontendLogContext
+from databricks.sql.telemetry.TelemetryClientContext import TelemetryClientContext
+from databricks.sql.telemetry.FrontendLogEntry import FrontendLogEntry
+from databricks.sql.auth.auth import AuthType
+from databricks.sql.auth.authenticators import (
+    DatabricksOAuthProvider,
+    ExternalAuthProvider,
+    AuthProvider,
+    AccessTokenAuthProvider,
+)
 
 import pandas
 
@@ -234,6 +250,32 @@ def read(self) -> Optional[OAuthToken]:
             server_hostname, **kwargs
         )
 
+        self.server_telemetry_enabled = True
+        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
+        self.telemetry_enabled = (
+            self.client_telemetry_enabled and self.server_telemetry_enabled
+        )
+        telemetry_batch_size = kwargs.get("telemetry_batch_size", 200)
+
+        if self.telemetry_enabled:
+            self.telemetry_client = TelemetryClient(
+                host=self.host,
+                connection_uuid="test-connection-uuid",
+                auth_provider=auth_provider,
+                is_authenticated=(
+                    isinstance(auth_provider, AccessTokenAuthProvider)
+                    or isinstance(auth_provider, DatabricksOAuthProvider)
+                    or isinstance(auth_provider, ExternalAuthProvider)
+                    or (
+                        isinstance(auth_provider, AuthProvider)
+                        and hasattr(auth_provider, "_header_factory")
+                    )
+                ),
+                batch_size=telemetry_batch_size,
+            )
+        else:
+            self.telemetry_client = NoopTelemetryClient()
+
         user_agent_entry = kwargs.get("user_agent_entry")
         if user_agent_entry is None:
             user_agent_entry = kwargs.get("_user_agent_entry")
@@ -419,6 +461,9 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
+        if hasattr(self, "telemetry_client"):
+            self.telemetry_client.close()
+
     def commit(self):
         """No-op because Databricks does not support transactions"""
         pass

src/databricks/sql/telemetry/DriverConnectionParameters.py

@@ -0,0 +1,24 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.HostDetails import HostDetails
+from databricks.sql.telemetry.enums.AuthMech import AuthMech
+from databricks.sql.telemetry.enums.AuthFlow import AuthFlow
+from databricks.sql.telemetry.enums.DatabricksClientType import DatabricksClientType
+
+
+@dataclass
+class DriverConnectionParameters:
+    http_path: str
+    driver_mode: DatabricksClientType
+    host_details: HostDetails
+    auth_mech: AuthMech
+    auth_flow: AuthFlow
+    auth_scope: str
+    discovery_url: str
+    allowed_volume_ingestion_paths: str
+    enable_complex_datatype_support: bool
+    azure_tenant_id: str
+    socket_timeout: int
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/DriverErrorInfo.py

@@ -0,0 +1,11 @@
+import json
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class DriverErrorInfo:
+    error_name: str
+    stack_trace: str
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/DriverSystemConfiguration.py

@@ -0,0 +1,37 @@
+import json
+from dataclasses import dataclass, asdict
+import platform
+import sys
+import locale
+from databricks.sql import __version__
+
+
+@dataclass
+class DriverSystemConfiguration:
+    driver_version: str
+    os_name: str
+    os_version: str
+    os_arch: str
+    runtime_name: str
+    runtime_version: str
+    runtime_vendor: str
+    client_app_name: str
+    locale_name: str
+    driver_name: str
+    char_set_encoding: str
+
+    def __init__(self):
+        self.driver_version = __version__
+        self.os_name = platform.system()
+        self.os_version = platform.version()
+        self.os_arch = platform.machine()
+        self.runtime_name = platform.python_implementation()
+        self.runtime_version = platform.python_version()
+        self.runtime_vendor = sys.implementation.name
+        self.client_app_name = "databricks-sql-python"
+        self.locale_name = locale.getdefaultlocale()[0]
+        self.driver_name = "databricks-sql-python"
+        self.char_set_encoding = "UTF-8"
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/DriverVolumeOperation.py

@@ -0,0 +1,14 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.enums.DriverVolumeOperationType import (
+    DriverVolumeOperationType,
+)
+
+
+@dataclass
+class DriverVolumeOperation:
+    volume_operation_type: DriverVolumeOperationType
+    volume_path: str
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/FrontendLogContext.py

@@ -0,0 +1,11 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.TelemetryClientContext import TelemetryClientContext
+
+
+@dataclass
+class FrontendLogContext:
+    client_context: TelemetryClientContext
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/FrontendLogEntry.py

@@ -0,0 +1,11 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.TelemetryEvent import TelemetryEvent
+
+
+@dataclass
+class FrontendLogEntry:
+    sql_driver_log: TelemetryEvent
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/HostDetails.py

@@ -0,0 +1,11 @@
+import json
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class HostDetails:
+    host_url: str
+    port: int
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/NoopTelemetryClient.py

@@ -0,0 +1,11 @@
+class NoopTelemetryClient:
+    # A no-operation telemetry client that implements the same interface but does nothing
+
+    def export_event(self, event):
+        pass
+
+    def flush(self):
+        pass
+
+    def close(self):
+        pass

src/databricks/sql/telemetry/SqlExecutionEvent.py

@@ -0,0 +1,15 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.enums.StatementType import StatementType
+from databricks.sql.telemetry.enums.ExecutionResultFormat import ExecutionResultFormat
+
+
+@dataclass
+class SqlExecutionEvent:
+    statement_type: StatementType
+    is_compressed: bool
+    execution_result: ExecutionResultFormat
+    retry_count: int
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/TelemetryClient.py

@@ -0,0 +1,36 @@
+import threading
+import time
+import json
+import requests
+from concurrent.futures import ThreadPoolExecutor
+
+
+class TelemetryClient:
+    def __init__(
+        self,
+        host,
+        connection_uuid,
+        auth_provider=None,
+        is_authenticated=False,
+        batch_size=200,
+    ):
+        self.host = host
+        self.connection_uuid = connection_uuid
+        self.auth_provider = auth_provider
+        self.is_authenticated = is_authenticated
+        self.batch_size = batch_size
+        self.events_batch = []
+        self.lock = threading.Lock()
+        self.executor = ThreadPoolExecutor(
+            max_workers=5
+        )  # Thread pool for async operations
+        self.DriverConnectionParameters = None
+
+    def export_event(self, event):
+        pass
+
+    def flush(self):
+        pass
+
+    def close(self):
+        pass

src/databricks/sql/telemetry/TelemetryClientContext.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, asdict
+import json
+
+
+@dataclass
+class TelemetryClientContext:
+    timestamp_millis: int
+    user_agent: str
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/TelemetryEvent.py

@@ -0,0 +1,25 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.DriverSystemConfiguration import DriverSystemConfiguration
+from databricks.sql.telemetry.DriverConnectionParameters import (
+    DriverConnectionParameters,
+)
+from databricks.sql.telemetry.DriverVolumeOperation import DriverVolumeOperation
+from databricks.sql.telemetry.SqlExecutionEvent import SqlExecutionEvent
+from databricks.sql.telemetry.DriverErrorInfo import DriverErrorInfo
+
+
+@dataclass
+class TelemetryEvent:
+    session_id: str
+    sql_statement_id: str
+    system_configuration: DriverSystemConfiguration
+    driver_connection_params: DriverConnectionParameters
+    auth_type: str
+    vol_operation: DriverVolumeOperation
+    sql_operation: SqlExecutionEvent
+    error_info: DriverErrorInfo
+    latency: int
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/TelemetryFrontendLog.py

@@ -0,0 +1,15 @@
+import json
+from dataclasses import dataclass, asdict
+from databricks.sql.telemetry.FrontendLogContext import FrontendLogContext
+from databricks.sql.telemetry.FrontendLogEntry import FrontendLogEntry
+
+
+@dataclass
+class TelemetryFrontendLog:
+    workspace_id: int
+    frontend_log_event_id: str
+    context: FrontendLogContext
+    entry: FrontendLogEntry
+
+    def to_json(self):
+        return json.dumps(asdict(self))

src/databricks/sql/telemetry/TelemetryHelper.py

@@ -0,0 +1,32 @@
+import platform
+import sys
+import uuid
+import time
+from typing import Optional
+
+from databricks.sql import __version__
+from databricks.sql.telemetry.DriverSystemConfiguration import DriverSystemConfiguration
+
+
+class TelemetryHelper:
+
+    # Singleton instance of DriverSystemConfiguration
+    _DRIVER_SYSTEM_CONFIGURATION = None
+
+    @classmethod
+    def getDriverSystemConfiguration(cls) -> DriverSystemConfiguration:
+        if cls._DRIVER_SYSTEM_CONFIGURATION is None:
+            cls._DRIVER_SYSTEM_CONFIGURATION = DriverSystemConfiguration(
+                driverName="Databricks SQL Python Connector",
+                driverVersion=__version__,
+                runtimeName=f"Python {sys.version.split()[0]}",
+                runtimeVendor=platform.python_implementation(),
+                runtimeVersion=platform.python_version(),
+                osName=platform.system(),
+                osVersion=platform.release(),
+                osArch=platform.machine(),
+                clientAppName=None,
+                localeName=f"{platform.system()}_{platform.release()}",
+                charSetEncoding=sys.getdefaultencoding(),
+            )
+        return cls._DRIVER_SYSTEM_CONFIGURATION

src/databricks/sql/telemetry/TelemetryRequest.py

@@ -0,0 +1,13 @@
+import json
+from dataclasses import dataclass, asdict
+from typing import List, Optional
+
+
+@dataclass
+class TelemetryRequest:
+    uploadTime: int
+    items: List[str]
+    protoLogs: Optional[List[str]]
+
+    def to_json(self):
+        return json.dumps(asdict(self))