added doc strings to latency_logger, abstracted export_telemetry_log

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

Author: saishreeeee

src/databricks/sql/telemetry/latency_logger.py

@@ -11,10 +11,32 @@
 
 
 class TelemetryExtractor:
+    """
+    Base class for extracting telemetry information from various object types.
+
+    This class serves as a proxy that delegates attribute access to the wrapped object
+    while providing a common interface for extracting telemetry-related data.
+    """
+
     def __init__(self, obj):
+        """
+        Initialize the extractor with an object to wrap.
+
+        Args:
+            obj: The object to extract telemetry information from.
+        """
         self._obj = obj
 
     def __getattr__(self, name):
+        """
+        Delegate attribute access to the wrapped object.
+
+        Args:
+            name (str): The name of the attribute to access.
+
+        Returns:
+            The attribute value from the wrapped object.
+        """
         return getattr(self._obj, name)
 
     def get_session_id_hex(self):
@@ -37,6 +59,13 @@ def get_retry_count(self):
 
 
 class CursorExtractor(TelemetryExtractor):
+    """
+    Telemetry extractor specialized for Cursor objects.
+
+    Extracts telemetry information from database cursor objects, including
+    statement IDs, session information, compression settings, and result formats.
+    """
+
     def get_statement_id(self) -> Optional[str]:
         return self.query_id
 
@@ -72,6 +101,13 @@ def get_statement_type(self) -> StatementType:
 
 
 class ResultSetExtractor(TelemetryExtractor):
+    """
+    Telemetry extractor specialized for ResultSet objects.
+
+    Extracts telemetry information from database result set objects, including
+    operation IDs, session information, compression settings, and result formats.
+    """
+
     def get_statement_id(self) -> Optional[str]:
         if self.command_id:
             return str(UUID(bytes=self.command_id.operationId.guid))
@@ -106,15 +142,63 @@ def get_retry_count(self) -> int:
 
 
 def get_extractor(obj):
+    """
+    Factory function to create the appropriate telemetry extractor for an object.
+
+    Determines the object type and returns the corresponding specialized extractor
+    that can extract telemetry information from that object type.
+
+    Args:
+        obj: The object to create an extractor for. Can be a Cursor, ResultSet,
+             or any other object.
+
+    Returns:
+        TelemetryExtractor: A specialized extractor instance:
+            - CursorExtractor for Cursor objects
+            - ResultSetExtractor for ResultSet objects
+            - TelemetryExtractor (base) for all other objects
+    """
     if obj.__class__.__name__ == "Cursor":
         return CursorExtractor(obj)
     elif obj.__class__.__name__ == "ResultSet":
         return ResultSetExtractor(obj)
     else:
-        return TelemetryExtractor(obj)
+        raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")
 
 
 def log_latency():
+    """
+    Decorator for logging execution latency and telemetry information.
+
+    This decorator measures the execution time of a method and sends telemetry
+    data about the operation, including latency, statement information, and
+    execution context.
+
+    The decorator automatically:
+    - Measures execution time using high-precision performance counters
+    - Extracts telemetry information from the method's object (self)
+    - Creates a SqlExecutionEvent with execution details
+    - Sends the telemetry data asynchronously via TelemetryClient
+
+    Usage:
+        @log_latency()
+        def execute(self, query):
+            # Method implementation
+            pass
+
+        @log_latency()
+        def fetchall(self):
+            # Method implementation
+            pass
+
+    Returns:
+        function: A decorator that wraps methods to add latency logging.
+
+    Note:
+        The wrapped method's object (self) must be compatible with the
+        telemetry extractor system (e.g., Cursor or ResultSet objects).
+    """
+
     def decorator(func):
         @functools.wraps(func)
         def wrapper(self, *args, **kwargs):

src/databricks/sql/telemetry/telemetry_client.py

@@ -248,14 +248,24 @@ def _telemetry_request_callback(self, future):
         except Exception as e:
             logger.debug("Telemetry request failed with exception: %s", e)
 
