rename proxy specific attrs with proxy prefix

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/utils/http_client.py

@@ -30,8 +30,8 @@ class SeaHttpClient:
     retry_policy: Union[DatabricksRetryPolicy, int]
     _pool: Optional[Union[HTTPConnectionPool, HTTPSConnectionPool]]
     proxy_uri: Optional[str]
-    realhost: Optional[str]
-    realport: Optional[int]
+    proxy_host: Optional[str]
+    proxy_port: Optional[int]
     proxy_auth: Optional[Dict[str, str]]
 
     def __init__(
@@ -136,15 +136,15 @@ def __init__(
 
         if proxy:
             parsed_proxy = urllib.parse.urlparse(proxy)
-            self.realhost = self.host
-            self.realport = self.port
+            self.proxy_host = self.host
+            self.proxy_port = self.port
             self.proxy_uri = proxy
             self.host = parsed_proxy.hostname
             self.port = parsed_proxy.port or (443 if self.scheme == "https" else 80)
             self.proxy_auth = self._basic_proxy_auth_headers(parsed_proxy)
         else:
-            self.realhost = None
-            self.realport = None
+            self.proxy_host = None
+            self.proxy_port = None
             self.proxy_auth = None
             self.proxy_uri = None
 
@@ -186,8 +186,8 @@ def _open(self):
                 proxy_headers=self.proxy_auth,
             )
             self._pool = proxy_manager.connection_from_host(
-                host=self.realhost,
-                port=self.realport,
+                host=self.proxy_host,
+                port=self.proxy_port,
                 scheme=self.scheme,
                 pool_kwargs=pool_kwargs,
             )
@@ -201,7 +201,7 @@ def close(self):
 
     def using_proxy(self) -> bool:
         """Check if proxy is being used."""
-        return self.realhost is not None
+        return self.proxy_host is not None
 
     def set_retry_command_type(self, command_type: CommandType):
         """Set the command type for retry policy decision making."""

tests/e2e/test_concurrent_telemetry.py

@@ -6,9 +6,13 @@
 import pytest
 
 from databricks.sql.telemetry.models.enums import StatementType
-from databricks.sql.telemetry.telemetry_client import TelemetryClient, TelemetryClientFactory
+from databricks.sql.telemetry.telemetry_client import (
+    TelemetryClient,
+    TelemetryClientFactory,
+)
 from tests.e2e.test_driver import PySQLPytestTestCase
 
+
 def run_in_threads(target, num_threads, pass_index=False):
     """Helper to run target function in multiple threads."""
     threads = [
@@ -22,7 +26,6 @@ def run_in_threads(target, num_threads, pass_index=False):
 
 
 class TestE2ETelemetry(PySQLPytestTestCase):
-    
     @pytest.fixture(autouse=True)
     def telemetry_setup_teardown(self):
         """
@@ -31,7 +34,7 @@ def telemetry_setup_teardown(self):
         this robust and automatic.
         """
         try:
-            yield 
+            yield
         finally:
             if TelemetryClientFactory._executor:
                 TelemetryClientFactory._executor.shutdown(wait=True)
@@ -68,20 +71,25 @@ def callback_wrapper(self_client, future, sent_count):
                 captured_futures.append(future)
             original_callback(self_client, future, sent_count)
 
-        with patch.object(TelemetryClient, "_send_telemetry", send_telemetry_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+        with patch.object(
+            TelemetryClient, "_send_telemetry", send_telemetry_wrapper
+        ), patch.object(
+            TelemetryClient, "_telemetry_request_callback", callback_wrapper
+        ):
 
             def execute_query_worker(thread_id):
                 """Each thread creates a connection and executes a query."""
 
                 time.sleep(random.uniform(0, 0.05))
-                
-                with self.connection(extra_params={"force_enable_telemetry": True}) as conn:
+
+                with self.connection(
+                    extra_params={"force_enable_telemetry": True}
+                ) as conn:
                     # Capture the session ID from the connection before executing the query
                     session_id_hex = conn.get_session_id_hex()
                     with capture_lock:
                         captured_session_ids.append(session_id_hex)
-                    
+
                     with conn.cursor() as cursor:
                         cursor.execute(f"SELECT {thread_id}")
                         # Capture the statement ID after executing the query
@@ -97,7 +105,10 @@ def execute_query_worker(thread_id):
             start_time = time.time()
             expected_event_count = num_threads
 
-            while len(captured_futures) < expected_event_count and time.time() - start_time < timeout_seconds:
+            while (
+                len(captured_futures) < expected_event_count
+                and time.time() - start_time < timeout_seconds
+            ):
                 time.sleep(0.1)
 
             done, not_done = wait(captured_futures, timeout=timeout_seconds)
@@ -115,30 +126,37 @@ def execute_query_worker(thread_id):
 
             assert not captured_exceptions
             assert len(captured_responses) > 0
-            
+
             total_successful_events = 0
             for response in captured_responses:
                 assert "errors" not in response or not response["errors"]
                 if "numProtoSuccess" in response:
                     total_successful_events += response["numProtoSuccess"]
             assert total_successful_events == num_threads * 2
 
-            assert len(captured_telemetry) == num_threads * 2  # 2 events per thread (initial_telemetry_log, latency_log (execute))
+            assert (
+                len(captured_telemetry) == num_threads * 2
+            )  # 2 events per thread (initial_telemetry_log, latency_log (execute))
             assert len(captured_session_ids) == num_threads  # One session ID per thread
-            assert len(captured_statement_ids) == num_threads  # One statement ID per thread (per query)
+            assert (
+                len(captured_statement_ids) == num_threads
+            )  # One statement ID per thread (per query)
 
             # Separate initial logs from latency logs
             initial_logs = [
-                e for e in captured_telemetry
+                e
+                for e in captured_telemetry
                 if e.entry.sql_driver_log.operation_latency_ms is None
                 and e.entry.sql_driver_log.driver_connection_params is not None
                 and e.entry.sql_driver_log.system_configuration is not None
             ]
             latency_logs = [
-                e for e in captured_telemetry
-                if e.entry.sql_driver_log.operation_latency_ms is not None 
-                and e.entry.sql_driver_log.sql_statement_id is not None 
-                and e.entry.sql_driver_log.sql_operation.statement_type == StatementType.QUERY
+                e
+                for e in captured_telemetry
+                if e.entry.sql_driver_log.operation_latency_ms is not None
+                and e.entry.sql_driver_log.sql_statement_id is not None
+                and e.entry.sql_driver_log.sql_operation.statement_type
+                == StatementType.QUERY
             ]
 
             # Verify counts
@@ -171,4 +189,4 @@ def execute_query_worker(thread_id):
             for event in latency_logs:
                 log = event.entry.sql_driver_log
                 assert log.sql_statement_id in captured_statement_ids
-                assert log.session_id in captured_session_ids
+                assert log.session_id in captured_session_ids

tests/unit/test_telemetry.py

@@ -30,7 +30,7 @@ def mock_telemetry_client():
         auth_provider=auth_provider,
         host_url="test-host.com",
         executor=executor,
-        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
     )
 
 
@@ -215,7 +215,7 @@ def test_client_lifecycle_flow(self):
             session_id_hex=session_id_hex,
             auth_provider=auth_provider,
             host_url="test-host.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -240,7 +240,7 @@ def test_disabled_telemetry_flow(self):
             session_id_hex=session_id_hex,
             auth_provider=None,
             host_url="test-host.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -260,7 +260,7 @@ def test_factory_error_handling(self):
                 session_id_hex=session_id,
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
-                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
             )
 
         # Should fall back to NoopTelemetryClient
@@ -279,7 +279,7 @@ def test_factory_shutdown_flow(self):
                 session_id_hex=session,
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
-                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
             )
 
         # Factory should be initialized
@@ -342,9 +342,7 @@ def _mock_ff_response(self, mock_requests_get, enabled: bool):
         mock_requests_get.return_value = mock_response
 
     @patch("databricks.sql.common.feature_flag.requests.get")
-    def test_telemetry_enabled_when_flag_is_true(
-        self, mock_requests_get, MockSession
-    ):
+    def test_telemetry_enabled_when_flag_is_true(self, mock_requests_get, MockSession):
         """Telemetry should be ON when enable_telemetry=True and server flag is 'true'."""
         self._mock_ff_response(mock_requests_get, enabled=True)
         mock_session_instance = MockSession.return_value
@@ -405,4 +403,4 @@ def test_telemetry_disabled_when_flag_request_fails(
         assert conn.telemetry_enabled is False
         mock_requests_get.assert_called_once()
         client = TelemetryClientFactory.get_telemetry_client("test-session-ff-fail")
-        assert isinstance(client, NoopTelemetryClient)
+        assert isinstance(client, NoopTelemetryClient)

tests/unit/test_telemetry_retry.py

@@ -51,7 +51,7 @@ def get_client(self, session_id, num_retries=3):
             session_id_hex=session_id,
             auth_provider=None,
             host_url="test.databricks.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)