files are downloaded atomically using path.rename()

Author: jewettaijfc

CHANGELOG.md

@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 or `Absorber` classes (or when invoking `pml()`, `stable_pml()`, or `absorber()` functions)
 with fewer layers than recommended.
 - Warnings and error messages originating from `Structure`, `Source`, or `Monitor` classes now refer to problematic objects by their user-supplied `name` attribute, alongside their index.
+- File downloads are atomic. Interruptions or failures during download will no longer result in incomplete files.
 
 ### Fixed
 - Arrow lengths are now scaled consistently in the X and Y directions, and their lengths no longer exceed the height of the plot window.

tests/test_web/test_s3utils.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+import tidy3d
+from tidy3d.web.core import s3utils
+
+
+@pytest.fixture
+def mock_S3STSToken(monkeypatch):
+    mock_token = MagicMock()
+    mock_token.cloud_path = ""
+    mock_token.user_credential = ""
+    mock_token.get_bucket = lambda: ""
+    mock_token.get_s3_key = lambda: ""
+    mock_token.is_expired = lambda: False
+    mock_token.get_client = lambda: tidy3d.web.core.s3utils.boto3.client()
+    monkeypatch.setattr(
+        target=tidy3d.web.core.s3utils, name="_S3STSToken", value=MagicMock(return_value=mock_token)
+    )
+    return mock_token
+
+
+@pytest.fixture
+def mock_get_s3_sts_token(monkeypatch):
+    def _mock_get_s3_sts_token(resource_id, remote_filename):
+        return s3utils._S3STSToken(resource_id, remote_filename)
+
+    monkeypatch.setattr(
+        target=tidy3d.web.core.s3utils, name="get_s3_sts_token", value=_mock_get_s3_sts_token
+    )
+    return _mock_get_s3_sts_token
+
+
+@pytest.fixture
+def mock_s3_client(monkeypatch):
+    """
+    Fixture that provides a generic mock S3 client.
+    Method-specific side_effects are omitted here and are specified later in the unit tests.
+    """
+    mock_client = MagicMock()
+    # Patch the `client` as it is imported within `tidy3d.web.core.s3utils.boto3` so that
+    # whenever it's invoked (for example with "s3"), it returns our `mock_client`.
+    monkeypatch.setattr(
+        target=tidy3d.web.core.s3utils.boto3,
+        name="client",
+        value=MagicMock(return_value=mock_client),
+    )
+    return mock_client
+
+
+def test_download_s3_file_success(mock_s3_client, mock_get_s3_sts_token, mock_S3STSToken, tmp_path):
+    """Tests a successful download."""
+    destination_path = tmp_path / "downloaded_file.txt"
+    expected_content = "abcdefg"
+
+    def simulate_download_success(Bucket, Key, Filename, Callback, Config, **kwargs):
+        with open(Filename, "w") as f:
+            f.write(expected_content)
+        return None
+
+    mock_s3_client.download_file.side_effect = simulate_download_success
+    mock_S3STSToken.get_bucket = lambda: "test-bucket"
+    mock_S3STSToken.get_s3_key = lambda: "test-key"
+
+    s3utils.download_file(
+        resource_id="1234567890",
+        remote_filename=destination_path.name,
+        to_file=str(destination_path),
+        verbose=False,
+        progress_callback=None,
+    )
+
+    # Check that mock_s3_client.download_file() was invoked with the correct arguments.
+    mock_s3_client.download_file.assert_called_once()
+    call_args, call_kwargs = mock_s3_client.download_file.call_args
+    assert call_kwargs["Bucket"] == "test-bucket"
+    assert call_kwargs["Key"] == "test-key"
+    assert call_kwargs["Filename"].endswith(s3utils.IN_TRANSIT_SUFFIX)
+    assert destination_path.exists()
+    with open(destination_path) as f:
+        assert f.read() == expected_content
+    for p in destination_path.parent.iterdir():
+        assert not p.name.endswith(s3utils.IN_TRANSIT_SUFFIX)  # no temporary files are present
+
+
+def test_download_s3_file_raises_oserror(
+    mock_s3_client, mock_get_s3_sts_token, mock_S3STSToken, tmp_path
+):
+    """Tests download failing with an ``OSError`` (No space left on device)."""
+    destination_path = tmp_path / "downloaded_file.txt"
+
+    def simulate_download_failure(Bucket, Key, Filename, Callback, Config, **kwargs):
+        with open(Filename, "w") as f:
+            f.write("abc")
+        raise OSError("No space left on device")
+
+    mock_s3_client.download_file.side_effect = simulate_download_failure
+    mock_S3STSToken.get_bucket = lambda: "test-bucket"
+    mock_S3STSToken.get_s3_key = lambda: "test-key"
+
+    with pytest.raises(OSError, match="No space left on device"):
+        s3utils.download_file(
+            resource_id="1234567890",
+            remote_filename=destination_path.name,
+            to_file=str(destination_path),
+            verbose=False,
+            progress_callback=None,
+        )
+
+    # Check that mock_s3_client.download_file() was invoked with the correct arguments.
+    mock_s3_client.download_file.assert_called_once()
+    call_args, call_kwargs = mock_s3_client.download_file.call_args
+    assert call_kwargs["Bucket"] == "test-bucket"
+    assert call_kwargs["Key"] == "test-key"
+    assert call_kwargs["Filename"].endswith(s3utils.IN_TRANSIT_SUFFIX)
+    # Since downloading failed, no new files should exist locally.
+    assert not destination_path.exists()
+    for p in destination_path.parent.iterdir():
+        assert not p.name.endswith(s3utils.IN_TRANSIT_SUFFIX)  # no temporary files are present

