Add functionality for export of latency logs via telemetry

src/databricks/sql/client.py

@@ -61,7 +61,8 @@
     DriverConnectionParameters,
     HostDetails,
 )
-
+from databricks.sql.telemetry.latency_logger import log_latency
+from databricks.sql.telemetry.models.enums import StatementType
 
 logger = logging.getLogger(__name__)
 
@@ -745,6 +746,7 @@ def _handle_staging_operation(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency(StatementType.SQL)
     def _handle_staging_put(
         self, presigned_url: str, local_file: str, headers: Optional[dict] = None
     ):
@@ -784,6 +786,7 @@ def _handle_staging_put(
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
+    @log_latency(StatementType.SQL)
     def _handle_staging_get(
         self, local_file: str, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -811,6 +814,7 @@ def _handle_staging_get(
         with open(local_file, "wb") as fp:
             fp.write(r.content)
 
+    @log_latency(StatementType.SQL)
     def _handle_staging_remove(
         self, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -824,6 +828,7 @@ def _handle_staging_remove(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency(StatementType.QUERY)
     def execute(
         self,
         operation: str,
@@ -914,6 +919,7 @@ def execute(
 
         return self
 
+    @log_latency(StatementType.QUERY)
     def execute_async(
         self,
         operation: str,
@@ -1039,6 +1045,7 @@ def executemany(self, operation, seq_of_parameters):
             self.execute(operation, parameters)
         return self
 
+    @log_latency(StatementType.METADATA)
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
@@ -1062,6 +1069,7 @@ def catalogs(self) -> "Cursor":
         )
         return self
 
+    @log_latency(StatementType.METADATA)
     def schemas(
         self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> "Cursor":
@@ -1090,6 +1098,7 @@ def schemas(
         )
         return self
 
+    @log_latency(StatementType.METADATA)
     def tables(
         self,
         catalog_name: Optional[str] = None,
@@ -1125,6 +1134,7 @@ def tables(
         )
         return self
 
+    @log_latency(StatementType.METADATA)
     def columns(
         self,
         catalog_name: Optional[str] = None,
@@ -1379,6 +1389,7 @@ 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)
@@ -1391,6 +1402,7 @@ 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)
@@ -1433,6 +1445,7 @@ 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.
@@ -1475,6 +1488,7 @@ 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.
@@ -1500,6 +1514,7 @@ 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()
@@ -1526,6 +1541,7 @@ 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/telemetry/latency_logger.py

@@ -0,0 +1,231 @@
+import time
+import functools
+from typing import Optional
+from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
+from databricks.sql.telemetry.models.event import (
+    SqlExecutionEvent,
+)
+from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
+from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
+from uuid import UUID
+
+
+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):
+        pass
+
+    def get_statement_id(self):
+        pass
+
+    def get_is_compressed(self):
+        pass
+
+    def get_execution_result(self):
+        pass
+
+    def get_retry_count(self):
+        pass
+
+
+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
+
+    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(self) -> ExecutionResultFormat:
+        if self.active_result_set is None:
+            return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+        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.thrift_backend, "retry_policy")
+            and self.thrift_backend.retry_policy
+        ):
+            return len(self.thrift_backend.retry_policy.history)
+        return 0
+
+
+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))
+        return None
+
+    def get_session_id_hex(self) -> Optional[str]:
+        return self.connection.get_session_id_hex()
+
+    def get_is_compressed(self) -> bool:
+        return self.lz4_compressed
+
+    def get_execution_result(self) -> ExecutionResultFormat:
+        if isinstance(self.results, ColumnQueue):
+            return ExecutionResultFormat.COLUMNAR_INLINE
+        elif isinstance(self.results, CloudFetchQueue):
+            return ExecutionResultFormat.EXTERNAL_LINKS
+        elif isinstance(self.results, ArrowQueue):
+            return ExecutionResultFormat.INLINE_ARROW
+        return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+    def get_retry_count(self) -> int:
+        if (
+            hasattr(self.thrift_backend, "retry_policy")
+            and self.thrift_backend.retry_policy
+        ):
+            return len(self.thrift_backend.retry_policy.history)
+        return 0
+
+
+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
+            - Throws an NotImplementedError for all other objects
+    """
+    if obj.__class__.__name__ == "Cursor":
+        return CursorExtractor(obj)
+    elif obj.__class__.__name__ == "ResultSet":
+        return ResultSetExtractor(obj)
+    else:
+        raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")
+
+
+def log_latency(statement_type: StatementType = StatementType.NONE):
+    """
+    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
+
+    Args:
+        statement_type (StatementType): The type of SQL statement being executed.
+
+    Usage:
+        @log_latency(StatementType.SQL)
+        def execute(self, query):
+            # 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):
+            start_time = time.perf_counter()
+            result = None
+            try:
+                result = func(self, *args, **kwargs)
+                return result
+            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)
+                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),
+                    retry_count=_safe_call(extractor.get_retry_count),
+                )
+
+                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,
+                )
+
+        return wrapper
+
+    return decorator