perf: Optimize telemetry latency logging to reduce overhead (databricks#715)

 perf: Optimize telemetry latency logging to reduce overhead

Optimizations implemented:
1. Eliminated extractor pattern - replaced wrapper classes with direct
   attribute access functions, removing object creation overhead
2. Added feature flag early exit - checks cached telemetry_enabled flag
   to skip heavy work when telemetry is disabled
3. Simplified code structure with early returns for better readability


Signed-off-by: Samikshya Chand <samikshya.chand@databricks.com>

Author: samikshya-db

src/databricks/sql/common/feature_flag.py

@@ -165,8 +165,9 @@ def get_instance(cls, connection: "Connection") -> FeatureFlagsContext:
             cls._initialize()
             assert cls._executor is not None
 
-            # Use the unique session ID as the key
-            key = connection.get_session_id_hex()
+            # Cache at HOST level - share feature flags across connections to same host
+            # Feature flags are per-host, not per-session
+            key = connection.session.host
             if key not in cls._context_map:
                 cls._context_map[key] = FeatureFlagsContext(
                     connection, cls._executor, connection.session.http_client
@@ -177,7 +178,8 @@ def get_instance(cls, connection: "Connection") -> FeatureFlagsContext:
     def remove_instance(cls, connection: "Connection"):
         """Removes the context for a given connection and shuts down the executor if no clients remain."""
         with cls._lock:
-            key = connection.get_session_id_hex()
+            # Use host as key to match get_instance
+            key = connection.session.host
             if key in cls._context_map:
                 cls._context_map.pop(key, None)
 

src/databricks/sql/telemetry/latency_logger.py

@@ -1,6 +1,6 @@
 import time
 import functools
-from typing import Optional
+from typing import Optional, Dict, Any
 import logging
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.telemetry.models.event import (
@@ -11,127 +11,141 @@
 logger = logging.getLogger(__name__)
 
 
-class TelemetryExtractor:
+def _extract_cursor_data(cursor) -> Dict[str, Any]:
     """
-    Base class for extracting telemetry information from various object types.
+    Extract telemetry data directly from a Cursor object.
 
-    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):
-        self._obj = obj
-
-    def __getattr__(self, name):
-        return getattr(self._obj, name)
-
-    def get_session_id_hex(self):
-        pass
-
-    def get_statement_id(self):
-        pass
-
-    def get_is_compressed(self):
-        pass
-
-    def get_execution_result_format(self):
-        pass
-
-    def get_retry_count(self):
-        pass
-
-    def get_chunk_id(self):
-        pass
+    OPTIMIZATION: Uses direct attribute access instead of wrapper objects.
+    This eliminates object creation overhead and method call indirection.
 
+    Args:
+        cursor: The Cursor object to extract data from
 
-class CursorExtractor(TelemetryExtractor):
+    Returns:
+        Dict with telemetry data (values may be None if extraction fails)
     """
-    Telemetry extractor specialized for Cursor objects.
-
-    Extracts telemetry information from database cursor objects, including
-    statement IDs, session information, compression settings, and result formats.
+    data = {}
+
+    # Extract statement_id (query_id) - direct attribute access
+    try:
+        data["statement_id"] = cursor.query_id
+    except (AttributeError, Exception):
+        data["statement_id"] = None
+
+    # Extract session_id_hex - direct method call
+    try:
+        data["session_id_hex"] = cursor.connection.get_session_id_hex()
+    except (AttributeError, Exception):
+        data["session_id_hex"] = None
+
+    # Extract is_compressed - direct attribute access
+    try:
+        data["is_compressed"] = cursor.connection.lz4_compression
+    except (AttributeError, Exception):
+        data["is_compressed"] = False
+
+    # Extract execution_result_format - inline logic
+    try:
+        if cursor.active_result_set is None:
+            data["execution_result"] = ExecutionResultFormat.FORMAT_UNSPECIFIED
+        else:
+            from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
+
+            results = cursor.active_result_set.results
+            if isinstance(results, ColumnQueue):
+                data["execution_result"] = ExecutionResultFormat.COLUMNAR_INLINE
+            elif isinstance(results, CloudFetchQueue):
+                data["execution_result"] = ExecutionResultFormat.EXTERNAL_LINKS
+            elif isinstance(results, ArrowQueue):
+                data["execution_result"] = ExecutionResultFormat.INLINE_ARROW
+            else:
+                data["execution_result"] = ExecutionResultFormat.FORMAT_UNSPECIFIED
+    except (AttributeError, Exception):
+        data["execution_result"] = ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+    # Extract retry_count - direct attribute access
+    try:
+        if hasattr(cursor.backend, "retry_policy") and cursor.backend.retry_policy:
+            data["retry_count"] = len(cursor.backend.retry_policy.history)
+        else:
+            data["retry_count"] = 0
+    except (AttributeError, Exception):
+        data["retry_count"] = 0
+
+    # chunk_id is always None for Cursor
+    data["chunk_id"] = None
+
+    return data
+
+
+def _extract_result_set_handler_data(handler) -> Dict[str, Any]:
     """
+    Extract telemetry data directly from a ResultSetDownloadHandler object.
 
-    def get_statement_id(self) -> Optional[str]:
-        return self.query_id
-
-    def get_session_id_hex(self) -> Optional[str]:
-        return self.connection.get_session_id_hex()
-
-    def get_is_compressed(self) -> bool:
-        return self.connection.lz4_compression
-
-    def get_execution_result_format(self) -> ExecutionResultFormat:
-        if self.active_result_set is None:
-            return ExecutionResultFormat.FORMAT_UNSPECIFIED
-
-        from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
-
-        if isinstance(self.active_result_set.results, ColumnQueue):
-            return ExecutionResultFormat.COLUMNAR_INLINE
-        elif isinstance(self.active_result_set.results, CloudFetchQueue):
-            return ExecutionResultFormat.EXTERNAL_LINKS
-        elif isinstance(self.active_result_set.results, ArrowQueue):
-            return ExecutionResultFormat.INLINE_ARROW
-        return ExecutionResultFormat.FORMAT_UNSPECIFIED
-
-    def get_retry_count(self) -> int:
-        if hasattr(self.backend, "retry_policy") and self.backend.retry_policy:
-            return len(self.backend.retry_policy.history)
-        return 0
-
-    def get_chunk_id(self):
-        return None
+    OPTIMIZATION: Uses direct attribute access instead of wrapper objects.
 
+    Args:
+        handler: The ResultSetDownloadHandler object to extract data from
 
-class ResultSetDownloadHandlerExtractor(TelemetryExtractor):
-    """
-    Telemetry extractor specialized for ResultSetDownloadHandler objects.
+    Returns:
+        Dict with telemetry data (values may be None if extraction fails)
     """
+    data = {}
 
-    def get_session_id_hex(self) -> Optional[str]:
-        return self._obj.session_id_hex
+    # Extract session_id_hex - direct attribute access
+    try:
+        data["session_id_hex"] = handler.session_id_hex
+    except (AttributeError, Exception):
+        data["session_id_hex"] = None
 
-    def get_statement_id(self) -> Optional[str]:
-        return self._obj.statement_id
+    # Extract statement_id - direct attribute access
+    try:
+        data["statement_id"] = handler.statement_id
+    except (AttributeError, Exception):
+        data["statement_id"] = None
 
-    def get_is_compressed(self) -> bool:
-        return self._obj.settings.is_lz4_compressed
+    # Extract is_compressed - direct attribute access
+    try:
+        data["is_compressed"] = handler.settings.is_lz4_compressed
+    except (AttributeError, Exception):
+        data["is_compressed"] = False
 
-    def get_execution_result_format(self) -> ExecutionResultFormat:
-        return ExecutionResultFormat.EXTERNAL_LINKS
+    # execution_result is always EXTERNAL_LINKS for result set handlers
+    data["execution_result"] = ExecutionResultFormat.EXTERNAL_LINKS
 
-    def get_retry_count(self) -> Optional[int]:
-        # standard requests and urllib3 libraries don't expose retry count
-        return None
+    # retry_count is not available for result set handlers
+    data["retry_count"] = None
+
+    # Extract chunk_id - direct attribute access
+    try:
+        data["chunk_id"] = handler.chunk_id
+    except (AttributeError, Exception):
+        data["chunk_id"] = None
 
-    def get_chunk_id(self) -> Optional[int]:
-        return self._obj.chunk_id
+    return data
 
 
-def get_extractor(obj):
+def _extract_telemetry_data(obj) -> Optional[Dict[str, Any]]:
     """
-    Factory function to create the appropriate telemetry extractor for an object.
+    Extract telemetry data from an object based on its type.
 
-    Determines the object type and returns the corresponding specialized extractor
-    that can extract telemetry information from that object type.
+    OPTIMIZATION: Returns a simple dict instead of creating wrapper objects.
+    This dict will be used to create the SqlExecutionEvent in the background thread.
 
     Args:
-        obj: The object to create an extractor for. Can be a Cursor,
-             ResultSetDownloadHandler, or any other object.
+        obj: The object to extract data from (Cursor, ResultSetDownloadHandler, etc.)
 
     Returns:
-        TelemetryExtractor: A specialized extractor instance:
-            - CursorExtractor for Cursor objects
-            - ResultSetDownloadHandlerExtractor for ResultSetDownloadHandler objects
-            - None for all other objects
+        Dict with telemetry data, or None if object type is not supported
     """
-    if obj.__class__.__name__ == "Cursor":
-        return CursorExtractor(obj)
-    elif obj.__class__.__name__ == "ResultSetDownloadHandler":
-        return ResultSetDownloadHandlerExtractor(obj)
+    obj_type = obj.__class__.__name__
+
+    if obj_type == "Cursor":
+        return _extract_cursor_data(obj)
+    elif obj_type == "ResultSetDownloadHandler":
+        return _extract_result_set_handler_data(obj)
     else:
-        logger.debug("No extractor found for %s", obj.__class__.__name__)
+        logger.debug("No telemetry extraction available for %s", obj_type)
         return None
 
 
@@ -143,12 +157,6 @@ def log_latency(statement_type: StatementType = StatementType.NONE):
     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
-
     Args:
         statement_type (StatementType): The type of SQL statement being executed.
 
@@ -162,54 +170,49 @@ def execute(self, query):
         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).
+        The wrapped method's object (self) must be a Cursor or
+        ResultSetDownloadHandler for telemetry data extraction.
     """
 
     def decorator(func):
         @functools.wraps(func)
         def wrapper(self, *args, **kwargs):
-            start_time = time.perf_counter()
-            result = None
+            start_time = time.monotonic()
             try:
-                result = func(self, *args, **kwargs)
-                return result
+                return func(self, *args, **kwargs)
             finally:
-
-                def _safe_call(func_to_call):
-                    """Calls a function and returns a default value on any exception."""
-                    try:
-                        return func_to_call()
-                    except Exception:
-                        return None
-
-                end_time = time.perf_counter()
-                duration_ms = int((end_time - start_time) * 1000)
-
-                extractor = get_extractor(self)
-
-                if extractor is not None:
-                    session_id_hex = _safe_call(extractor.get_session_id_hex)
-                    statement_id = _safe_call(extractor.get_statement_id)
-
-                    sql_exec_event = SqlExecutionEvent(
-                        statement_type=statement_type,
-                        is_compressed=_safe_call(extractor.get_is_compressed),
-                        execution_result=_safe_call(
-                            extractor.get_execution_result_format
-                        ),
-                        retry_count=_safe_call(extractor.get_retry_count),
-                        chunk_id=_safe_call(extractor.get_chunk_id),
-                    )
-
-                    telemetry_client = TelemetryClientFactory.get_telemetry_client(
-                        session_id_hex
-                    )
-                    telemetry_client.export_latency_log(
-                        latency_ms=duration_ms,
-                        sql_execution_event=sql_exec_event,
-                        sql_statement_id=statement_id,
-                    )
+                duration_ms = int((time.monotonic() - start_time) * 1000)
+
+                # Always log for debugging
+                logger.debug("%s completed in %dms", func.__name__, duration_ms)
+
+                # Fast check: use cached telemetry_enabled flag from connection
+                # Avoids dictionary lookup + instance check on every operation
+                connection = getattr(self, "connection", None)
+                if connection and getattr(connection, "telemetry_enabled", False):
+                    session_id_hex = connection.get_session_id_hex()
+                    if session_id_hex:
+                        # Telemetry enabled - extract and send
+                        telemetry_data = _extract_telemetry_data(self)
+                        if telemetry_data:
+                            sql_exec_event = SqlExecutionEvent(
+                                statement_type=statement_type,
+                                is_compressed=telemetry_data.get("is_compressed"),
+                                execution_result=telemetry_data.get("execution_result"),
+                                retry_count=telemetry_data.get("retry_count"),
+                                chunk_id=telemetry_data.get("chunk_id"),
+                            )
+
+                            telemetry_client = (
+                                TelemetryClientFactory.get_telemetry_client(
+                                    session_id_hex
+                                )
+                            )
+                            telemetry_client.export_latency_log(
+                                latency_ms=duration_ms,
+                                sql_execution_event=sql_exec_event,
+                                sql_statement_id=telemetry_data.get("statement_id"),
+                            )
 
         return wrapper