Skip to content

Commit 6a1274f

Browse files
fix type annotations
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent cc48caf commit 6a1274f

File tree

1 file changed

+21
-4
lines changed

1 file changed

+21
-4
lines changed

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

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,13 @@ class SeaHttpClient:
2525
and connection pooling, similar to the Thrift HTTP client but simplified.
2626
"""
2727

28+
retry_policy: Union[DatabricksRetryPolicy, int]
29+
_pool: Optional[Union[HTTPConnectionPool, HTTPSConnectionPool]]
30+
proxy_uri: Optional[str]
31+
realhost: Optional[str]
32+
realport: Optional[int]
33+
proxy_auth: Optional[Dict[str, str]]
34+
2835
def __init__(
2936
self,
3037
server_hostname: str,
@@ -105,7 +112,7 @@ def __init__(
105112
except (KeyError, AttributeError):
106113
proxy = None
107114
else:
108-
if urllib.request.proxy_bypass(self.host):
115+
if self.host and urllib.request.proxy_bypass(self.host):
109116
proxy = None
110117

111118
if proxy:
@@ -114,10 +121,13 @@ def __init__(
114121
self.realport = self.port
115122
self.proxy_uri = proxy
116123
self.host = parsed_proxy.hostname
117-
self.port = parsed_proxy.port
124+
self.port = parsed_proxy.port or (443 if self.scheme == "https" else 80)
118125
self.proxy_auth = self._basic_proxy_auth_headers(parsed_proxy)
119126
else:
120-
self.realhost = self.realport = self.proxy_auth = self.proxy_uri = None
127+
self.realhost = None
128+
self.realport = None
129+
self.proxy_auth = None
130+
self.proxy_uri = None
121131

122132
# Initialize connection pool
123133
self._pool = None
@@ -150,7 +160,7 @@ def _open(self):
150160
}
151161
)
152162

153-
if self.proxy_uri:
163+
if self.using_proxy():
154164
proxy_manager = ProxyManager(
155165
self.proxy_uri,
156166
num_pools=1,
@@ -170,6 +180,10 @@ def close(self):
170180
if self._pool:
171181
self._pool.clear()
172182

183+
def using_proxy(self) -> bool:
184+
"""Check if proxy is being used (for compatibility with Thrift client)."""
185+
return self.realhost is not None
186+
173187
def set_retry_command_type(self, command_type: CommandType):
174188
"""Set the command type for retry policy decision making."""
175189
if isinstance(self.retry_policy, DatabricksRetryPolicy):
@@ -232,6 +246,9 @@ def _make_request(
232246

233247
# When v3 retries are enabled, urllib3 handles retries internally via DatabricksRetryPolicy
234248
# When disabled, we let exceptions bubble up (similar to Thrift backend approach)
249+
if self._pool is None:
250+
raise RequestError("Connection pool not initialized", None)
251+
235252
response = self._pool.request(
236253
method=method.upper(),
237254
url=url,

0 commit comments

Comments
 (0)