-    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
-        logger.debug(
-            "Exporting initial telemetry log for connection %s", self._session_id_hex
-        )
+    def _export_telemetry_log(self, **telemetry_event_kwargs):
+        """
+        Common helper method for exporting telemetry logs.
+
+        Args:
+            **telemetry_event_kwargs: Keyword arguments to pass to TelemetryEvent constructor
+        """
+        logger.debug("Exporting telemetry log for connection %s", self._session_id_hex)
 
         try:
-            self._driver_connection_params = driver_connection_params
-            self._user_agent = user_agent
+            # Set common fields for all telemetry events
+            event_kwargs = {
+                "session_id": self._session_id_hex,
+                "system_configuration": TelemetryHelper.get_driver_system_configuration(),
+                "driver_connection_params": self._driver_connection_params,
+            }
+            # Add any additional fields passed in
+            event_kwargs.update(telemetry_event_kwargs)
 
             telemetry_frontend_log = TelemetryFrontendLog(
                 frontend_log_event_id=str(uuid.uuid4()),
@@ -265,72 +275,29 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
                         user_agent=self._user_agent,
                     )
                 ),
-                entry=FrontendLogEntry(
-                    sql_driver_log=TelemetryEvent(
-                        session_id=self._session_id_hex,
-                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
-                        driver_connection_params=self._driver_connection_params,
-                    )
-                ),
+                entry=FrontendLogEntry(sql_driver_log=TelemetryEvent(**event_kwargs)),
             )
 
             self._export_event(telemetry_frontend_log)
 
         except Exception as e:
-            logger.debug("Failed to export initial telemetry log: %s", e)
+            logger.debug("Failed to export telemetry log: %s", e)
+
+    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
+        self._driver_connection_params = driver_connection_params
+        self._user_agent = user_agent
+        self._export_telemetry_log()
 
     def export_failure_log(self, error_name, error_message):
-        logger.debug("Exporting failure log for connection %s", self._session_id_hex)
-        try:
-            error_info = DriverErrorInfo(
-                error_name=error_name, stack_trace=error_message
-            )
-            telemetry_frontend_log = TelemetryFrontendLog(
-                frontend_log_event_id=str(uuid.uuid4()),
-                context=FrontendLogContext(
-                    client_context=TelemetryClientContext(
-                        timestamp_millis=int(time.time() * 1000),
-                        user_agent=self._user_agent,
-                    )
-                ),
-                entry=FrontendLogEntry(
-                    sql_driver_log=TelemetryEvent(
-                        session_id=self._session_id_hex,
-                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
-                        driver_connection_params=self._driver_connection_params,
-                        error_info=error_info,
-                    )
-                ),
-            )
-            self._export_event(telemetry_frontend_log)
-        except Exception as e:
-            logger.debug("Failed to export failure log: %s", e)
+        error_info = DriverErrorInfo(error_name=error_name, stack_trace=error_message)
+        self._export_telemetry_log(error_info=error_info)
 
     def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
-        logger.debug("Exporting latency log for connection %s", self._session_id_hex)
-        try:
-            telemetry_frontend_log = TelemetryFrontendLog(
-                frontend_log_event_id=str(uuid.uuid4()),
-                context=FrontendLogContext(
-                    client_context=TelemetryClientContext(
-                        timestamp_millis=int(time.time() * 1000),
-                        user_agent=self._user_agent,
-                    )
-                ),
-                entry=FrontendLogEntry(
-                    sql_driver_log=TelemetryEvent(
-                        session_id=self._session_id_hex,
-                        system_configuration=TelemetryHelper.get_driver_system_configuration(),
-                        driver_connection_params=self._driver_connection_params,
-                        sql_statement_id=sql_statement_id,
-                        sql_operation=sql_execution_event,
-                        operation_latency_ms=latency_ms,
-                    )
-                ),
-            )
-            self._export_event(telemetry_frontend_log)
-        except Exception as e:
-            logger.debug("Failed to export latency log: %s", e)
+        self._export_telemetry_log(
+            sql_statement_id=sql_statement_id,
+            sql_operation=sql_execution_event,
+            operation_latency_ms=latency_ms,
+        )
 
     def close(self):
         """Flush remaining events before closing"""