Some more fixes and aligned tests

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/auth/auth.py

@@ -10,7 +10,7 @@
 from databricks.sql.auth.common import AuthType, ClientContext
 
 
-def get_auth_provider(cfg: ClientContext):
+def get_auth_provider(cfg: ClientContext, http_client):
     if cfg.credentials_provider:
         return ExternalAuthProvider(cfg.credentials_provider)
     elif cfg.auth_type == AuthType.AZURE_SP_M2M.value:
@@ -113,4 +113,4 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs)
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
     )
-    return get_auth_provider(cfg)
+    return get_auth_provider(cfg, http_client)

src/databricks/sql/auth/oauth.py

@@ -61,22 +61,6 @@ def refresh(self) -> Token:
         pass
 
 
-class IgnoreNetrcAuth(requests.auth.AuthBase):
-    """This auth method is a no-op.
-
-    We use it to force requestslib to not use .netrc to write auth headers
-    when making .post() requests to the oauth token endpoints, since these
-    don't require authentication.
-
-    In cases where .netrc is outdated or corrupt, these requests will fail.
-
-    See issue #121
-    """
-
-    def __call__(self, r):
-        return r
-
-
 class OAuthManager:
     def __init__(
         self,
@@ -103,7 +87,6 @@ def __fetch_well_known_config(self, hostname: str):
         known_config_url = self.idp_endpoint.get_openid_config_url(hostname)
 
         try:
-            from databricks.sql.common.unified_http_client import IgnoreNetrcAuth
             response = self.http_client.request('GET', url=known_config_url)
             # Convert urllib3 response to requests-like response for compatibility
             response.status_code = response.status
@@ -214,7 +197,6 @@ def __send_token_request(token_request_url, data):
             "Content-Type": "application/x-www-form-urlencoded",
         }
         # Use unified HTTP client
-        from databricks.sql.common.unified_http_client import IgnoreNetrcAuth
         response = self.http_client.request(
             'POST', url=token_request_url, body=data, headers=headers
         )

src/databricks/sql/backend/thrift_backend.py

@@ -105,6 +105,7 @@ def __init__(
         http_headers,
         auth_provider: AuthProvider,
         ssl_options: SSLOptions,
+        http_client=None,
         **kwargs,
     ):
         # Internal arguments in **kwargs:
@@ -145,10 +146,8 @@ def __init__(
         #  Number of threads for handling cloud fetch downloads. Defaults to 10
 
         logger.debug(
-            "ThriftBackend.__init__(server_hostname=%s, port=%s, http_path=%s)",
-            server_hostname,
-            port,
-            http_path,
+            "ThriftBackend.__init__(server_hostname=%s, port=%s, http_path=%s)"
+            % (server_hostname, port, http_path)
         )
 
         port = port or 443
@@ -177,8 +176,8 @@ def __init__(
         self._max_download_threads = kwargs.get("max_download_threads", 10)
 
         self._ssl_options = ssl_options
-
         self._auth_provider = auth_provider
+        self._http_client = http_client
 
         # Connector version 3 retry approach
         self.enable_v3_retries = kwargs.get("_enable_v3_retries", True)
@@ -1292,6 +1291,7 @@ def fetch_results(
             session_id_hex=self._session_id_hex,
             statement_id=command_id.to_hex_guid(),
             chunk_id=chunk_id,
+            http_client=self._http_client,
         )
 
         return (

src/databricks/sql/client.py

@@ -50,6 +50,9 @@
 from databricks.sql.session import Session
 from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId
 
+from databricks.sql.auth.common import ClientContext
+from databricks.sql.common.unified_http_client import UnifiedHttpClient
+
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TOpenSessionResp,
     TSparkParameter,
@@ -251,10 +254,14 @@ def read(self) -> Optional[OAuthToken]:
             "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
         )
 
+        client_context = self._build_client_context(server_hostname, **kwargs)
+        http_client = UnifiedHttpClient(client_context)
+
         try:
             self.session = Session(
                 server_hostname,
                 http_path,
+                http_client,
                 http_headers,
                 session_configuration,
                 catalog,
@@ -270,6 +277,7 @@ def read(self) -> Optional[OAuthToken]:
                 host_url=server_hostname,
                 http_path=http_path,
                 port=kwargs.get("_port", 443),
+                http_client=http_client,
                 user_agent=self.session.useragent_header
                 if hasattr(self, "session")
                 else None,
@@ -342,6 +350,46 @@ def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
 
         return value
 
+    def _build_client_context(self, server_hostname: str, **kwargs):
+        """Build ClientContext for HTTP client configuration."""
+        from databricks.sql.auth.common import ClientContext
+        from databricks.sql.types import SSLOptions
+        
+        # Extract SSL options
+        ssl_options = SSLOptions(
+            tls_verify=not kwargs.get("_tls_no_verify", False),
+            tls_verify_hostname=kwargs.get("_tls_verify_hostname", True),
+            tls_trusted_ca_file=kwargs.get("_tls_trusted_ca_file"),
+            tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
+            tls_client_cert_key_file=kwargs.get("_tls_client_cert_key_file"),
+            tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
+        )
+        
+        # Build user agent
+        user_agent_entry = kwargs.get("user_agent_entry", "")
+        if user_agent_entry:
+            user_agent = f"PyDatabricksSqlConnector/{__version__} ({user_agent_entry})"
+        else:
+            user_agent = f"PyDatabricksSqlConnector/{__version__}"
+        
+        return ClientContext(
+            hostname=server_hostname,
+            ssl_options=ssl_options,
+            socket_timeout=kwargs.get("_socket_timeout"),
+            retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count", 30),
+            retry_delay_min=kwargs.get("_retry_delay_min", 1.0),
+            retry_delay_max=kwargs.get("_retry_delay_max", 60.0),
+            retry_stop_after_attempts_duration=kwargs.get("_retry_stop_after_attempts_duration", 900.0),
+            retry_delay_default=kwargs.get("_retry_delay_default", 1.0),
+            retry_dangerous_codes=kwargs.get("_retry_dangerous_codes", []),
+            http_proxy=kwargs.get("_http_proxy"),
+            proxy_username=kwargs.get("_proxy_username"),
+            proxy_password=kwargs.get("_proxy_password"),
+            pool_connections=kwargs.get("_pool_connections", 1),
+            pool_maxsize=kwargs.get("_pool_maxsize", 1),
+            user_agent=user_agent,
+        )
+
     # The ideal return type for this method is perhaps Self, but that was not added until 3.11, and we support pre-3.11 pythons, currently.
     def __enter__(self) -> "Connection":
         return self

src/databricks/sql/session.py

@@ -22,6 +22,7 @@ def __init__(
         self,
         server_hostname: str,
         http_path: str,
+        http_client: UnifiedHttpClient,
         http_headers: Optional[List[Tuple[str, str]]] = None,
         session_configuration: Optional[Dict[str, Any]] = None,
         catalog: Optional[str] = None,
@@ -75,9 +76,8 @@ def __init__(
             tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
         )
 
-        # Create HTTP client configuration and unified HTTP client
-        self.client_context = self._build_client_context(server_hostname, **kwargs)
-        self.http_client = UnifiedHttpClient(self.client_context)
+        # Use the provided HTTP client (created in Connection)
+        self.http_client = http_client
 
         # Create auth provider with HTTP client context
         self.auth_provider = get_python_sql_connector_auth_provider(
@@ -95,26 +95,6 @@ def __init__(
 
         self.protocol_version = None
 
-    def _build_client_context(self, server_hostname: str, **kwargs) -> ClientContext:
-        """Build ClientContext with HTTP configuration from kwargs."""
-        return ClientContext(
-            hostname=server_hostname,
-            ssl_options=self.ssl_options,
-            socket_timeout=kwargs.get("_socket_timeout"),
-            retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count"),
-            retry_delay_min=kwargs.get("_retry_delay_min"),
-            retry_delay_max=kwargs.get("_retry_delay_max"),
-            retry_stop_after_attempts_duration=kwargs.get("_retry_stop_after_attempts_duration"),
-            retry_delay_default=kwargs.get("_retry_delay_default"),
-            retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"),
-            http_proxy=kwargs.get("http_proxy"),
-            proxy_username=kwargs.get("proxy_username"),
-            proxy_password=kwargs.get("proxy_password"),
-            pool_connections=kwargs.get("pool_connections"),
-            pool_maxsize=kwargs.get("pool_maxsize"),
-            user_agent=self.useragent_header,
-        )
-
     def _create_backend(
         self,
         server_hostname: str,
@@ -142,6 +122,7 @@ def _create_backend(
             "http_headers": all_headers,
             "auth_provider": auth_provider,
             "ssl_options": self.ssl_options,
+            "http_client": self.http_client,
             "_use_arrow_native_complex_types": _use_arrow_native_complex_types,
             **kwargs,
         }

src/databricks/sql/telemetry/telemetry_client.py

@@ -3,7 +3,6 @@
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Optional, TYPE_CHECKING
-from databricks.sql.common.http import TelemetryHttpClient
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -38,6 +37,8 @@
 from databricks.sql.telemetry.utils import BaseTelemetryClient
 from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
 
+from src.databricks.sql.common.unified_http_client import UnifiedHttpClient
+
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
 
@@ -511,7 +512,6 @@ def close(session_id_hex):
                 try:
                     TelemetryClientFactory._stop_flush_thread()
                     TelemetryClientFactory._executor.shutdown(wait=True)
-                    TelemetryHttpClient.close()
                 except Exception as e:
                     logger.debug("Failed to shutdown thread pool executor: %s", e)
                 TelemetryClientFactory._executor = None
@@ -524,6 +524,7 @@ def connection_failure_log(
         host_url: str,
         http_path: str,
         port: int,
+        http_client: UnifiedHttpClient,
         user_agent: Optional[str] = None,
     ):
         """Send error telemetry when connection creation fails, without requiring a session"""
@@ -536,6 +537,7 @@ def connection_failure_log(
             auth_provider=None,
             host_url=host_url,
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
+            http_client=http_client,
         )
 
         telemetry_client = TelemetryClientFactory.get_telemetry_client(