Add functionality for export of latency logs via telemetry

src/databricks/sql/client.py

@@ -61,7 +61,7 @@
     DriverConnectionParameters,
     HostDetails,
 )
-
+from databricks.sql.telemetry.latency_logger import log_latency
 
 logger = logging.getLogger(__name__)
 
@@ -745,6 +745,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
     ):
@@ -784,6 +785,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
     ):
@@ -811,6 +813,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
     ):
@@ -824,6 +827,7 @@ def _handle_staging_remove(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    @log_latency()
     def execute(
         self,
         operation: str,
@@ -914,6 +918,7 @@ def execute(
 
         return self
 
+    @log_latency()
     def execute_async(
         self,
         operation: str,
@@ -1039,6 +1044,7 @@ def executemany(self, operation, seq_of_parameters):
             self.execute(operation, parameters)
         return self
 
+    @log_latency()
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
@@ -1062,6 +1068,7 @@ def catalogs(self) -> "Cursor":
         )
         return self
 
+    @log_latency()
     def schemas(
         self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> "Cursor":
@@ -1090,6 +1097,7 @@ def schemas(
         )
         return self
 
+    @log_latency()
     def tables(
         self,
         catalog_name: Optional[str] = None,
@@ -1125,6 +1133,7 @@ def tables(
         )
         return self
 
+    @log_latency()
     def columns(
         self,
         catalog_name: Optional[str] = None,
@@ -1379,6 +1388,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 +1401,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,152 @@
+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:
+    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_statement_type(self):
+        pass
+
+    def get_is_compressed(self):
+        pass
+
+    def get_execution_result(self):
+        pass
+
+    def get_retry_count(self):
+        pass
+
+
+class CursorExtractor(TelemetryExtractor):
+    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
+
+    def get_statement_type(self) -> StatementType:
+        # TODO: Implement this
+        return StatementType.SQL
+
+
+class ResultSetExtractor(TelemetryExtractor):
+    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_statement_type(self) -> StatementType:
+        # TODO: Implement this
+        return StatementType.SQL
+
+    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):
+    if obj.__class__.__name__ == "Cursor":
+        return CursorExtractor(obj)
+    elif obj.__class__.__name__ == "ResultSet":
+        return ResultSetExtractor(obj)
+    else:
+        return TelemetryExtractor(obj)
+
+
+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)
+
+                extractor = get_extractor(self)
+                session_id_hex = extractor.get_session_id_hex()
+                statement_id = extractor.get_statement_id()
+
+                sql_exec_event = SqlExecutionEvent(
+                    statement_type=extractor.get_statement_type(),
+                    is_compressed=extractor.get_is_compressed(),
+                    execution_result=extractor.get_execution_result(),
+                    retry_count=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

src/databricks/sql/telemetry/telemetry_client.py

@@ -4,7 +4,7 @@
 import requests
 import logging
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -112,6 +112,10 @@ 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):
+        raise NotImplementedError("Subclasses must implement export_latency_log")
+
     @abstractmethod
     def close(self):
         raise NotImplementedError("Subclasses must implement close")
@@ -136,6 +140,9 @@ 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):
+        pass
+
     def close(self):
         pass
 
@@ -299,6 +306,32 @@ 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):
+        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)
@@ -431,6 +464,9 @@ def close(session_id_hex):
                 logger.debug(
                     "No more telemetry clients, shutting down thread pool executor"
                 )
-                TelemetryClientFactory._executor.shutdown(wait=True)
+                try:
+                    TelemetryClientFactory._executor.shutdown(wait=True)
+                except Exception as e:
+                    logger.debug("Failed to shutdown thread pool executor: %s", e)
                 TelemetryClientFactory._executor = None
                 TelemetryClientFactory._initialized = False

tests/unit/test_client.py

@@ -38,6 +38,11 @@ def new(cls):
         cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
         MockTExecuteStatementResp = MagicMock(spec=TExecuteStatementResp())
 
+        # Mock retry_policy with history attribute
+        mock_retry_policy = Mock()
+        mock_retry_policy.history = []
+        cls.apply_property_to_mock(ThriftBackendMock, retry_policy=mock_retry_policy)
+
         cls.apply_property_to_mock(
             MockTExecuteStatementResp,
             description=None,
@@ -69,6 +74,15 @@ def apply_property_to_mock(self, mock_obj, **kwargs):
             prop = PropertyMock(**kwargs)
             setattr(type(mock_obj), key, prop)
 
+    @classmethod
+    def mock_thrift_backend_with_retry_policy(cls): # Required for log_latency() decorator
+        """Create a simple thrift_backend mock with retry_policy for basic tests."""
+        mock_thrift_backend = Mock()
+        mock_retry_policy = Mock()
+        mock_retry_policy.history = []
+        mock_thrift_backend.retry_policy = mock_retry_policy
+        return mock_thrift_backend
+
 
 class ClientTestSuite(unittest.TestCase):
     """
@@ -318,7 +332,7 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
         mock_result_sets[1].fetchall.assert_called_once_with()
 
     def test_closed_cursor_doesnt_allow_operations(self):
-        cursor = client.Cursor(Mock(), Mock())
+        cursor = client.Cursor(Mock(), ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy())
         cursor.close()
 
         with self.assertRaises(Error) as e:
@@ -380,7 +394,7 @@ def test_get_schemas_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.schemas(**req_args)
@@ -403,7 +417,7 @@ def test_get_tables_parameters_passed_to_thrift_backend(self, mock_thrift_backen
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.tables(**req_args)
@@ -426,7 +440,7 @@ def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.columns(**req_args)