diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 5db76de9..f7c22a1e 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -78,7 +78,7 @@ def __init__( self.proxy_uri: str = proxy self.host = parsed.hostname self.port = parsed.port - self.proxy_auth = self.basic_proxy_auth_header(parsed) + self.proxy_auth = self.basic_proxy_auth_headers(parsed) else: self.realhost = self.realport = self.proxy_auth = None @@ -167,7 +167,7 @@ def flush(self): } if self.using_proxy() and self.scheme == "http" and self.proxy_auth is not None: - headers["Proxy-Authorization" : self.proxy_auth] + headers.update(self.proxy_auth) if self.__custom_headers: custom_headers = {key: val for key, val in self.__custom_headers.items()} @@ -190,7 +190,7 @@ def flush(self): self.headers = self.__resp.headers @staticmethod - def basic_proxy_auth_header(proxy): + def basic_proxy_auth_headers(proxy): if proxy is None or not proxy.username: return None ap = "%s:%s" % ( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index ee4fc4b7..4bcf84d2 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -167,11 +167,12 @@ def test_proxy_headers_are_set(self): parsed_proxy = urlparse(fake_proxy_spec) try: - result = THttpClient.basic_proxy_auth_header(parsed_proxy) + result = THttpClient.basic_proxy_auth_headers(parsed_proxy) except TypeError as e: assert False - assert isinstance(result, type(str())) + assert isinstance(result, type(dict())) + assert isinstance(result.get('proxy-authorization'), type(str())) @patch("databricks.sql.auth.thrift_http_client.THttpClient") @patch("databricks.sql.thrift_backend.create_default_context")