unify ssl proxy

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

Author: vikrantpuppala

src/databricks/sql/auth/thrift_http_client.py

@@ -15,6 +15,10 @@
 from urllib3.util import make_headers
 from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
 from databricks.sql.types import SSLOptions
+from databricks.sql.common.http_utils import (
+    detect_and_parse_proxy,
+    create_connection_pool,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -58,25 +62,22 @@ def __init__(
             self.path = parsed.path
             if parsed.query:
                 self.path += "?%s" % parsed.query
-        try:
-            proxy = urllib.request.getproxies()[self.scheme]
-        except KeyError:
-            proxy = None
-        else:
-            if urllib.request.proxy_bypass(self.host):
-                proxy = None
-        if proxy:
-            parsed = urllib.parse.urlparse(proxy)
+        
+        # Handle proxy settings using shared utility
+        proxy_uri, proxy_auth = detect_and_parse_proxy(self.scheme, self.host)
+        
+        if proxy_uri:
+            parsed_proxy = urllib.parse.urlparse(proxy_uri)
 
             # realhost and realport are the host and port of the actual request
             self.realhost = self.host
             self.realport = self.port
 
             # this is passed to ProxyManager
-            self.proxy_uri: str = proxy
-            self.host = parsed.hostname
-            self.port = parsed.port
-            self.proxy_auth = self.basic_proxy_auth_headers(parsed)
+            self.proxy_uri: str = proxy_uri
+            self.host = parsed_proxy.hostname
+            self.port = parsed_proxy.port
+            self.proxy_auth = proxy_auth
         else:
             self.realhost = self.realport = self.proxy_auth = None
 
@@ -105,40 +106,17 @@ def startRetryTimer(self):
         self.retry_policy and self.retry_policy.start_retry_timer()
 
     def open(self):
-
-        # self.__pool replaces the self.__http used by the original THttpClient
-        _pool_kwargs = {"maxsize": self.max_connections}
-
-        if self.scheme == "http":
-            pool_class = HTTPConnectionPool
-        elif self.scheme == "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)
+        """Initialize the connection pool using shared utility."""
+        self.__pool = create_connection_pool(
+            scheme=self.scheme,
+            host=self.realhost if self.using_proxy() else self.host,
+            port=self.realport if self.using_proxy() else self.port,
+            ssl_options=self._ssl_options,
+            proxy_uri=getattr(self, 'proxy_uri', None),
+            proxy_headers=self.proxy_auth,
+            retry_policy=self.retry_policy,
+            max_connections=self.max_connections,
+        )
 
     def close(self):
         self.__resp and self.__resp.drain_conn()
@@ -204,15 +182,9 @@ def flush(self):
             )
         )
 
-    @staticmethod
-    def basic_proxy_auth_headers(proxy):
-        if proxy is None or not proxy.username:
-            return None
-        ap = "%s:%s" % (
-            urllib.parse.unquote(proxy.username),
-            urllib.parse.unquote(proxy.password),
-        )
-        return make_headers(proxy_basic_auth=ap)
+    def using_proxy(self) -> bool:
+        """Check if proxy is being used."""
+        return self.realhost is not None
 
     def set_retry_command_type(self, value: CommandType):
         """Pass the provided CommandType to the retry policy"""

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

@@ -15,6 +15,11 @@
 from databricks.sql.exc import (
     RequestError,
 )
+from databricks.sql.common.http_utils import (
+    detect_and_parse_proxy,
+    create_retry_policy_from_kwargs,
+    create_connection_pool,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -121,27 +126,17 @@ def __init__(
             )
             self.retry_policy = 0
 
-        # Handle proxy settings
-        try:
-            # returns a dictionary of scheme -> proxy server URL mappings.
-            # https://docs.python.org/3/library/urllib.request.html#urllib.request.getproxies
-            proxy = urllib.request.getproxies().get(self.scheme)
-        except (KeyError, AttributeError):
-            # No proxy found or getproxies() failed - disable proxy
-            proxy = None
-        else:
-            # Proxy found, but check if this host should bypass proxy
-            if self.host and urllib.request.proxy_bypass(self.host):
-                proxy = None  # Host bypasses proxy per system rules
-
-        if proxy:
-            parsed_proxy = urllib.parse.urlparse(proxy)
+        # Handle proxy settings using shared utility
+        proxy_uri, proxy_auth = detect_and_parse_proxy(self.scheme, self.host)
+        
+        if proxy_uri:
+            parsed_proxy = urllib.parse.urlparse(proxy_uri)
             self.proxy_host = self.host
             self.proxy_port = self.port
-            self.proxy_uri = proxy
+            self.proxy_uri = proxy_uri
             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)
+            self.proxy_auth = proxy_auth
         else:
             self.proxy_host = None
             self.proxy_port = None
