fix mypy

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

Author: vikrantpuppala

src/databricks/sql/auth/thrift_http_client.py

@@ -23,6 +23,11 @@
 
 
 class THttpClient(thrift.transport.THttpClient.THttpClient):
+    realhost: Optional[str]
+    realport: Optional[int]
+    proxy_uri: Optional[str]
+    proxy_auth: Optional[Dict[str, str]]
+
     def __init__(
         self,
         auth_provider,
@@ -75,7 +80,7 @@ def __init__(
             self.realhost = self.host
             self.realport = self.port
             # this is passed to ProxyManager
-            self.proxy_uri: str = proxy_uri
+            self.proxy_uri = proxy_uri
             self.host = parsed_proxy.hostname
             self.port = parsed_proxy.port
             self.proxy_auth = proxy_auth

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

@@ -33,9 +33,9 @@ class SeaHttpClient:
     retry_policy: Union[DatabricksRetryPolicy, int]
     _pool: Optional[Union[HTTPConnectionPool, HTTPSConnectionPool]]
     proxy_uri: Optional[str]
-    proxy_host: Optional[str]
-    proxy_port: Optional[int]
     proxy_auth: Optional[Dict[str, str]]
+    realhost: Optional[str]
+    realport: Optional[int]
 
     def __init__(
         self,
@@ -172,8 +172,8 @@ def _open(self):
                 proxy_headers=self.proxy_auth,
             )
             self._pool = proxy_manager.connection_from_host(
-                host=self.proxy_host,
-                port=self.proxy_port,
+                host=self.realhost,
+                port=self.realport,
                 scheme=self.scheme,
                 pool_kwargs=pool_kwargs,
             )

src/databricks/sql/common/http_utils.py

@@ -60,7 +60,7 @@ def detect_and_parse_proxy(
     return proxy, proxy_headers
 
 
-def _generate_negotiate_headers(proxy_hostname: str) -> Optional[Dict[str, str]]:
+def _generate_negotiate_headers(proxy_hostname: Optional[str]) -> Optional[Dict[str, str]]:
     """Generate Kerberos/SPNEGO authentication headers"""
     try:
         from requests_kerberos import HTTPKerberosAuth

src/databricks/sql/common/unified_http_client.py

@@ -231,7 +231,7 @@ def request_context(
             **kwargs: Additional arguments passed to urllib3 request
 
         Yields:
-            urllib3.HTTPResponse: The HTTP response object
+            urllib3.BaseHTTPResponse: The HTTP response object
         """
         logger.debug(
             "Making %s request to %s", method, urllib.parse.urlparse(url).netloc
@@ -268,7 +268,7 @@ def request(
         url: str,
         headers: Optional[Dict[str, str]] = None,
         **kwargs,
-    ) -> urllib3.HTTPResponse:
+    ) -> urllib3.BaseHTTPResponse:
         """
         Make an HTTP request.
 
@@ -279,12 +279,12 @@ def request(
             **kwargs: Additional arguments passed to urllib3 request
 
         Returns:
-            urllib3.HTTPResponse: The HTTP response object with data and metadata pre-loaded
+            urllib3.BaseHTTPResponse: The HTTP response object with data and metadata pre-loaded
         """
         with self.request_context(method, url, headers=headers, **kwargs) as response:
             # Read the response data to ensure it's available after context exit
             # Note: status and headers remain accessible after close(), only data needs caching
-            response._body = response.data
+            response.read()
             return response
 
     def using_proxy(self) -> bool: