Perf update

Author: jprakash-db

src/databricks/sql/client.py

@@ -1514,7 +1514,7 @@ 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()
         self._next_row_index += results.num_rows
-        
+
         partial_result_chunks = [results]
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()

src/databricks/sql/cloudfetch/downloader.py

@@ -89,7 +89,7 @@ def run(self) -> DownloadedFile:
         ResultSetDownloadHandler._validate_link(
             self.link, self.settings.link_expiry_buffer_secs
         )
-        
+
         with self._http_client.execute(
             method=HttpMethod.GET,
             url=self.link.fileLink,

tests/unit/test_telemetry.py

@@ -8,7 +8,7 @@
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient
+    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
@@ -24,7 +24,7 @@ def mock_telemetry_client():
     session_id = str(uuid.uuid4())
     auth_provider = AccessTokenAuthProvider("test-token")
     executor = MagicMock()
-    
+
     return TelemetryClient(
         telemetry_enabled=True,
         session_id_hex=session_id,
@@ -43,7 +43,7 @@ def test_noop_client_behavior(self):
         client1 = NoopTelemetryClient()
         client2 = NoopTelemetryClient()
         assert client1 is client2
-        
+
         # Test that all methods can be called without exceptions
         client1.export_initial_telemetry_log(MagicMock(), "test-agent")
         client1.export_failure_log("TestError", "Test message")
@@ -58,76 +58,76 @@ def test_event_batching_and_flushing_flow(self, mock_telemetry_client):
         """Test the complete event batching and flushing flow."""
         client = mock_telemetry_client
         client._batch_size = 3  # Small batch for testing
-        
+
         # Mock the network call
-        with patch.object(client, '_send_telemetry') as mock_send:
+        with patch.object(client, "_send_telemetry") as mock_send:
             # Add events one by one - should not flush yet
             client._export_event("event1")
             client._export_event("event2")
             mock_send.assert_not_called()
             assert len(client._events_batch) == 2
-            
+
             # Third event should trigger flush
             client._export_event("event3")
             mock_send.assert_called_once()
             assert len(client._events_batch) == 0  # Batch cleared after flush
-    
-    @patch('requests.post')
+
+    @patch("requests.post")
     def test_network_request_flow(self, mock_post, mock_telemetry_client):
         """Test the complete network request flow with authentication."""
         mock_post.return_value.status_code = 200
         client = mock_telemetry_client
-        
+
         # Create mock events
         mock_events = [MagicMock() for _ in range(2)]
         for i, event in enumerate(mock_events):
             event.to_json.return_value = f'{{"event": "{i}"}}'
-        
+
         # Send telemetry
         client._send_telemetry(mock_events)
-        
+
         # Verify request was submitted to executor
         client._executor.submit.assert_called_once()
         args, kwargs = client._executor.submit.call_args
-        
+
         # Verify correct function and URL
         assert args[0] == requests.post
-        assert args[1] == 'https://test-host.com/telemetry-ext'
-        assert kwargs['headers']['Authorization'] == 'Bearer test-token'
-        
+        assert args[1] == "https://test-host.com/telemetry-ext"
+        assert kwargs["headers"]["Authorization"] == "Bearer test-token"
+
         # Verify request body structure
-        request_data = kwargs['data']
+        request_data = kwargs["data"]
         assert '"uploadTime"' in request_data
         assert '"protoLogs"' in request_data
 
     def test_telemetry_logging_flows(self, mock_telemetry_client):
         """Test all telemetry logging methods work end-to-end."""
         client = mock_telemetry_client
-        
-        with patch.object(client, '_export_event') as mock_export:
+
+        with patch.object(client, "_export_event") as mock_export:
             # Test initial log
             client.export_initial_telemetry_log(MagicMock(), "test-agent")
             assert mock_export.call_count == 1
-            
+
             # Test failure log
             client.export_failure_log("TestError", "Error message")
             assert mock_export.call_count == 2
-            
+
             # Test latency log
             client.export_latency_log(150, "EXECUTE_STATEMENT", "stmt-123")
             assert mock_export.call_count == 3
 
     def test_error_handling_resilience(self, mock_telemetry_client):
         """Test that telemetry errors don't break the client."""
         client = mock_telemetry_client
-        
+
         # Test that exceptions in telemetry don't propagate
-        with patch.object(client, '_export_event', side_effect=Exception("Test error")):
+        with patch.object(client, "_export_event", side_effect=Exception("Test error")):
             # These should not raise exceptions
             client.export_initial_telemetry_log(MagicMock(), "test-agent")
             client.export_failure_log("TestError", "Error message")
             client.export_latency_log(100, "EXECUTE_STATEMENT", "stmt-123")
-        
+
         # Test executor submission failure
         client._executor.submit.side_effect = Exception("Thread pool error")
         client._send_telemetry([MagicMock()])  # Should not raise
@@ -140,7 +140,7 @@ def test_system_configuration_caching(self):
         """Test that system configuration is cached and contains expected data."""
         config1 = TelemetryHelper.get_driver_system_configuration()
         config2 = TelemetryHelper.get_driver_system_configuration()
-        
+
         # Should be cached (same instance)
         assert config1 is config2
 
@@ -153,7 +153,7 @@ def test_auth_mechanism_detection(self):
             (MagicMock(), AuthMech.OTHER),  # Unknown provider
             (None, None),
         ]
-        
+
         for provider, expected in test_cases:
             assert TelemetryHelper.get_auth_mechanism(provider) == expected
 
@@ -163,19 +163,25 @@ def test_auth_flow_detection(self):
         oauth_with_tokens = MagicMock(spec=DatabricksOAuthProvider)
         oauth_with_tokens._access_token = "test-access-token"
         oauth_with_tokens._refresh_token = "test-refresh-token"
-        assert TelemetryHelper.get_auth_flow(oauth_with_tokens) == AuthFlow.TOKEN_PASSTHROUGH
-        
+        assert (
+            TelemetryHelper.get_auth_flow(oauth_with_tokens)
+            == AuthFlow.TOKEN_PASSTHROUGH
+        )
+
         # Test OAuth with browser-based auth
         oauth_with_browser = MagicMock(spec=DatabricksOAuthProvider)
         oauth_with_browser._access_token = None
         oauth_with_browser._refresh_token = None
         oauth_with_browser.oauth_manager = MagicMock()
-        assert TelemetryHelper.get_auth_flow(oauth_with_browser) == AuthFlow.BROWSER_BASED_AUTHENTICATION
-        
+        assert (
+            TelemetryHelper.get_auth_flow(oauth_with_browser)
+            == AuthFlow.BROWSER_BASED_AUTHENTICATION
+        )
+
         # Test non-OAuth provider
         pat_auth = AccessTokenAuthProvider("test-token")
         assert TelemetryHelper.get_auth_flow(pat_auth) is None
-        
+
         # Test None auth provider
         assert TelemetryHelper.get_auth_flow(None) is None
 
@@ -202,56 +208,58 @@ def test_client_lifecycle_flow(self):
         """Test complete client lifecycle: initialize -> use -> close."""
         session_id_hex = "test-session"
         auth_provider = AccessTokenAuthProvider("token")
-        
+
         # Initialize enabled client
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=True,
             session_id_hex=session_id_hex,
             auth_provider=auth_provider,
-            host_url="test-host.com"
+            host_url="test-host.com",
         )
-        
+
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, TelemetryClient)
         assert client._session_id_hex == session_id_hex
-        
+
         # Close client
-        with patch.object(client, 'close') as mock_close:
+        with patch.object(client, "close") as mock_close:
             TelemetryClientFactory.close(session_id_hex)
             mock_close.assert_called_once()
-        
+
         # Should get NoopTelemetryClient after close
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, NoopTelemetryClient)
 
     def test_disabled_telemetry_flow(self):
         """Test that disabled telemetry uses NoopTelemetryClient."""
         session_id_hex = "test-session"
-        
+
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=False,
             session_id_hex=session_id_hex,
             auth_provider=None,
-            host_url="test-host.com"
+            host_url="test-host.com",
         )
-        
+
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, NoopTelemetryClient)
 
     def test_factory_error_handling(self):
         """Test that factory errors fall back to NoopTelemetryClient."""
         session_id = "test-session"
-        
+
         # Simulate initialization error
-        with patch('databricks.sql.telemetry.telemetry_client.TelemetryClient', 
-                  side_effect=Exception("Init error")):
+        with patch(
+            "databricks.sql.telemetry.telemetry_client.TelemetryClient",
+            side_effect=Exception("Init error"),
+        ):
             TelemetryClientFactory.initialize_telemetry_client(
                 telemetry_enabled=True,
                 session_id_hex=session_id,
                 auth_provider=AccessTokenAuthProvider("token"),
-                host_url="test-host.com"
+                host_url="test-host.com",
             )
-        
+
         # Should fall back to NoopTelemetryClient
         client = TelemetryClientFactory.get_telemetry_client(session_id)
         assert isinstance(client, NoopTelemetryClient)
@@ -260,25 +268,25 @@ def test_factory_shutdown_flow(self):
         """Test factory shutdown when last client is removed."""
         session1 = "session-1"
         session2 = "session-2"
-        
+
         # Initialize multiple clients
         for session in [session1, session2]:
             TelemetryClientFactory.initialize_telemetry_client(
                 telemetry_enabled=True,
                 session_id_hex=session,
                 auth_provider=AccessTokenAuthProvider("token"),
-                host_url="test-host.com"
+                host_url="test-host.com",
             )
-        
+
         # Factory should be initialized
         assert TelemetryClientFactory._initialized is True
         assert TelemetryClientFactory._executor is not None
-        
+
         # Close first client - factory should stay initialized
         TelemetryClientFactory.close(session1)
         assert TelemetryClientFactory._initialized is True
-        
+
         # Close second client - factory should shut down
         TelemetryClientFactory.close(session2)
         assert TelemetryClientFactory._initialized is False
-        assert TelemetryClientFactory._executor is None
+        assert TelemetryClientFactory._executor is None