@@ -152,47 +147,18 @@ def __init__(
         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.proxy_host,
-                port=self.proxy_port,
-                scheme=self.scheme,
-                pool_kwargs=pool_kwargs,
-            )
-        else:
-            self._pool = pool_class(self.host, self.port, **pool_kwargs)
+        """Initialize the connection pool using shared utility."""
+        self._pool = create_connection_pool(
+            scheme=self.scheme,
+            host=self.proxy_host if self.using_proxy() else self.host,
+            port=self.proxy_port if self.using_proxy() else self.port,
+            ssl_options=self.ssl_options,
+            proxy_uri=self.proxy_uri,
+            proxy_headers=self.proxy_auth,
+            retry_policy=self.retry_policy,
+            max_connections=self.max_connections,
+        )
 
     def close(self):
         """Close the connection pool."""

src/databricks/sql/common/unified_http_client.py

@@ -12,6 +12,11 @@
 from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 from databricks.sql.exc import RequestError
 from databricks.sql.common.http import HttpMethod
+from databricks.sql.common.http_utils import (
+    detect_and_parse_proxy,
+    create_retry_policy_from_kwargs,
+    create_connection_pool,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -39,37 +44,6 @@ def __init__(self, client_context):
 
     def _setup_pool_manager(self):
         """Set up the urllib3 PoolManager with configuration from ClientContext."""
-
-        # SSL context setup
-        ssl_context = None
-        if self.config.ssl_options:
-            ssl_context = ssl.create_default_context()
-
-            # Configure SSL verification
-            if not self.config.ssl_options.tls_verify:
-                ssl_context.check_hostname = False
-                ssl_context.verify_mode = ssl.CERT_NONE
-            elif not self.config.ssl_options.tls_verify_hostname:
-                ssl_context.check_hostname = False
-                ssl_context.verify_mode = ssl.CERT_REQUIRED
-
-            # Load custom CA file if specified
-            if self.config.ssl_options.tls_trusted_ca_file:
-                ssl_context.load_verify_locations(
-                    self.config.ssl_options.tls_trusted_ca_file
-                )
-
-            # Load client certificate if specified
-            if (
-                self.config.ssl_options.tls_client_cert_file
-                and self.config.ssl_options.tls_client_cert_key_file
-            ):
-                ssl_context.load_cert_chain(
-                    self.config.ssl_options.tls_client_cert_file,
-                    self.config.ssl_options.tls_client_cert_key_file,
-                    self.config.ssl_options.tls_client_cert_key_password,
-                )
-
         # Create retry policy
         self._retry_policy = DatabricksRetryPolicy(
             delay_min=self.config.retry_delay_min,
@@ -85,32 +59,32 @@ def _setup_pool_manager(self):
         self._retry_policy._command_type = None
         self._retry_policy._retry_start_time = None
 
-        # Common pool manager kwargs
-        pool_kwargs = {
-            "num_pools": self.config.pool_connections,
-            "maxsize": self.config.pool_maxsize,
-            "retries": self._retry_policy,
-            "timeout": urllib3.Timeout(
-                connect=self.config.socket_timeout, read=self.config.socket_timeout
+        
+        parsed_url = urllib.parse.urlparse(self.config.hostname)
+        self.scheme = parsed_url.scheme
+        # Detect proxy using shared utility
+        proxy_uri, proxy_headers = detect_and_parse_proxy(self.scheme, self.config.hostname)
+        
+
+        # Create pool 
+        additional_kwargs = {}
+        if self.config.socket_timeout:
+            additional_kwargs["timeout"] = urllib3.Timeout(
+                connect=self.config.socket_timeout, 
+                read=self.config.socket_timeout
             )
-            if self.config.socket_timeout
-            else None,
-            "ssl_context": ssl_context,
-        }
-
-        # Create proxy or regular pool manager
-        if self.config.http_proxy:
-            proxy_headers = None
-            if self.config.proxy_username and self.config.proxy_password:
-                proxy_headers = make_headers(
-                    proxy_basic_auth=f"{self.config.proxy_username}:{self.config.proxy_password}"
-                )
-
-            self._pool_manager = ProxyManager(
-                self.config.http_proxy, proxy_headers=proxy_headers, **pool_kwargs
-            )
-        else:
-            self._pool_manager = PoolManager(**pool_kwargs)
+        
+        self._pool_manager = create_connection_pool(
+            scheme=self.scheme,
+            host=self.config.hostname,
+            port=443,
+            ssl_options=self.config.ssl_options,
+            proxy_uri=proxy_uri,
+            proxy_headers=proxy_headers,
+            retry_policy=self._retry_policy,
+            max_connections=self.config.pool_maxsize,
+            **additional_kwargs
+        )
 
     def _prepare_headers(
         self, headers: Optional[Dict[str, str]] = None