Adding support for http/s proxies

NERSC · Oct 9, 2024 · 07eff60 · 07eff60
1 parent 5c8bc9d
commit 07eff60
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 6 deletions.
diff --git a/scripts/run.py b/scripts/run.py
@@ -34,6 +34,7 @@ def run_unasync():
         "AsyncJobSacct": "JobSacct",
         "AsyncJobSqueue": "JobSqueue",
         "AsyncOAuth2Client": "OAuth2Client",
+        "AsyncHTTPTransport": "HTTPTransport",
         "aclose": "close",
         "_ASYNC_SLEEP": "_SLEEP",
     }

diff --git a/src/sfapi_client/_async/client.py b/src/sfapi_client/_async/client.py
@@ -6,6 +6,7 @@
 from authlib.integrations.httpx_client.oauth2_client import AsyncOAuth2Client
 from authlib.oauth2.rfc7523 import PrivateKeyJWT
 import httpx
+import urllib
 import tenacity
 from authlib.jose import JsonWebKey
 
@@ -260,6 +261,11 @@ async def __aenter__(self):
 
     async def _http_client(self):
         headers = {"accept": "application/json"}
+        # Get the users http/https proxies to use for the client
+        proxy_mounts = {
+            f"{_type}://": httpx.AsyncHTTPTransport(proxy=_proxy)
+            for _type, _proxy in urllib.request.getproxies().items()
+        }
         # If we have a client_id then we need to use the OAuth2 client
         if self._client_id is not None:
             if self.__http_client is None:
@@ -272,6 +278,7 @@ async def _http_client(self):
                     token_endpoint=self._token_url,
                     timeout=10.0,
                     headers=headers,
+                    mounts=proxy_mounts
                 )
 
                 await self.__http_client.fetch_token()
@@ -285,7 +292,7 @@ async def _http_client(self):
             if self._access_token is not None:
                 headers.update({"Authorization": f"Bearer {self._access_token}"})
 
-            self.__http_client = httpx.AsyncClient(headers=headers)
+            self.__http_client = httpx.AsyncClient(headers=headers, mounts=proxy_mounts)
 
         return self.__http_client
 

diff --git a/src/sfapi_client/_sync/client.py b/src/sfapi_client/_sync/client.py
@@ -6,6 +6,7 @@
 from authlib.integrations.httpx_client.oauth2_client import OAuth2Client
 from authlib.oauth2.rfc7523 import PrivateKeyJWT
 import httpx
+import urllib
 import tenacity
 from authlib.jose import JsonWebKey
 
@@ -260,6 +261,11 @@ def __enter__(self):
 
     def _http_client(self):
         headers = {"accept": "application/json"}
+        # Get the users http/https proxies to use for the client
+        proxy_mounts = {
+            f"{_type}://": httpx.HTTPTransport(proxy=_proxy)
+            for _type, _proxy in urllib.request.getproxies().items()
+        }
         # If we have a client_id then we need to use the OAuth2 client
         if self._client_id is not None:
             if self.__http_client is None:
@@ -272,6 +278,7 @@ def _http_client(self):
                     token_endpoint=self._token_url,
                     timeout=10.0,
                     headers=headers,
+                    mounts=proxy_mounts
                 )
 
                 self.__http_client.fetch_token()
@@ -285,7 +292,7 @@ def _http_client(self):
             if self._access_token is not None:
                 headers.update({"Authorization": f"Bearer {self._access_token}"})
 
-            self.__http_client = httpx.Client(headers=headers)
+            self.__http_client = httpx.Client(headers=headers, mounts=proxy_mounts)
 
         return self.__http_client
 
@@ -383,10 +390,7 @@ def get(self, url: str, params: Dict[str, Any] = {}) -> httpx.Response:
         stop=tenacity.stop_after_attempt(MAX_RETRY),
     )
     def post(
-        self,
-        url: str,
-        data: Dict[str, Any] = None,
-        json: Dict[str, Any] = None
+        self, url: str, data: Dict[str, Any] = None, json: Dict[str, Any] = None
     ) -> httpx.Response:
         client = self._http_client()