Replace requests with aiohttp

requirements.txt

@@ -2,19 +2,13 @@ aiohttp==3.6.2
 aiohttp-xmlrpc==0.8.1
 async-timeout==3.0.1
 attrs==19.3.0
-certifi==2019.11.28
 chardet==3.0.4
 filelock==3.0.12
 idna==2.9
-idna-ssl==1.1.0
 lxml==4.5.0
 multidict==4.7.5
 packaging==20.1
 pyparsing==2.4.6
-python-dateutil==2.8.1
-requests==2.23.0
 setuptools==45.2.0
 six==1.14.0
-typing-extensions==3.7.4.1
-urllib3==1.25.8
 yarl==1.4.2

requirements_test.txt

@@ -7,7 +7,7 @@ flake8-bugbear==20.1.4
 freezegun==0.3.15
 pre-commit==2.1.0
 pytest==5.3.5
+pytest-asyncio==0.10.0
 pytest-timeout==1.3.4
-pytest-cache==1.0
 setuptools==45.2.0
 tox==3.14.5

setup.cfg

@@ -16,7 +16,7 @@ project_urls =
     Source Code = https://github.com/pypa/bandersnatch
     Change Log = https://github.com/pypa/bandersnatch/CHANGES.md
 url = https://github.com/pypa/bandersnatch/
-version = 4.0.0.dev0
+version = 4.0.0.dev1
 
 [options]
 include_package_data = True
@@ -25,7 +25,6 @@ install_requires =
     aiohttp-xmlrpc
     filelock
     packaging
-    requests
     setuptools>40.0.0
 package_dir =
     =src

src/bandersnatch/__init__.py

@@ -19,7 +19,7 @@ def version_str(self) -> str:
     major=4,
     minor=0,
     micro=0,
-    releaselevel="dev0",
+    releaselevel="dev1",
     serial=0,  # Not currently in use with Bandersnatch versioning
 )
 __version__ = __version_info__.version_str

src/bandersnatch/main.py

@@ -4,8 +4,10 @@
 import logging.config
 import shutil
 import sys
+from configparser import ConfigParser
 from pathlib import Path
 from tempfile import gettempdir
+from typing import Optional
 
 import bandersnatch.configuration
 import bandersnatch.delete
@@ -89,7 +91,36 @@ def _verify_parser(subparsers: argparse._SubParsersAction) -> None:
     v.set_defaults(op="verify")
 
 
