DRIVERS-3032 Make MongoDB downloads more robust (#613)

Co-authored-by: Ezra Chung <[email protected]>

Author: blink1073

.evergreen/mongodl.py

@@ -609,30 +609,35 @@ def download_file(self, url: str) -> DownloadResult:
         if modtime:
             headers["If-Modified-Since"] = modtime
         digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:4]
-        dest = self._dirpath / "files" / digest / PurePosixPath(url).name
+        file_name = PurePosixPath(url).name
+        dest = self._dirpath / "files" / digest / file_name
         if not dest.exists():
             headers = {}
         req = urllib.request.Request(url, headers=headers)
 
         try:
-            resp = urllib.request.urlopen(req, context=SSL_CONTEXT)
+            resp = urllib.request.urlopen(req, context=SSL_CONTEXT, timeout=30)
         except urllib.error.HTTPError as e:
             if e.code != 304:
                 raise RuntimeError(f"Failed to download [{url}]") from e
             assert dest.is_file(), (
                 "The download cache is missing an expected file",
                 dest,
             )
+            LOGGER.info("Using cached file %s", file_name)
             return DownloadResult(False, dest)
 
         _mkdir(dest.parent)
         got_etag = resp.getheader("ETag")
         got_modtime = resp.getheader("Last-Modified")
+        got_len = int(resp.getheader("Content-Length"))
         with dest.open("wb") as of:
-            buf = resp.read(1024 * 1024 * 4)
-            while buf:
-                of.write(buf)
-                buf = resp.read(1024 * 1024 * 4)
+            shutil.copyfileobj(resp, of, length=got_len)
+        file_size = dest.stat().st_size
+        if file_size != got_len:
+            raise RuntimeError(
+                f"File size: {file_size} does not match download size: {got_len}"
+            )
         self._db(
             "INSERT OR REPLACE INTO mdl_http_downloads (url, etag, last_modified) "
             "VALUES (:url, :etag, :mtime)",
@@ -895,7 +900,8 @@ def _dl_component(
             return _expand_archive(
                 cached, out_dir, pattern, strip_components, test=test
             )
-        except Exception:
+        except Exception as e:
+            LOGGER.exception(e)
             if not retrier.retry():
                 raise
 

.evergreen/mongosh_dl.py

@@ -25,9 +25,11 @@
 from mongodl import LOGGER as DL_LOGGER
 from mongodl import (
     SSL_CONTEXT,
+    Cache,
     DownloadRetrier,
     ExpandResult,
     _expand_archive,
+    default_cache_dir,
     infer_arch,
 )
 
@@ -40,7 +42,7 @@ def _get_latest_version():
     url = "https://api.github.com/repos/mongodb-js/mongosh/releases"
     req = urllib.request.Request(url, headers=headers)
     try:
-        resp = urllib.request.urlopen(req, context=SSL_CONTEXT)
+        resp = urllib.request.urlopen(req, context=SSL_CONTEXT, timeout=30)
     except Exception:
         return _get_latest_version_git()
 
@@ -73,6 +75,7 @@ def _get_latest_version_git():
 
 
 def _download(
+    cache: Cache,
     out_dir: Path,
     version: str,
     target: str,
@@ -108,24 +111,15 @@ def _download(
     if no_download:
         return ExpandResult.Okay
 
-    req = urllib.request.Request(dl_url)
     retrier = DownloadRetrier(retries)
     while True:
         try:
-            resp = urllib.request.urlopen(req)
-            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as fp:
-                four_mebibytes = 1024 * 1024 * 4
-                buf = resp.read(four_mebibytes)
-                while buf:
-                    fp.write(buf)
-                    buf = resp.read(four_mebibytes)
-                fp.close()
-                resp = _expand_archive(
-                    Path(fp.name), out_dir, pattern, strip_components, test=test
-                )
-                os.remove(fp.name)
-            return resp
-        except Exception:
+            cached = cache.download_file(dl_url).path
+            return _expand_archive(
+                cached, out_dir, pattern, strip_components, test=test
+            )
+        except Exception as e:
+            LOGGER.exception(e)
             if not retrier.retry():
                 raise
 
@@ -140,6 +134,12 @@ def main(argv=None):
     parser.add_argument(
         "--quiet", "-q", action="store_true", help="Whether to log at the WARNING level"
     )
+    parser.add_argument(
+        "--cache-dir",
+        type=Path,
+        default=default_cache_dir(),
+        help="Directory where download caches and metadata will be stored",
+    )
     dl_grp = parser.add_argument_group(
         "Download arguments",
         description="Select what to download and extract. "
@@ -220,7 +220,10 @@ def main(argv=None):
     elif args.quiet:
         LOGGER.setLevel(logging.WARNING)
         DL_LOGGER.setLevel(logging.WARNING)
+
+    cache = Cache.open_in(args.cache_dir)
     result = _download(
+        cache,
         out,
         version=args.version,
         target=target,