From 10deb8ea2f1a97bdce4183fd0077c99400e7d496 Mon Sep 17 00:00:00 2001
From: Stefano Scafiti <stefano.scafiti96@gmail.com>
Date: Mon, 23 May 2022 20:44:35 +0200
Subject: [PATCH] feat: use atomic writes to avoid data corruption in case of
 crashes (#468)

(cherry picked from commit 2a2abe52e098136f7183eaaca21efbf42bdef76c)
---
 tinydb/storages.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/tinydb/storages.py b/tinydb/storages.py
index 804ff70f..cb35df76 100644
--- a/tinydb/storages.py
+++ b/tinydb/storages.py
@@ -6,6 +6,7 @@
 import io
 import json
 import os
+import tempfile
 from abc import ABC, abstractmethod
 from typing import Dict, Any, Optional
 
@@ -125,25 +126,31 @@ def read(self) -> Optional[Dict[str, Dict[str, Any]]]:
             return json.load(self._handle)
 
     def write(self, data: Dict[str, Dict[str, Any]]):
-        # Move the cursor to the beginning of the file just in case
-        self._handle.seek(0)
+        file_name = self._handle.name
+
+        # Create a temporary file in the same folder
+        temp_file = tempfile.NamedTemporaryFile(mode=self._mode, prefix=file_name, delete=False)
 
         # Serialize the database state using the user-provided arguments
         serialized = json.dumps(data, **self.kwargs)
 
         # Write the serialized data to the file
         try:
-            self._handle.write(serialized)
+            temp_file.write(serialized)
         except io.UnsupportedOperation:
             raise IOError('Cannot write to the database. Access mode is "{0}"'.format(self._mode))
 
         # Ensure the file has been written
-        self._handle.flush()
-        os.fsync(self._handle.fileno())
+        temp_file.flush()
+        os.fsync(temp_file.fileno())
+
+        # Replace the current file with the temporary file
+        temp_file.close()
+        os.rename(temp_file.name, file_name)
 
-        # Remove data that is behind the new cursor in case the file has
-        # gotten shorter
-        self._handle.truncate()
+        # Reopen the file
+        self._handle.close()
+        self._handle = open(file_name, mode=self._mode, encoding=self._handle.encoding)
 
 
 class MemoryStorage(Storage):