-def main() -> int:
+async def async_main(args: argparse.Namespace, config: ConfigParser) -> int:
+    if args.op.lower() == "delete":
+        return await bandersnatch.delete.delete_packages(config, args)
+    elif args.op.lower() == "verify":
+        return await bandersnatch.verify.metadata_verify(config, args)
+
+    if args.force_check:
+        status_file = Path(config.get("mirror", "directory")) / "status"
+        if status_file.exists():
+            tmp_status_file = Path(gettempdir()) / "status"
+            try:
+                shutil.move(status_file, tmp_status_file)
+                logger.debug(
+                    "Force bandersnatch to check everything against the master PyPI"
+                    + f" - status file moved to {tmp_status_file}"
+                )
+            except OSError as e:
+                logger.error(
+                    f"Could not move status file ({status_file} to "
+                    + f" {tmp_status_file}): {e}"
+                )
+        else:
+            logger.info(
+                f"No status file to move ({status_file}) - Full sync will occur"
+            )
+
+    return await bandersnatch.mirror.mirror(config)
+
+
+def main(loop: Optional[asyncio.AbstractEventLoop] = None) -> int:
     parser = argparse.ArgumentParser(description="PyPI PEP 381 mirroring client.")
     parser.add_argument(
         "--version", action="version", version=f"%(prog)s {bandersnatch.__version__}"
@@ -140,39 +171,12 @@ def main() -> int:
     if config.has_option("mirror", "log-config"):
         logging.config.fileConfig(str(Path(config.get("mirror", "log-config"))))
 
-    if args.op in ("delete", "verify"):
-        loop = asyncio.get_event_loop()
-        op_coro = (
-            bandersnatch.delete.delete_packages
-            if args.op == "delete"
-            else bandersnatch.verify.metadata_verify
-        )
-        try:
-            return loop.run_until_complete(op_coro(config, args))
-        finally:
-            loop.close()
-    else:
-        if args.force_check:
-            status_file = Path(config.get("mirror", "directory")) / "status"
-            if status_file.exists():
-                tmp_status_file = Path(gettempdir()) / "status"
-                try:
-                    shutil.move(status_file, tmp_status_file)
-                    logger.debug(
-                        "Force bandersnatch to check everything against the master PyPI"
-                        + f" - status file moved to {tmp_status_file}"
-                    )
-                except OSError as e:
-                    logger.error(
-                        f"Could not move status file ({status_file} to "
-                        + f" {tmp_status_file}): {e}"
-                    )
-            else:
-                logger.info(
-                    f"No status file to move ({status_file}) - Full sync will occur"
-                )
-
-        return bandersnatch.mirror.mirror(config)
+    # TODO: Go to asyncio.run() when >= 3.7
+    loop = loop or asyncio.get_event_loop()
+    try:
+        return loop.run_until_complete(async_main(args, config))
+    finally:
+        loop.close()
 
 
 if __name__ == "__main__":

src/bandersnatch/master.py

@@ -1,41 +1,55 @@
 import asyncio
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, AsyncGenerator, Dict, Optional
 
-import requests
-from aiohttp import ClientTimeout
+import aiohttp
 from aiohttp_xmlrpc.client import ServerProxy
 
 import bandersnatch
 
 from .utils import USER_AGENT
 
 logger = logging.getLogger(__name__)
+PYPI_SERIAL_HEADER = "X-PYPI-LAST-SERIAL"
 
 
 class StalePage(Exception):
     """We got a page back from PyPI that doesn't meet our expected serial."""
 
 
+class XmlRpcError(aiohttp.ClientError):
+    """Issue getting package listing from PyPI Repository"""
+
+
 class Master:
-    def __init__(self, url, timeout=10.0) -> None:
+    def __init__(self, url: str, timeout: float = 10.0) -> None:
+        self.loop = asyncio.get_event_loop()
+        self.timeout = timeout
         self.url = url
         if self.url.startswith("http://"):
             err = f"Master URL {url} is not https scheme"
             logger.error(err)
             raise ValueError(err)
 
-        self.loop = asyncio.get_event_loop()
-        self.timeout = timeout
-        self.session = requests.Session()
-        self.session.headers.update({"User-Agent": USER_AGENT})
+    async def __aenter__(self) -> "Master":
+        logger.debug("Initializing Master's aiohttp ClientSession")
+        custom_headers = {"User-Agent": USER_AGENT}
+        skip_headers = {"User-Agent"}
+        self.session = aiohttp.ClientSession(
+            headers=custom_headers,
+            skip_auto_headers=skip_headers,
+            timeout=self.timeout,
+            trust_env=True,
+        )
+        return self
 
-    def get(self, path, required_serial, **kw):
-        logger.debug(f"Getting {path} (serial {required_serial})")
-        if not path.startswith(("https://", "http://")):
-            path = self.url + path
-        r = self.session.get(path, timeout=self.timeout, **kw)
-        r.raise_for_status()
+    async def __aexit__(self, *exc) -> None:
+        logger.debug("Closing Master's aiohttp ClientSession")
+        await self.session.close()
+
+    async def check_for_stale_cache(
+        self, path: str, required_serial: Optional[int], got_serial: Optional[int]
+    ) -> None:
         # The PYPI-LAST-SERIAL header allows us to identify cached entries,
         # e.g. via the public CDN or private, transparent mirrors and avoid us
         # injecting stale entries into the mirror without noticing.
@@ -44,33 +58,52 @@ def get(self, path, required_serial, **kw):
             # want you to think really hard before passing in None. This is a
             # really important check to achieve consistency and you should only
             # leave it out if you know what you're doing.
-            got_serial = int(r.headers["X-PYPI-LAST-SERIAL"])
-            if got_serial < required_serial:
+            if not got_serial or got_serial < required_serial:
                 logger.debug(
-                    "Expected PyPI serial {} for request {} but got {}".format(
-                        required_serial, path, got_serial
-                    )
+                    f"Expected PyPI serial {required_serial} for request {path} "
+                    + f"but got {got_serial}"
                 )
+
                 # HACK: The following attempts to purge the cache of the page we
                 # just tried to fetch. This works around PyPI's caches sometimes
                 # returning a stale serial for a package. Ideally, this should
                 # be fixed on the PyPI side, at which point the following code
                 # should be removed.
                 logger.debug(f"Issuing a PURGE for {path} to clear the cache")
                 try:
-                    self.session.request("PURGE", path, timeout=self.timeout)
-                except requests.exceptions.HTTPError:
+                    async with self.session.request("PURGE", path):
+                        pass
+                except aiohttp.ClientError:
                     logger.warning(
                         "Got an error when attempting to clear the cache", exc_info=True
                     )
 
                 raise StalePage(
-                    "Expected PyPI serial {} for request {} but got {}. "
-                    + "HTTP PURGE has been issued to the request url".format(
-                        required_serial, path, got_serial
-                    )
+                    f"Expected PyPI serial {required_serial} for request {path} "
+                    + f"but got {got_serial}. "
+                    + "HTTP PURGE has been issued to the request url"
                 )
-        return r
+
+    async def get(
+        self, path: str, required_serial: Optional[int], **kw: Any
+    ) -> AsyncGenerator[aiohttp.ClientResponse, None]:
+        logger.debug(f"Getting {path} (serial {required_serial})")
+        if not path.startswith(("https://", "http://")):
+            path = self.url + path
+
+        timeout = aiohttp.ClientTimeout(total=300)
+        if "timeout" in kw:
+            timeout = aiohttp.ClientTimeout(total=kw["timeout"])
+            del kw["timeout"]
+
+        async with self.session.get(path, timeout=timeout, **kw) as r:
+            got_serial = (
+                int(r.headers[PYPI_SERIAL_HEADER])
+                if PYPI_SERIAL_HEADER in r.headers
+                else None
+            )
+            await self.check_for_stale_cache(path, required_serial, got_serial)
+            yield r
 
     @property
     def xmlrpc_url(self) -> str:
@@ -91,13 +124,13 @@ async def _gen_custom_headers(self) -> Dict[str, str]:
 
     async def _gen_xmlrpc_client(self) -> ServerProxy:
         custom_headers = await self._gen_custom_headers()
-        timeout = ClientTimeout(total=self.timeout)
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
         client = ServerProxy(
             self.xmlrpc_url, loop=self.loop, headers=custom_headers, timeout=timeout
         )
         return client
 
-    # TODO: Add an async decorator to aiohttp-xmlrpc to replace this function
+    # TODO: Add an async context manager to aiohttp-xmlrpc to replace this function
     async def rpc(self, method_name: str, serial: int = 0) -> Any:
         try:
             client = await self._gen_xmlrpc_client()
@@ -112,9 +145,12 @@ async def rpc(self, method_name: str, serial: int = 0) -> Any:
             await client.client.close()
 
     async def all_packages(self) -> Optional[Dict[str, int]]:
-        return await self.rpc("list_packages_with_serial")
+        all_packages_with_serial = await self.rpc("list_packages_with_serial")
+        if not all_packages_with_serial:
+            raise XmlRpcError("Unable to get full list of packages")
+        return all_packages_with_serial
 
-    async def changed_packages(self, last_serial: int) -> Optional[Dict[str, int]]:
+    async def changed_packages(self, last_serial: int) -> Dict[str, int]:
         changelog = await self.rpc("changelog_since_serial", last_serial)
         if changelog is None:
             changelog = []