unify ssl proxy

Signed-off-by: Vikrant Puppala <[email protected]>

Author: vikrantpuppala

src/databricks/sql/auth/thrift_http_client.py

@@ -68,11 +68,9 @@ def __init__(
 
         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_uri
             self.host = parsed_proxy.hostname
@@ -112,7 +110,7 @@ def open(self):
             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_uri=self.proxy_uri,
             proxy_headers=self.proxy_auth,
             retry_policy=self.retry_policy,
             max_connections=self.max_connections,

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

@@ -131,17 +131,14 @@ def __init__(
 
         if proxy_uri:
             parsed_proxy = urllib.parse.urlparse(proxy_uri)
-            self.proxy_host = self.host
-            self.proxy_port = self.port
+            self.realhost = self.host
+            self.realport = self.port
             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 = proxy_auth
         else:
-            self.proxy_host = None
-            self.proxy_port = None
-            self.proxy_auth = None
-            self.proxy_uri = None
+            self.realhost = self.realport = self.proxy_auth = None
 
         # Initialize connection pool
         self._pool = None
@@ -151,8 +148,8 @@ def _open(self):
         """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,
+            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=self.proxy_uri,
             proxy_headers=self.proxy_auth,

src/databricks/sql/common/unified_http_client.py

@@ -14,7 +14,6 @@
 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,
 )
 
@@ -62,9 +61,24 @@ def _setup_pool_manager(self):
 
         parsed_url = urllib.parse.urlparse(self.config.hostname)
         self.scheme = parsed_url.scheme
+        self.host = parsed_url.hostname
+        self.port = parsed_url.port
+
         # Detect proxy using shared utility
-        proxy_uri, proxy_headers = detect_and_parse_proxy(self.scheme, self.config.hostname)
+        proxy_uri, proxy_auth = detect_and_parse_proxy(self.scheme, self.config.hostname)
 
+        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_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
 
         # Create pool 
         additional_kwargs = {}
@@ -76,11 +90,11 @@ def _setup_pool_manager(self):
 
         self._pool_manager = create_connection_pool(
             scheme=self.scheme,
-            host=self.config.hostname,
-            port=443,
+            host=self.realhost if self.using_proxy() else self.host,
+            port=self.realport if self.using_proxy() else self.port,
             ssl_options=self.config.ssl_options,
             proxy_uri=proxy_uri,
-            proxy_headers=proxy_headers,
+            proxy_headers=proxy_auth,
             retry_policy=self._retry_policy,
             max_connections=self.config.pool_maxsize,
             **additional_kwargs