Fix: Eliminate uses of requests library (#7633)

.mypy.ini

@@ -73,6 +73,7 @@ modules =
     azul.field_type,
     azul.source,
     scripts.kibana_proxy,
+    urllib3_mock,
 
 
 packages =

scripts/request_flooder.py

@@ -11,11 +11,12 @@
 import sys
 import time
 
-import requests
-
 from azul.args import (
     AzulArgumentHelpFormatter,
 )
+from azul.http import (
+    http_client,
+)
 from azul.lib import (
     R,
 )
@@ -57,10 +58,6 @@ def parse_args(argv):
                         type=int,
                         default='300',
                         help='Total duration of the test in seconds.')
-    parser.add_argument('--log-headers',
-                        default=False,
-                        action='store_true',
-                        help='Include response headers in log output')
     args = parser.parse_args(argv)
     args.method = args.method.upper()
     assert args.method in ['HEAD', 'GET', 'PUT'], R(
@@ -75,16 +72,12 @@ def parse_args(argv):
     return args
 
 
-def request_url(method: str, url: str, log_headers: bool) -> int:
-    log.info('Making %s request to %r', method, url)
-    start_time = time.time()
-    response = requests.request(method=method, url=url)
-    duration = time.time() - start_time
-    if log_headers:
-        log.info('… with response headers %r', response.headers)
-    log.info('Got %i response after %.3fs from %s to %s',
-             response.status_code, duration, method, url)
-    return response.status_code
+http = http_client(log=log)
+
+
+def request_url(method: str, url: str) -> int:
+    response = http.request(method=method, url=url)
+    return response.status
 
 
 def main(argv):
@@ -98,7 +91,7 @@ def main(argv):
         end_time = start_time + args.duration
         while time.time() < end_time:
             time.sleep(sleep_delay)
-            futures.append(tpe.submit(request_url, args.method, args.url, args.log_headers))
+            futures.append(tpe.submit(request_url, args.method, args.url))
         for f in as_completed(futures):
             assert f.result() in [200, 429]
 

src/azul/health.py

@@ -26,7 +26,7 @@
 from furl import (
     furl,
 )
-import requests
+import urllib3
 
 from azul import (
     CatalogName,
@@ -40,6 +40,10 @@
 from azul.deployment import (
     aws,
 )
+from azul.http import (
+    HasCachedHttpClient,
+    raise_on_status,
+)
 from azul.lib import (
     R,
     cache,
@@ -174,7 +178,7 @@ def _make_response(self, body: JSON) -> Response:
 
 
 @attr.s(frozen=True, kw_only=True, auto_attribs=True)
-class Health:
+class Health(HasCachedHttpClient):
     """
     Encapsulates information about the health status of an Azul deployment. All
     aspects of health are exposed as lazily loaded properties. Instantiating the
@@ -260,14 +264,10 @@ def progress(self) -> JSON:
     def _api_endpoint(self, entity_type: str) -> JSON:
         relative_url = furl(path=('index', entity_type), args={'size': '1'})
         url = str(config.service_endpoint.join(relative_url))
-        log.info('Making HEAD request to %s', url)
-        start = time.time()
-        response = requests.api.head(url)
-        log.info('Got %s response after %.3fs from HEAD request to %s',
-                 response.status_code, time.time() - start, url)
+        response = self._http_client.request('HEAD', url)
         try:
-            response.raise_for_status()
-        except requests.exceptions.HTTPError as e:
+            raise_on_status(response)
+        except urllib3.exceptions.HTTPError as e:
             return {'up': False, 'error': repr(e)}
         else:
             return {'up': True}
@@ -298,9 +298,8 @@ def _lambda(self, lambda_name) -> JSON:
         try:
             url = config.lambda_endpoint(lambda_name).set(path='/health/basic',
                                                           args={'catalog': self.catalog})
-            log.info('Requesting %r', url)
-            response = requests.api.get(str(url))
-            response.raise_for_status()
+            response = self._http_client.request('GET', str(url))
+            raise_on_status(response)
             up = response.json()['up']
         except Exception as e:
             return {

src/azul/http.py

@@ -174,6 +174,12 @@ def http_client(log: logging.Logger | None = None) -> HttpClient:
     return StatusRetryHttpClient(client)
 
 
+def raise_on_status(response: urllib3.BaseHTTPResponse) -> None:
+    if not 200 <= response.status <= 399:
+        msg = f'{response.reason} for url: {response.url}'
+        raise urllib3.exceptions.HTTPError(msg)
+
+
 class LimitedTimeoutException(Exception):
 
     def __init__(self, url: furl, timeout: float):

src/azul/plugins/repository/dss/__init__.py

@@ -18,7 +18,6 @@
 from more_itertools import (
     one,
 )
-import requests
 
 from azul import (
     config,
@@ -31,6 +30,7 @@
 )
 from azul.http import (
     HasCachedHttpClient,
+    raise_on_status,
 )
 from azul.indexer import (
     SourcedBundleFQID,
@@ -208,7 +208,7 @@ def validate_version(self, version: str) -> None:
         parse_dcp2_version(version)
 
 
-class DSSFileDownload(RepositoryFileDownload):
+class DSSFileDownload(RepositoryFileDownload, HasCachedHttpClient):
     _location: str | None = None
     _retry_after: int | None = None
 
@@ -222,8 +222,8 @@ def update(self, authentication: Authentication | None) -> None:
                                                 file_version=self.file.version,
                                                 replica=self.replica,
                                                 token=self.token)
-        dss_response = requests.get(dss_url, allow_redirects=False)
-        if dss_response.status_code == 301:
+        dss_response = self._http_client.request('GET', dss_url, redirect=False)
+        if dss_response.status == 301:
             retry_after = int(dss_response.headers.get('Retry-After'))
             location = dss_response.headers['Location']
 
@@ -233,7 +233,7 @@ def update(self, authentication: Authentication | None) -> None:
             self.replica = one(query['replica'])
             self.file = attrs.evolve(self.file, version=one(query['version']))
             self._retry_after = retry_after
-        elif dss_response.status_code == 302:
+        elif dss_response.status == 302:
             location = dss_response.headers['Location']
             # Remove once https://github.com/HumanCellAtlas/data-store/issues/1837 is resolved
             if True:
@@ -256,7 +256,7 @@ def update(self, authentication: Authentication | None) -> None:
                                                          Params=params)
             self._location = location
         else:
-            dss_response.raise_for_status()
+            raise_on_status(dss_response)
             assert False
 
     @property

src/azul/service/drs_controller.py

@@ -12,6 +12,7 @@
 from datetime import (
     datetime,
 )
+import logging
 from typing import (
     Any,
 )
@@ -27,7 +28,7 @@
 from more_itertools import (
     one,
 )
-import requests
+import urllib3
 
 from azul import (
     config,
@@ -38,6 +39,9 @@
     drs_object_uri,
     drs_object_url_path,
 )
+from azul.http import (
+    HasCachedHttpClient,
+)
 from azul.lib import (
     cached_property,
     mutable_furl,
@@ -60,8 +64,10 @@
     IndexService,
 )
 
+log = logging.getLogger(__name__)
+
 
-class DRSController(ServiceController):
+class DRSController(ServiceController, HasCachedHttpClient):
 
     @cached_property
     def _service(self) -> IndexService:
@@ -205,22 +211,22 @@ def get_object(self, file_uuid, query_params):
             # We only want direct URLs for Google
             extra_params = dict(query_params, directurl=access_method.replica == 'gcp')
             response = self._dss_get_file(file_uuid, access_method.replica, **extra_params)
-            if response.status_code == 301:
+            if response.status == 301:
                 retry_url = response.headers['location']
                 query = urllib.parse.urlparse(retry_url).query
                 query = urllib.parse.parse_qs(query, strict_parsing=True)
                 token = one(query['token'])
                 # We use the encoded token string as the key for our access ID.
                 access_id = encode_access_id(token, access_method.replica)
                 drs_object.add_access_method(access_method, access_id=access_id)
-            elif response.status_code == 302:
+            elif response.status == 302:
                 retry_url = response.headers['location']
                 if access_method.replica == 'gcp':
                     assert retry_url.startswith('gs:')
                 drs_object.add_access_method(access_method, url=retry_url)
             else:
                 # For errors, just proxy DSS response
-                return Response(response.text, status_code=response.status_code)
+                return Response(response.data, status_code=response.status)
         return Response(drs_object.to_json())
 
     def get_object_access(self, access_id, file_uuid, query_params):
@@ -238,24 +244,33 @@ def get_object_access(self, access_id, file_uuid, query_params):
                 'directurl': replica == 'gcp',
                 'token': token
             })
-            if response.status_code == 301:
-                headers = {'retry-after': response.headers['retry-after']}
+            if response.status == 301:
+                headers: dict[str, str | list[str]] = {
+                    'retry-after': response.headers['retry-after']
+                }
                 # DRS says no body for 202 responses
                 return Response(body='', status_code=202, headers=headers)
-            elif response.status_code == 302:
+            elif response.status == 302:
                 retry_url = response.headers['location']
                 return Response(self._access_url(retry_url))
             else:
                 # For errors, just proxy DSS response
-                return Response(response.text, status_code=response.status_code)
+                return Response(response.data, status_code=response.status)
 
-    def _dss_get_file(self, file_uuid, replica, **kwargs):
+    def _dss_get_file(self,
+                      file_uuid,
+                      replica,
+                      **kwargs
+                      ) -> urllib3.BaseHTTPResponse:
         dss_params = {
             'replica': replica,
             **kwargs
         }
         url = self.dss_file_url(file_uuid)
-        return requests.api.get(str(url), params=dss_params, allow_redirects=False)
+        return self._http_client.request('GET',
+                                         str(url),
+                                         fields=dss_params,
+                                         redirect=False)
 
     @classmethod
     def dss_file_url(cls, file_uuid: str) -> mutable_furl:
@@ -267,7 +282,7 @@ class GatewayTimeoutError(ChaliceViewError):
 
 
 @dataclass
-class DRSObject:
+class DRSObject(HasCachedHttpClient):
     """"
     Used to build up a https://ga4gh.github.io/data-repository-service-schemas/docs/#_drsobject
     """
@@ -293,7 +308,7 @@ def add_access_method(self,
     def to_json(self) -> JSON:
         args = _url_query(replica='aws', version=self.version)
         url = DRSController.dss_file_url(self.uuid).add(args=args)
-        headers = requests.api.head(str(url)).headers
+        headers = self._http_client.request('HEAD', str(url)).headers
         version = headers['x-dss-version']
         if self.version is not None:
             assert version == self.version

src/humancellatlas/data/metadata/helpers/schema_validation.py

@@ -15,8 +15,11 @@
     Registry,
     Resource,
 )
-import requests
 
+from azul.http import (
+    HasCachedHttpClient,
+    raise_on_status,
+)
 from azul.lib import (
     R,
     cached_property,
@@ -28,7 +31,7 @@
 log = logging.getLogger(__name__)
 
 
-class SchemaValidator:
+class SchemaValidator(HasCachedHttpClient):
 
     def validate_json(self, file_json: JSON, file_name: str):
         try:
@@ -45,8 +48,8 @@ def validate_json(self, file_json: JSON, file_name: str):
 
     @lru_cache(maxsize=None)
     def _download_json_file(self, file_url: str) -> JSON:
-        response = requests.get(file_url, allow_redirects=False)
-        response.raise_for_status()
+        response = self._http_client.request('GET', file_url, redirect=False)
+        raise_on_status(response)
         return response.json()
 
     def _retrieve_resource(self, resource_url: str) -> Resource:

test/app_test_case.py

@@ -18,14 +18,16 @@
 from furl import (
     furl,
 )
-import requests
-
+import urllib3
 from azul import (
     config,
 )
 from azul.chalice import (
     AzulChaliceApp,
 )
+from azul.http import (
+    raise_on_status,
+)
 from azul.lib import (
     mutable_furl,
 )
@@ -129,7 +131,7 @@ def setUp(self):
         while True:
             try:
                 response = self._ping()
-                response.raise_for_status()
+                raise_on_status(response)
             except Exception:
                 if time.time() > deadline:
                     raise
@@ -138,8 +140,14 @@ def setUp(self):
             else:
                 break
 
-    def _ping(self):
-        return requests.get(str(self.base_url.set(path='/health/basic')))
+    def _ping(self) -> urllib3.BaseHTTPResponse:
+        return self._http_client.urlopen('GET',
+                                         str(self.base_url.set(path='/health/basic')),
+                                         retries=urllib3.Retry(connect=2,
+                                                               read=2,
+                                                               status=0,
+                                                               redirect=0,
+                                                               status_forcelist={500}))
 
     def chalice_config(self):
         return ChaliceConfig.create(lambda_timeout=config.api_gateway_lambda_timeout)