Make Notebook.backup() crash-safe

src/labapi/tree/notebook.py

@@ -10,7 +10,7 @@
 
 from os import PathLike
 from pathlib import Path
-from tempfile import TemporaryDirectory
+from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import TYPE_CHECKING, Literal
 
 from typing_extensions import override
@@ -205,7 +205,26 @@ def backup(
             raise
 
         path = Path(destination)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        with stream, path.open("wb") as file:
-            file.writelines(stream)
+        with stream:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            # Write to a temp file in the destination directory and atomically
+            # replace it, so an interrupted download never truncates an
+            # existing archive or leaves a corrupt file at the destination.
+            tmp: Path | None = None
+            try:
+                with NamedTemporaryFile(
+                    dir=path.parent,
+                    prefix=f"{path.name}.",
+                    suffix=".part",
+                    delete=False,
+                ) as file:
+                    tmp = Path(file.name)
+                    file.writelines(stream)
+                tmp.replace(path)
+            except BaseException:
+                # Remove the partial temp file on any failure (including
+                # KeyboardInterrupt), then re-raise; the destination is untouched.
+                if tmp is not None:
+                    tmp.unlink(missing_ok=True)
+                raise
         return path

tests/tree/test_backup.py

@@ -72,3 +72,46 @@ def test_backup_propagates_other_errors(client, notebook: LA.Notebook, tmp_path)
         notebook.backup(tmp_path / "nb.7z")
 
     assert exc_info.value is raw_error
+
+
+def _failing_stream(client) -> Mock:
+    """Replace stream_api_get with a stream that raises partway through."""
+
+    def chunks():
+        yield b"partial"
+        raise RuntimeError("network drop")
+
+    response = Mock()
+    response.iter_content.return_value = chunks()
+    stream = Mock(return_value=StreamingResponse(response))
+    client.stream_api_get = stream
+    return stream
+
+
+def test_backup_interrupted_preserves_existing_archive(
+    client, notebook: LA.Notebook, tmp_path
+):
+    """A mid-download failure must not truncate the existing archive or leave a partial."""
+    dest = tmp_path / "nb.7z"
+    dest.write_bytes(b"OLD-GOOD-ARCHIVE")
+    _failing_stream(client)
+
+    with pytest.raises(RuntimeError, match="network drop"):
+        notebook.backup(dest)
+
+    assert dest.read_bytes() == b"OLD-GOOD-ARCHIVE"
+    assert list(tmp_path.glob("*.part")) == []
+
+
+def test_backup_closes_stream_when_mkdir_fails(client, notebook: LA.Notebook, tmp_path):
+    """If destination-directory creation fails, the HTTP stream is still closed."""
+    blocker = tmp_path / "blocker"
+    blocker.write_bytes(b"i am a file, not a directory")
+    response = Mock()
+    response.iter_content.return_value = [b"data"]
+    client.stream_api_get = Mock(return_value=StreamingResponse(response))
+
+    with pytest.raises(FileExistsError):
+        notebook.backup(blocker / "nb.7z")
+
+    response.close.assert_called_once_with()