Revert "cleaner HTTP client using requests.sessions"

This reverts commit b7a4677.

Author: varun-edachali-dbx

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

@@ -1,13 +1,13 @@
 import json
 import logging
 import ssl
+import urllib.parse
 import urllib.request
 from typing import Dict, Any, Optional, List, Tuple, Union
 from urllib.parse import urljoin
 
-import requests
-from requests.adapters import HTTPAdapter
-from requests.exceptions import RequestException, HTTPError, ConnectionError
+from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
+from urllib3.util import make_headers
 from urllib3.exceptions import MaxRetryError
 
 from databricks.sql.auth.authenticators import AuthProvider
@@ -23,46 +23,20 @@
 logger = logging.getLogger(__name__)
 
 
-class SSLContextAdapter(HTTPAdapter):
-    """
-    An HTTP adapter that uses a custom SSLContext to handle advanced SSL settings,
-    including client certificate key passwords.
-    """
-
-    def __init__(self, ssl_options: SSLOptions, **kwargs):
-        self.ssl_context = self._create_ssl_context(ssl_options)
-        super().__init__(**kwargs)
-
-    def _create_ssl_context(self, ssl_options: SSLOptions) -> ssl.SSLContext:
-        """
-        Build a custom SSLContext based on the provided SSLOptions.
-        """
-        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
-        if not ssl_options.tls_verify:
-            context.check_hostname = False
-            context.verify_mode = ssl.CERT_NONE
-        elif ssl_options.tls_trusted_ca_file:
-            context.load_verify_locations(cafile=ssl_options.tls_trusted_ca_file)
-        if ssl_options.tls_client_cert_file:
-            context.load_cert_chain(
-                certfile=ssl_options.tls_client_cert_file,
-                keyfile=ssl_options.tls_client_cert_key_file,
-                password=ssl_options.tls_client_cert_key_password,
-            )
-        return context
-
-    def init_poolmanager(self, *args, **kwargs):
-        kwargs["ssl_context"] = self.ssl_context
-        return super().init_poolmanager(*args, **kwargs)
-
-
 class SeaHttpClient:
     """
-    HTTP client for Statement Execution API (SEA), using the requests library.
+    HTTP client for Statement Execution API (SEA).
+
+    This client uses urllib3 for robust HTTP communication with retry policies
+    and connection pooling, similar to the Thrift HTTP client but simplified.
     """
 
     retry_policy: Union[DatabricksRetryPolicy, int]
-    _session: requests.Session
+    _pool: Optional[Union[HTTPConnectionPool, HTTPSConnectionPool]]
+    proxy_uri: Optional[str]
+    realhost: Optional[str]
+    realport: Optional[int]
+    proxy_auth: Optional[Dict[str, str]]
 
     def __init__(
         self,
@@ -74,16 +48,39 @@ def __init__(
         ssl_options: SSLOptions,
         **kwargs,
     ):
+        """
+        Initialize the SEA HTTP client.
+
+        Args:
+            server_hostname: Hostname of the Databricks server
+            port: Port number for the connection
+            http_path: HTTP path for the connection
+            http_headers: List of HTTP headers to include in requests
+            auth_provider: Authentication provider
+            ssl_options: SSL configuration options
+            **kwargs: Additional keyword arguments including retry policy settings
+        """
+
         self.server_hostname = server_hostname
         self.port = port or 443
+        self.http_path = http_path
         self.auth_provider = auth_provider
         self.ssl_options = ssl_options
-        self.scheme = "https"
-        self.base_url = f"{self.scheme}://{server_hostname}:{self.port}"
-        self._session = requests.Session()
+
+        # Build base URL
+        self.base_url = f"https://{server_hostname}:{self.port}"
+
+        # Parse URL for proxy handling
+        parsed_url = urllib.parse.urlparse(self.base_url)
+        self.scheme = parsed_url.scheme
+        self.host = parsed_url.hostname
+        self.port = parsed_url.port or (443 if self.scheme == "https" else 80)
+
+        # Setup headers
         self.headers: Dict[str, str] = dict(http_headers)
         self.headers.update({"Content-Type": "application/json"})
-        self._session.headers.update(self.headers)
+
+        # Extract retry policy settings
         self._retry_delay_min = kwargs.get("_retry_delay_min", 1.0)
         self._retry_delay_max = kwargs.get("_retry_delay_max", 60.0)
         self._retry_stop_after_attempts_count = kwargs.get(
@@ -94,36 +91,23 @@ def __init__(
         )
         self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0)
         self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
+
+        # Connection pooling settings
         self.max_connections = kwargs.get("max_connections", 10)
-        self._configure_proxies()
-        self.enable_v3_retries = kwargs.get("_enable_v3_retries", True)
-        self._configure_retries_and_ssl(**kwargs)
 
-    def _configure_proxies(self):
-        try:
-            proxy = urllib.request.getproxies().get(self.scheme)
-        except (KeyError, AttributeError):
-            proxy = None
-        else:
-            if self.server_hostname and urllib.request.proxy_bypass(
-                self.server_hostname
-            ):
-                proxy = None
-        if proxy:
-            self._session.proxies = {"http": proxy, "https": proxy}
+        # Setup retry policy
+        self.enable_v3_retries = kwargs.get("_enable_v3_retries", True)
 
-    def _configure_retries_and_ssl(self, **kwargs):
         if self.enable_v3_retries:
             urllib3_kwargs = {"allowed_methods": ["GET", "POST", "DELETE"]}
             _max_redirects = kwargs.get("_retry_max_redirects")
             if _max_redirects:
                 if _max_redirects > self._retry_stop_after_attempts_count:
                     logger.warning(
-                        "_retry_max_redirects > _retry_stop_after_attempts_count "
-                        "so it will have no effect!"
+                        "_retry_max_redirects > _retry_stop_after_attempts_count so it will have no affect!"
                     )
                 urllib3_kwargs["redirect"] = _max_redirects
-                self._session.max_redirects = _max_redirects
+
             self.retry_policy = DatabricksRetryPolicy(
                 delay_min=self._retry_delay_min,
                 delay_max=self._retry_delay_max,
@@ -133,34 +117,104 @@ def _configure_retries_and_ssl(self, **kwargs):
                 force_dangerous_codes=self.force_dangerous_codes,
                 urllib3_kwargs=urllib3_kwargs,
             )
-            retry_strategy = self.retry_policy
         else:
+            # Legacy behavior - no automatic retries
             logger.warning(
                 "Legacy retry behavior is enabled for this connection."
                 " This behaviour is not supported for the SEA backend."
             )
             self.retry_policy = 0
-            retry_strategy = 0
-        adapter = SSLContextAdapter(
-            ssl_options=self.ssl_options,
-            pool_connections=self.max_connections,
-            max_retries=retry_strategy,
-        )
-        self._session.mount("https://", adapter)
-        self._session.mount("http://", adapter)
+
+        # Handle proxy settings
+        try:
+            proxy = urllib.request.getproxies().get(self.scheme)
+        except (KeyError, AttributeError):
+            proxy = None
+        else:
+            if self.host and urllib.request.proxy_bypass(self.host):
+                proxy = None
+
+        if proxy:
+            parsed_proxy = urllib.parse.urlparse(proxy)
+            self.realhost = self.host
+            self.realport = 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_auth = None
+            self.proxy_uri = None
+
+        # Initialize connection pool
+        self._pool = None
+        self._open()
+
+    def _basic_proxy_auth_headers(self, proxy_parsed) -> Optional[Dict[str, str]]:
+        """Create basic auth headers for proxy if credentials are provided."""
+        if proxy_parsed is None or not proxy_parsed.username:
+            return None
+        ap = f"{urllib.parse.unquote(proxy_parsed.username)}:{urllib.parse.unquote(proxy_parsed.password)}"
+        return make_headers(proxy_basic_auth=ap)
+
+    def _open(self):
+        """Initialize the connection pool."""
+        pool_kwargs = {"maxsize": self.max_connections}
+
+        if self.scheme == "http":
+            pool_class = HTTPConnectionPool
+        else:  # https
+            pool_class = HTTPSConnectionPool
+            pool_kwargs.update(
+                {
+                    "cert_reqs": ssl.CERT_REQUIRED
+                    if self.ssl_options.tls_verify
+                    else ssl.CERT_NONE,
+                    "ca_certs": self.ssl_options.tls_trusted_ca_file,
+                    "cert_file": self.ssl_options.tls_client_cert_file,
+                    "key_file": self.ssl_options.tls_client_cert_key_file,
+                    "key_password": self.ssl_options.tls_client_cert_key_password,
+                }
+            )
+
+        if self.using_proxy():
+            proxy_manager = ProxyManager(
+                self.proxy_uri,
+                num_pools=1,
+                proxy_headers=self.proxy_auth,
+            )
+            self._pool = proxy_manager.connection_from_host(
+                host=self.realhost,
+                port=self.realport,
+                scheme=self.scheme,
+                pool_kwargs=pool_kwargs,
+            )
+        else:
+            self._pool = pool_class(self.host, self.port, **pool_kwargs)
 
     def close(self):
-        self._session.close()
+        """Close the connection pool."""
+        if self._pool:
+            self._pool.clear()
+
+    def using_proxy(self) -> bool:
+        """Check if proxy is being used (for compatibility with Thrift client)."""
+        return self.realhost is not None
 
     def set_retry_command_type(self, command_type: CommandType):
+        """Set the command type for retry policy decision making."""
         if isinstance(self.retry_policy, DatabricksRetryPolicy):
             self.retry_policy.command_type = command_type
 
     def start_retry_timer(self):
+        """Start the retry timer for duration-based retry limits."""
         if isinstance(self.retry_policy, DatabricksRetryPolicy):
             self.retry_policy.start_retry_timer()
 
     def _get_auth_headers(self) -> Dict[str, str]:
+        """Get authentication headers from the auth provider."""
         headers: Dict[str, str] = {}
         self.auth_provider.add_headers(headers)
         return headers
@@ -171,51 +225,91 @@ def _make_request(
         path: str,
         data: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
-        full_url = urljoin(self.base_url, path)
-        auth_headers = self._get_auth_headers()
+        """
+        Make an HTTP request to the SEA endpoint.
+
+        Args:
+            method: HTTP method (GET, POST, DELETE)
+            path: API endpoint path
+            data: Request payload data
+
+        Returns:
+            Dict[str, Any]: Response data parsed from JSON
+
+        Raises:
+            RequestError: If the request fails after retries
+        """
+
+        # Prepare headers
+        headers = {**self.headers, **self._get_auth_headers()}
+
+        # Prepare request body
+        body = json.dumps(data).encode("utf-8") if data else b""
+        if body:
+            headers["Content-Length"] = str(len(body))
+
+        # Set command type for retry policy
         command_type = self._get_command_type_from_path(path, method)
         self.set_retry_command_type(command_type)
         self.start_retry_timer()
-        logger.debug(f"Making {method} request to {full_url}")
+
+        logger.debug(f"Making {method} request to {path}")
+
+        # When v3 retries are enabled, urllib3 handles retries internally via DatabricksRetryPolicy
+        # When disabled, we let exceptions bubble up (similar to Thrift backend approach)
+        if self._pool is None:
+            raise RequestError("Connection pool not initialized", None)
+
         try:
-            with self._session.request(
+            response = self._pool.request(
                 method=method.upper(),
-                url=full_url,
-                json=data,
-                headers=auth_headers,
-            ) as response:
-                logger.debug(f"Response status: {response.status_code}")
-                response.raise_for_status()
-                return response.json()
-        except requests.exceptions.ConnectionError as e:
-            # Check if the first argument of the ConnectionError is a MaxRetryError
-            if e.args and isinstance(e.args[0], MaxRetryError):
-                # We want to raise the original MaxRetryError, not the wrapper
-                original_error = e.args[0]
-                logger.error(
-                    f"SEA HTTP request failed with MaxRetryError: {original_error}"
-                )
-                raise original_error
-            else:
-                logger.error(f"SEA HTTP request failed with ConnectionError: {e}")
-                raise RequestError("Error during request to server.", None, None, e)
-        except RequestException as e:
-            error_message = f"Error during request to server: {e}"
+                url=path,
+                body=body,
+                headers=headers,
+                preload_content=False,
+                retries=self.retry_policy,
+            )
+        except MaxRetryError as e:
+            # urllib3 MaxRetryError should bubble up for redirect tests to catch
+            logger.error(f"SEA HTTP request failed with MaxRetryError: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"SEA HTTP request failed with exception: {e}")
+            error_message = f"Error during request to server. {e}"
+            # Construct RequestError with proper 3-argument format (message, context, error)
             raise RequestError(error_message, None, None, e)
 
+        logger.debug(f"Response status: {response.status}")
+
+        # Handle successful responses
+        if 200 <= response.status < 300:
+            return response.json()
+
+        error_message = f"SEA HTTP request failed with status {response.status}"
+
+        raise RequestError(error_message, None)
+
     def _get_command_type_from_path(self, path: str, method: str) -> CommandType:
+        """
+        Determine the command type based on the API path and method.
+
+        This helps the retry policy make appropriate decisions for different
+        types of SEA operations.
+        """
         path = path.lower()
         method = method.upper()
+
         if "/statements" in path:
             if method == "POST" and path.endswith("/statements"):
                 return CommandType.EXECUTE_STATEMENT
             elif "/cancel" in path:
-                return CommandType.OTHER
+                return CommandType.OTHER  # Cancel operation
             elif method == "DELETE":
                 return CommandType.CLOSE_OPERATION
             elif method == "GET":
                 return CommandType.GET_OPERATION_STATUS
         elif "/sessions" in path:
             if method == "DELETE":
                 return CommandType.CLOSE_SESSION
+
         return CommandType.OTHER