fix flaky crime download

pipeline/download/crime.py

@@ -15,6 +15,7 @@ import json
 import re
 import shutil
 import sys
+import time
 import zipfile
 from dataclasses import asdict, dataclass
 from datetime import UTC, datetime
@@ -35,6 +36,11 @@ ARCHIVE_LINK_RE = re.compile(
 )
 VALID_MD5_RE = re.compile(r"^[0-9a-fA-F]{32}$")
 MONTH_RE = re.compile(r"^\d{4}-\d{2}$")
+# Each archive is up to ~1.6 GB, so a transient failure (network blip, or the
+# work dir being removed underneath us by a concurrent run) must not abort the
+# whole multi-archive download. Retry a few times with a short backoff.
+DOWNLOAD_ATTEMPTS = 4
+RETRY_BACKOFF_SECONDS = 5
 STREET_CRIME_CSV_RE = re.compile(r"^\d{4}-\d{2}-.+-street\.csv$")
 CONTAINED_RANGE_RE = re.compile(
     r"Contains data from (?P<start_month>[A-Za-z]+) (?P<start_year>\d{4}) "
@@ -265,6 +271,10 @@ def file_md5(path: Path) -> str:
     return digest.hexdigest()
 
 
+class _RetryableDownloadError(Exception):
+    """A download attempt failed in a way that is worth retrying."""
+
+
 def download_archive(
     archive: CrimeArchive,
     archive_dir: Path,
@@ -273,7 +283,12 @@ def download_archive(
     force: bool,
     timeout: float,
 ) -> Path:
-    """Download one archive ZIP, resuming an existing .part file when possible."""
+    """Download one archive ZIP, resuming an existing .part file when possible.
+
+    Transient failures (network errors, or the work dir being removed
+    underneath us by a concurrent run) are retried with a short backoff so a
+    single hiccup does not abort a whole multi-archive download.
+    """
     dest = archive_dir / archive.filename
     partial = dest.with_suffix(dest.suffix + ".part")
 
@@ -297,6 +312,43 @@ def download_archive(
             print(f"{archive.filename}: already downloaded")
             return dest
 
+    last_error: Exception | None = None
+    for attempt in range(1, DOWNLOAD_ATTEMPTS + 1):
+        # The work dir can be removed underneath us (e.g. a concurrent run's
+        # startup cleanup), which is exactly what makes the final rename fail
+        # with ENOENT even after the download reached 100%. Recreate it so the
+        # (re)download has somewhere to land; a surviving .part is resumed.
+        archive_dir.mkdir(parents=True, exist_ok=True)
+        try:
+            _stream_to_partial(archive, partial, timeout=timeout)
+            partial.replace(dest)
+            if verify and archive.md5 is not None:
+                actual_md5 = file_md5(dest)
+                if actual_md5 != archive.md5:
+                    dest.unlink(missing_ok=True)
+                    raise _RetryableDownloadError(
+                        f"{archive.filename}: MD5 mismatch: "
+                        f"expected {archive.md5}, got {actual_md5}"
+                    )
+            return dest
+        except (httpx.HTTPError, OSError, _RetryableDownloadError) as error:
+            last_error = error
+            if attempt < DOWNLOAD_ATTEMPTS:
+                delay = RETRY_BACKOFF_SECONDS * attempt
+                print(
+                    f"{archive.filename}: download attempt {attempt} failed "
+                    f"({error}); retrying in {delay}s",
+                    file=sys.stderr,
+                )
+                time.sleep(delay)
+
+    raise RuntimeError(
+        f"{archive.filename}: download failed after {DOWNLOAD_ATTEMPTS} attempts"
+    ) from last_error
+
+
+def _stream_to_partial(archive: CrimeArchive, partial: Path, *, timeout: float) -> None:
+    """Stream one archive into its .part file, resuming a valid prefix."""
     resume_from = partial.stat().st_size if partial.exists() else 0
     headers = {"Range": f"bytes={resume_from}-"} if resume_from else {}
 
@@ -332,18 +384,6 @@ def download_archive(
                 output.write(chunk)
                 progress.update(len(chunk))
 
-    partial.replace(dest)
-
-    if verify and archive.md5 is not None:
-        actual_md5 = file_md5(dest)
-        if actual_md5 != archive.md5:
-            dest.unlink(missing_ok=True)
-            raise RuntimeError(
-                f"{archive.filename}: MD5 mismatch: expected {archive.md5}, got {actual_md5}"
-            )
-
-    return dest
-
 
 def _is_street_crime_csv(path: PurePosixPath | Path) -> bool:
     return STREET_CRIME_CSV_RE.fullmatch(path.name) is not None

pipeline/download/test_crime.py

@@ -1,7 +1,13 @@
+import hashlib
+import shutil
 from zipfile import ZipFile
 
+import pytest
+
+from pipeline.download import crime
 from pipeline.download.crime import (
     CrimeArchive,
+    download_archive,
     extract_csvs,
     prepare_archive_dir,
     prune_unused_csvs,
@@ -10,6 +16,19 @@ from pipeline.download.crime import (
 )
 
 
+def _make_archive(md5: str | None = None) -> CrimeArchive:
+    return CrimeArchive(
+        month="2023-05",
+        label="May 2023",
+        url="https://data.police.uk/data/archive/2023-05.zip",
+        filename="2023-05.zip",
+        size="1.6 GB",
+        contained_range="",
+        md5=md5,
+        raw_md5=md5 or "",
+    )
+
+
 def test_parse_archives_reads_monthly_zip_links_only():
     html = """
     <p><a href="/data/archive/latest.zip">latest.zip</a></p>
@@ -194,3 +213,47 @@ def test_prune_unused_csvs_removes_non_street_csvs(tmp_path):
     assert street.exists()
     assert not outcomes.exists()
     assert not stop_search.exists()
+
+
+def test_download_archive_recovers_when_workdir_wiped_mid_download(
+    tmp_path, monkeypatch
+):
+    # Reproduces the crash where a completed .part vanished before the rename
+    # (a concurrent run wiped _download_tmp), so partial.replace raised ENOENT
+    # even though the download had reached 100%.
+    monkeypatch.setattr(crime.time, "sleep", lambda *_: None)
+    archive_dir = tmp_path / "_download_tmp"
+    archive_dir.mkdir()
+    payload = b"real-zip-bytes" * 4096
+    archive = _make_archive(md5=hashlib.md5(payload).hexdigest())
+
+    calls = {"n": 0}
+
+    def fake_stream(_archive, partial, *, timeout):
+        calls["n"] += 1
+        partial.write_bytes(payload)
+        if calls["n"] == 1:
+            # Simulate the concurrent cleanup: the rename target disappears.
+            shutil.rmtree(partial.parent)
+
+    monkeypatch.setattr(crime, "_stream_to_partial", fake_stream)
+
+    dest = download_archive(archive, archive_dir, verify=True, force=False, timeout=1.0)
+
+    assert calls["n"] == 2
+    assert dest.read_bytes() == payload
+
+
+def test_download_archive_gives_up_after_repeated_failures(tmp_path, monkeypatch):
+    monkeypatch.setattr(crime.time, "sleep", lambda *_: None)
+    archive_dir = tmp_path / "_download_tmp"
+
+    def always_fail(_archive, _partial, *, timeout):
+        raise OSError("connection reset")
+
+    monkeypatch.setattr(crime, "_stream_to_partial", always_fail)
+
+    with pytest.raises(RuntimeError, match=f"after {crime.DOWNLOAD_ATTEMPTS} attempts"):
+        download_archive(
+            _make_archive(), archive_dir, verify=False, force=False, timeout=1.0
+        )