latency logs funcitionality

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -61,7 +61,7 @@
     DriverConnectionParameters,
     HostDetails,
 )
-
+from databricks.sql.telemetry.latency_logger import log_latency
 
 logger = logging.getLogger(__name__)
 
@@ -758,6 +758,7 @@ def _handle_staging_operation(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def _handle_staging_put(
         self, presigned_url: str, local_file: str, headers: Optional[dict] = None
     ):
@@ -797,6 +798,7 @@ def _handle_staging_put(
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
+    @log_latency()
     def _handle_staging_get(
         self, local_file: str, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -824,6 +826,7 @@ def _handle_staging_get(
         with open(local_file, "wb") as fp:
             fp.write(r.content)
 
+    @log_latency()
     def _handle_staging_remove(
         self, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -837,6 +840,7 @@ def _handle_staging_remove(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def execute(
         self,
         operation: str,
@@ -927,6 +931,7 @@ def execute(
 
         return self
 
+    @log_latency()
     def execute_async(
         self,
         operation: str,
@@ -1052,6 +1057,7 @@ def executemany(self, operation, seq_of_parameters):
             self.execute(operation, parameters)
         return self
 
+    @log_latency()
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
@@ -1075,6 +1081,7 @@ def catalogs(self) -> "Cursor":
         )
         return self
 
+    @log_latency()
     def schemas(
         self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> "Cursor":
@@ -1103,6 +1110,7 @@ def schemas(
         )
         return self
 
+    @log_latency()
     def tables(
         self,
         catalog_name: Optional[str] = None,
@@ -1138,6 +1146,7 @@ def tables(
         )
         return self
 
+    @log_latency()
     def columns(
         self,
         catalog_name: Optional[str] = None,
@@ -1173,6 +1182,7 @@ def columns(
         )
         return self
 
+    @log_latency()
     def fetchall(self) -> List[Row]:
         """
         Fetch all (remaining) rows of a query result, returning them as a sequence of sequences.
@@ -1206,6 +1216,7 @@ def fetchone(self) -> Optional[Row]:
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def fetchmany(self, size: int) -> List[Row]:
         """
         Fetch the next set of rows of a query result, returning a sequence of sequences (e.g. a
@@ -1231,6 +1242,7 @@ def fetchmany(self, size: int) -> List[Row]:
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def fetchall_arrow(self) -> "pyarrow.Table":
         self._check_not_closed()
         if self.active_result_set:
@@ -1241,6 +1253,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def fetchmany_arrow(self, size) -> "pyarrow.Table":
         self._check_not_closed()
         if self.active_result_set:
@@ -1406,6 +1419,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)
@@ -1418,6 +1432,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)

src/databricks/sql/telemetry/latency_logger.py

@@ -0,0 +1,131 @@
+import time
+import functools
+from typing import Optional
+from uuid import UUID
+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
+
+# Helper to get statement_id/query_id from instance if available
+def _get_statement_id(instance) -> Optional[str]:
+    """
+    Get statement ID from an instance using various methods:
+    1. For Cursor: Use query_id property which returns UUID from active_op_handle
+    2. For ResultSet: Use command_id which contains operationId
+
+    Note: ThriftBackend itself doesn't have a statement ID since one backend
+    can handle multiple concurrent operations/cursors.
+    """
+    if hasattr(instance, "query_id"):
+        return instance.query_id
+
+    if hasattr(instance, "command_id") and instance.command_id:
+        return str(UUID(bytes=instance.command_id.operationId.guid))
+
+    return None
+
+
+def _get_session_id_hex(instance) -> Optional[str]:
+    if hasattr(instance, "connection") and instance.connection:
+        return instance.connection.get_session_id_hex()
+    if hasattr(instance, "get_session_id_hex"):
+        return instance.get_session_id_hex()
+    return None
+
+
+def _get_statement_type(func_name: str) -> StatementType:  # TODO: implement this
+    return StatementType.SQL
+
+
+def _get_is_compressed(instance) -> bool:
+    """
+    Get compression status from instance:
+    1. Direct lz4_compression attribute (Connection)
+    2. Through connection attribute (Cursor/ResultSet)
+    3. Through thrift_backend attribute (Cursor)
+    """
+    if hasattr(instance, "lz4_compression"):
+        return instance.lz4_compression
+    if hasattr(instance, "connection") and instance.connection:
+        return instance.connection.lz4_compression
+    if hasattr(instance, "thrift_backend") and instance.thrift_backend:
+        return instance.thrift_backend.lz4_compressed
+    return False
+
+
+def _get_execution_result(instance) -> ExecutionResultFormat:
+    """
+    Get execution result format from instance:
+    1. For ResultSet: Check if using cloud fetch (external_links) or arrow/columnar format
+    2. For Cursor: Check through active_result_set
+    3. For ThriftBackend: Check result format from server
+    """
+    if hasattr(instance, "_use_cloud_fetch") and instance._use_cloud_fetch:
+        return ExecutionResultFormat.EXTERNAL_LINKS
+
+    if hasattr(instance, "active_result_set") and instance.active_result_set:
+        if isinstance(instance.active_result_set.results, ColumnQueue):
+            return ExecutionResultFormat.COLUMNAR_INLINE
+        return ExecutionResultFormat.INLINE_ARROW
+
+    if hasattr(instance, "thrift_backend") and instance.thrift_backend:
+        if hasattr(instance.thrift_backend, "_use_arrow_native_complex_types"):
+            return ExecutionResultFormat.INLINE_ARROW
+
+    return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+
+def _get_retry_count(instance) -> int:
+    """
+    Get retry count from instance by checking retry_policy.history length.
+    The retry_policy is only accessible through thrift_backend.
+    """
+    # TODO: implement this
+
+    return 0
+
+
+def log_latency():
+    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:
+                end_time = time.perf_counter()
+                duration_ms = int((end_time - start_time) * 1000)
+
+                session_id_hex = _get_session_id_hex(self)
+
+                if session_id_hex:
+                    statement_id = _get_statement_id(self)
+                    statement_type = _get_statement_type(func.__name__)
+                    is_compressed = _get_is_compressed(self)
+                    execution_result = _get_execution_result(self)
+                    retry_count = _get_retry_count(self)
+
+                    sql_exec_event = SqlExecutionEvent(
+                        statement_type=statement_type,
+                        is_compressed=is_compressed,
+                        execution_result=execution_result,
+                        retry_count=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

src/databricks/sql/telemetry/telemetry_client.py

@@ -112,6 +112,12 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
     def export_failure_log(self, error_name, error_message):
         raise NotImplementedError("Subclasses must implement export_failure_log")
 
+    @abstractmethod
+    def export_latency_log(
+        self, latency_ms, sql_execution_event, sql_statement_id=None
+    ):
+        raise NotImplementedError("Subclasses must implement export_latency_log")
+
     @abstractmethod
     def close(self):
         raise NotImplementedError("Subclasses must implement close")
@@ -136,6 +142,11 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
     def export_failure_log(self, error_name, error_message):
         pass
 
+    def export_latency_log(
+        self, latency_ms, sql_execution_event, sql_statement_id=None
+    ):
+        pass
+
     def close(self):
         pass
 
@@ -299,6 +310,34 @@ def export_failure_log(self, error_name, error_message):
         except Exception as e:
             logger.debug("Failed to export failure log: %s", e)
 
+    def export_latency_log(
+        self, latency_ms, sql_execution_event, sql_statement_id=None
+    ):
+        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)
+
     def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)

src/databricks/sql/thrift_backend.py

@@ -583,6 +583,9 @@ def open_session(self, session_configuration, catalog, schema):
             self._transport.close()
             raise
 
+    def get_session_id_hex(self) -> str:
+        return self._session_id_hex
+
     def close_session(self, session_handle) -> None:
         req = ttypes.TCloseSessionReq(sessionHandle=session_handle)
         try: