Merge pull request #330 from laura-king/file-cleanup

Remove temporary file in JSON backend store function

docs/source/upcoming_release_notes/304-remove_temporary_file.rst

@@ -0,0 +1,23 @@
+304 remove_temporary_file
+#################
+
+API Changes
+-----------
+- N/A
+
+Features
+--------
+- N/A
+
+Bugfixes
+--------
+- N/A
+
+Maintenance
+-----------
+- Adds error handling for the temporary file created when initializing a json backend object.
+- Changes format of temporary file name generation to contain only a unique hash.
+
+Contributors
+------------
+- laura-king

happi/backends/json_db.py

@@ -2,14 +2,12 @@
 Backend implemenation using the ``simplejson`` package.
 """
 import contextlib
-import getpass
 import logging
 import math
 import os
 import os.path
 import re
 import shutil
-import time
 import uuid
 from typing import Any, Callable, Optional, Union
 
@@ -138,27 +136,29 @@ def store(self, db: dict[str, ItemMeta]) -> None:
             Dictionary to store in JSON.
         """
         temp_path = self._temp_path()
-
-        with open(temp_path, 'w') as fd:
-            json.dump(db, fd, sort_keys=True, indent=4)
-
-        if os.path.exists(self.path):
-            shutil.copymode(self.path, temp_path)
-        shutil.move(temp_path, self.path)
+        try:
+            with open(temp_path, 'w') as fd:
+                json.dump(db, fd, sort_keys=True, indent=4)
+
+            if os.path.exists(self.path):
+                shutil.copymode(self.path, temp_path)
+            shutil.move(temp_path, self.path)
+        except BaseException as ex:
+            logger.debug('JSON db move failed: %s', ex, exc_info=ex)
+            # remove temporary file
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+            raise
 
     def _temp_path(self) -> str:
         """
         Return a temporary path to write the json file to during "store".
 
-        Includes a username and a timestamp for sorting
-        (in the cases where the temp file isn't cleaned up)
-        and a hash for uniqueness
+        Includes a hash for uniqueness
         (in the cases where multiple temp files are written at once).
         """
         directory = os.path.dirname(self.path)
         filename = (
-            f".{getpass.getuser()}"
-            f"_{int(time.time())}"
             f"_{str(uuid.uuid4())[:8]}"
             f"_{os.path.basename(self.path)}"
         )

happi/tests/test_backends.py

@@ -203,6 +203,25 @@ def test_json_tempfile_uniqueness():
     assert len(set(tempfiles)) == len(tempfiles)
 
 
+def test_json_tempfile_remove(monkeypatch):
+    # Set consistent temppath
+    jb = JSONBackend("testing.json", initialize=False)
+    temp_path = jb._temp_path()
+    jb._temp_path = lambda: temp_path
+
+    # Ensure file is created, then throw error through patching
+    def shutil_move_patch(*args, **kwargs):
+        assert os.path.isfile(os.path.join(os.getcwd(), temp_path))
+        raise RuntimeError('Simulated testing error.')
+
+    monkeypatch.setattr('shutil.move', shutil_move_patch)
+
+    # Test, and ensure file is deleted appropriately
+    with pytest.raises(RuntimeError):
+        jb.initialize()
+    assert os.path.exists(temp_path) is False
+
+
 @requires_questionnaire
 def test_qs_find(mockqsbackend):
     assert len(list(mockqsbackend.find(dict(beamline='TST')))) == 14