test fixes

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

Author: saishreeeee

src/databricks/sql/telemetry/latency_logger.py

@@ -253,7 +253,7 @@ def _safe_call(func_to_call):
                         statement_type=statement_type,
                         is_compressed=_safe_call(extractor.get_is_compressed),
                         execution_result=_safe_call(extractor.get_execution_result),
-                        retry_count=extractor.get_retry_count(),
+                        retry_count=_safe_call(extractor.get_retry_count),
                         chunk_id=_safe_call(extractor.get_chunk_id),
                     )
 

tests/unit/test_downloader.py

@@ -27,7 +27,7 @@ def test_run_link_expired(self, mock_time):
         # Already expired
         result_link.expiryTime = 999
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
 
         with self.assertRaises(Error) as context:
@@ -43,7 +43,7 @@ def test_run_link_past_expiry_buffer(self, mock_time):
         # Within the expiry buffer time
         result_link.expiryTime = 1004
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
 
         with self.assertRaises(Error) as context:
@@ -63,7 +63,7 @@ def test_run_get_response_not_ok(self, mock_time, mock_session):
         result_link = Mock(expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
         with self.assertRaises(requests.exceptions.HTTPError) as context:
             d.run()
@@ -82,7 +82,7 @@ def test_run_uncompressed_successful(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
         file = d.run()
 
@@ -105,7 +105,7 @@ def test_run_compressed_successful(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
         file = d.run()
 
@@ -121,7 +121,7 @@ def test_download_connection_error(self, mock_time, mock_session):
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
         with self.assertRaises(ConnectionError):
             d.run()
@@ -136,7 +136,7 @@ def test_download_timeout(self, mock_time, mock_session):
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
+            settings, result_link, ssl_options=SSLOptions(), chunk_id=0
         )
         with self.assertRaises(TimeoutError):
             d.run()

tests/unit/test_thrift_backend.py

@@ -649,7 +649,7 @@ def test_handle_execute_response_sets_compression_in_direct_results(
                 ssl_options=SSLOptions(),
             )
 
-            execute_response, _ = thrift_backend._handle_execute_response(
+            execute_response, _, _ = thrift_backend._handle_execute_response(
                 t_execute_resp, Mock()
             )
             self.assertEqual(execute_response.lz4_compressed, lz4Compressed)
@@ -892,6 +892,7 @@ def test_handle_execute_response_can_handle_without_direct_results(
                 (
                     execute_response,
                     _,
+                    _,
                 ) = thrift_backend._handle_execute_response(execute_resp, Mock())
                 self.assertEqual(
                     execute_response.status,
@@ -927,7 +928,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
                     auth_provider=AuthProvider(),
                     ssl_options=SSLOptions(),
                 )
-                thrift_backend._results_message_to_execute_response = Mock()
+                thrift_backend._results_message_to_execute_response = Mock(return_value=(Mock(), Mock()))
 
                 thrift_backend._handle_execute_response(execute_resp, Mock())
 
@@ -965,7 +966,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             )
         )
         thrift_backend = self._make_fake_thrift_backend()
-        execute_response, _ = thrift_backend._handle_execute_response(
+        execute_response, _, _ = thrift_backend._handle_execute_response(
             t_execute_resp, Mock()
         )
 
@@ -997,7 +998,7 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
             )
         )
         thrift_backend = self._make_fake_thrift_backend()
-        _, _ = thrift_backend._handle_execute_response(t_execute_resp, Mock())
+        _, _, _ = thrift_backend._handle_execute_response(t_execute_resp, Mock())
 
         self.assertEqual(
             hive_schema_mock,
@@ -1046,6 +1047,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                 (
                     execute_response,
                     has_more_rows_result,
+                    _
                 ) = thrift_backend._handle_execute_response(execute_resp, Mock())
 
                 self.assertEqual(is_direct_results, has_more_rows_result)
@@ -1179,7 +1181,7 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.execute_command(
@@ -1215,7 +1217,7 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_catalogs(Mock(), 100, 200, cursor_mock)
@@ -1248,7 +1250,7 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_schemas(
@@ -1290,7 +1292,7 @@ def test_get_tables_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_tables(
@@ -1336,7 +1338,7 @@ def test_get_columns_calls_client_and_handle_execute_response(
             ssl_options=SSLOptions(),
         )
         thrift_backend._handle_execute_response = Mock()
-        thrift_backend._handle_execute_response.return_value = (Mock(), Mock())
+        thrift_backend._handle_execute_response.return_value = (Mock(), Mock(), Mock())
         cursor_mock = Mock()
 
         result = thrift_backend.get_columns(
@@ -1682,7 +1684,7 @@ def test_handle_execute_response_sets_active_op_handle(self):
         thrift_backend = self._make_fake_thrift_backend()
         thrift_backend._check_direct_results_for_error = Mock()
         thrift_backend._wait_until_command_done = Mock()
-        thrift_backend._results_message_to_execute_response = Mock()
+        thrift_backend._results_message_to_execute_response = Mock(return_value=(Mock(), Mock()))
 
         # Create a mock response with a real operation handle
         mock_resp = Mock()
@@ -2254,6 +2256,7 @@ def test_execute_command_sets_complex_type_fields_correctly(
         mock_handle_execute_response.return_value = (
             mock_execute_response,
             mock_arrow_schema,
+            Mock()
         )
 
         # Iterate through each possible combination of native types (True, False and unset)