added test for export_latency_log, made mock of thrift backend with retry policy

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -1185,7 +1185,6 @@ 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.
@@ -1219,7 +1218,6 @@ 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
@@ -1245,7 +1243,6 @@ 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:
@@ -1256,7 +1253,6 @@ 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:
@@ -1380,7 +1376,11 @@ def get_execution_result(self) -> ExecutionResultFormat:
         return ExecutionResultFormat.FORMAT_UNSPECIFIED
 
     def get_retry_count(self) -> int:
-        # return len(self.thrift_backend.retry_policy.history)
+        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, func_name: str) -> StatementType:
@@ -1712,5 +1712,9 @@ def get_statement_type(self, func_name: str) -> StatementType:
         return StatementType.SQL
 
     def get_retry_count(self) -> int:
-        # return len(self.thrift_backend.retry_policy.history)
+        if (
+            hasattr(self.thrift_backend, "retry_policy")
+            and self.thrift_backend.retry_policy
+        ):
+            return len(self.thrift_backend.retry_policy.history)
         return 0

src/databricks/sql/telemetry/telemetry_client.py

@@ -113,9 +113,7 @@ 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
-    ):
+    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
         raise NotImplementedError("Subclasses must implement export_latency_log")
 
     @abstractmethod
@@ -310,9 +308,7 @@ 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
-    ):
+    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(

tests/unit/test_client.py

@@ -39,6 +39,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,
@@ -70,6 +75,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):
     """
@@ -319,7 +333,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:
@@ -399,7 +413,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)
@@ -422,7 +436,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)
@@ -445,7 +459,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)

tests/unit/test_fetches.py

@@ -31,6 +31,15 @@ def make_arrow_queue(batch):
         _, table = FetchTests.make_arrow_table(batch)
         queue = ArrowQueue(table, len(batch))
         return queue
+    
+    @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
 
     @staticmethod
     def make_dummy_result_set_from_initial_results(initial_results):
@@ -39,7 +48,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
         arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
         rs = client.ResultSet(
             connection=Mock(),
-            thrift_backend=None,
+            thrift_backend=FetchTests.mock_thrift_backend_with_retry_policy(),
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
@@ -79,7 +88,7 @@ def fetch_results(
 
             return results, batch_index < len(batch_list)
 
-        mock_thrift_backend = Mock()
+        mock_thrift_backend = FetchTests.mock_thrift_backend_with_retry_policy()
         mock_thrift_backend.fetch_results = fetch_results
         num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
 

tests/unit/test_telemetry.py

@@ -96,6 +96,12 @@ def test_export_failure_log(self, noop_telemetry_client):
             error_name="TestError", error_message="Test error message"
         )
 
+    def test_export_latency_log(self, noop_telemetry_client):
+        """Test that export_latency_log does nothing."""
+        noop_telemetry_client.export_latency_log(
+            latency_ms=100, sql_execution_event="EXECUTE_STATEMENT", sql_statement_id="test-id"
+        )
+
     def test_close(self, noop_telemetry_client):
         """Test that close does nothing."""
         noop_telemetry_client.close()
@@ -181,6 +187,40 @@ def test_export_failure_log(
 
         client._export_event.assert_called_once_with(mock_frontend_log.return_value)
 
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryFrontendLog")
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHelper.get_driver_system_configuration")
+    @patch("databricks.sql.telemetry.telemetry_client.uuid.uuid4")
+    @patch("databricks.sql.telemetry.telemetry_client.time.time")
+    def test_export_latency_log(
+        self, 
+        mock_time, 
+        mock_uuid4, 
+        mock_get_driver_config, 
+        mock_frontend_log,
+        telemetry_client_setup
+    ):
+        """Test exporting latency telemetry log."""
+        mock_time.return_value = 3000
+        mock_uuid4.return_value = "test-latency-uuid"
+        mock_get_driver_config.return_value = "test-driver-config"
+        mock_frontend_log.return_value = MagicMock()
+
+        client = telemetry_client_setup["client"]
+        client._export_event = MagicMock()
+        
+        client._driver_connection_params = "test-connection-params"
+        client._user_agent = "test-user-agent"
+        
+        latency_ms = 150
+        sql_execution_event = "test-execution-event"
+        sql_statement_id = "test-statement-id"
+        
+        client.export_latency_log(latency_ms, sql_execution_event, sql_statement_id)
+        
+        mock_frontend_log.assert_called_once()
+        
+        client._export_event.assert_called_once_with(mock_frontend_log.return_value)
+
     def test_export_event(self, telemetry_client_setup):
         """Test exporting an event."""
         client = telemetry_client_setup["client"]
@@ -311,6 +351,11 @@ def test_telemetry_client_exception_handling(self, telemetry_client_setup):
             # Should not raise exception
             client.export_failure_log("TestError", "Test error message")
 
+        # Test export_latency_log with exception
+        with patch.object(client, '_export_event', side_effect=Exception("Test error")):
+            # Should not raise exception
+            client.export_latency_log(100, "EXECUTE_STATEMENT", "test-statement-id")
+        
         # Test _send_telemetry with exception
         with patch.object(client._executor, 'submit', side_effect=Exception("Test error")):
             # Should not raise exception