tidy3d/web/core/s3utils.py

@@ -29,6 +29,8 @@
 from .file_util import extract_gzip_file
 from .http_util import http
 
+IN_TRANSIT_SUFFIX = ".tmp"
+
 
 class _UserCredential(BaseModel):
     """Stores information about user credentials."""
@@ -312,12 +314,12 @@ def download_file(
     # set to_file if None
     if not to_file:
         path = pathlib.Path(resource_id)
-        to_file = path / remote_basename
+        to_path = path / remote_basename
     else:
-        to_file = pathlib.Path(to_file)
+        to_path = pathlib.Path(to_file)
 
-    # make the leading directories in the 'to_file', if any
-    to_file.parent.mkdir(parents=True, exist_ok=True)
+    # make the leading directories in the 'to_path', if any
+    to_path.parent.mkdir(parents=True, exist_ok=True)
 
     def _download(_callback: Callable) -> None:
         """Perform the download with a callback function.
@@ -327,14 +329,25 @@ def _download(_callback: Callable) -> None:
         _callback : Callable[[float], None]
             Callback function for download, accepts ``bytes_in_chunk``
         """
-
-        client.download_file(
-            Bucket=token.get_bucket(),
-            Filename=str(to_file),
-            Key=token.get_s3_key(),
-            Callback=_callback,
-            Config=_s3_config,
-        )
+        # Caller can assume the existence of the file means download succeeded.
+        # So make sure this file does not exist until that assumption is true.
+        to_path.unlink(missing_ok=True)
+        # Download to a temporary file.
+        try:
+            fd, tmp_file_path_str = tempfile.mkstemp(suffix=IN_TRANSIT_SUFFIX, dir=to_path.parent)
+            os.close(fd)  # `tempfile.mkstemp()` creates and opens a randomly named file.  close it.
+            to_path_tmp = pathlib.Path(tmp_file_path_str)
+            client.download_file(
+                Bucket=token.get_bucket(),
+                Filename=tmp_file_path_str,
+                Key=token.get_s3_key(),
+                Callback=_callback,
+                Config=_s3_config,
+            )
+            to_path_tmp.rename(to_file)
+        except Exception as e:
+            to_path_tmp.unlink(missing_ok=True)  # Delete incompletely downloaded file.
+            raise e
 
     if progress_callback is not None:
         _download(progress_callback)
@@ -355,7 +368,7 @@ def _callback(bytes_in_chunk):
         else:
             _download(lambda bytes_in_chunk: None)
 
-    return to_file
+    return to_path
 
 
 def download_gz_file(
@@ -394,24 +407,24 @@ def download_gz_file(
 
     # Otherwise, download and unzip
     # The tempfile is set as ``hdf5.gz`` so that the mock download in the webapi tests works
-    tmp_file, tmp_file_path = tempfile.mkstemp(".hdf5.gz")
+    tmp_file, tmp_file_path_str = tempfile.mkstemp(".hdf5.gz")
     os.close(tmp_file)
 
     # make the leading directories in the 'to_file', if any
-    to_file = pathlib.Path(to_file)
-    to_file.parent.mkdir(parents=True, exist_ok=True)
+    to_path = pathlib.Path(to_file)
+    to_path.parent.mkdir(parents=True, exist_ok=True)
     try:
         download_file(
             resource_id,
             remote_filename,
-            to_file=tmp_file_path,
+            to_file=tmp_file_path_str,
             verbose=verbose,
             progress_callback=progress_callback,
         )
-        if os.path.exists(tmp_file_path):
-            extract_gzip_file(tmp_file_path, to_file)
+        if os.path.exists(tmp_file_path_str):
+            extract_gzip_file(tmp_file_path_str, to_path)
         else:
             raise WebError(f"Failed to download and extract '{remote_filename}'.")
     finally:
-        os.unlink(tmp_file_path)
-    return to_file
+        os.unlink(tmp_file_path_str)
+    return to_path