From a89027f0ba05930c38d8ea5a315304f815e88c9a Mon Sep 17 00:00:00 2001
From: Rebecca Breu <rebecca@rbreu.de>
Date: Sun, 19 May 2024 22:38:42 +0200
Subject: [PATCH] Fix hang when saving an open bee file that has been removed

---
 CHANGELOG.rst             |  9 +++++++++
 beeref/fileio/__init__.py |  2 +-
 beeref/fileio/sql.py      | 24 +++++++++++++++++-------
 tests/fileio/test_sql.py  | 11 +++++++++++
 4 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8947cc2..bc2756b 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -12,6 +12,15 @@ Added
 * Added option to export all images from scene (File -> Export Images)
 
 
+Fixed
+-----
+
+* Fixed a case where adding/importing an image would hang while
+  reading unsupported exif data (#111)
+* Fixed a hang when saving an open bee file that has been removed
+  since being opened
+
+
 0.3.3 - 2024-05-05
 ==================
 
diff --git a/beeref/fileio/__init__.py b/beeref/fileio/__init__.py
index 4307ca9..5834d97 100644
--- a/beeref/fileio/__init__.py
+++ b/beeref/fileio/__init__.py
@@ -49,7 +49,7 @@ def save_bee(filename, scene, create_new=False, worker=None):
     logger.debug(f'Create new: {create_new}')
     io = SQLiteIO(filename, scene, create_new, worker=worker)
     io.write()
-    logger.info('Saved!')
+    logger.info('End save')
 
 
 def load_images(filenames, pos, scene, worker):
diff --git a/beeref/fileio/sql.py b/beeref/fileio/sql.py
index f4c7b5f..de98686 100644
--- a/beeref/fileio/sql.py
+++ b/beeref/fileio/sql.py
@@ -52,7 +52,7 @@ def handle_sqlite_errors(func):
     def wrapper(self, *args, **kwargs):
         try:
             func(self, *args, **kwargs)
-        except (sqlite3.Error, BeeFileIOError) as e:
+        except Exception as e:
             logger.exception(f'Error while reading/writing {self.filename}')
             try:
                 # Try to roll back transaction if there is any
@@ -80,6 +80,7 @@ def __init__(self, filename, scene, create_new=False, readonly=False,
         self.filename = filename
         self.readonly = readonly
         self.worker = worker
+        self.retry = False
 
     def __del__(self):
         self._close_connection()
@@ -109,7 +110,13 @@ def _establish_connection(self):
         self._connection = sqlite3.connect(uri, uri=True)
         self._cursor = self.connection.cursor()
         if not self.create_new:
-            self._migrate()
+            try:
+                self._migrate()
+            except Exception:
+                # Updating a file failed; try creating it from scratch instead
+                logger.exception('Error migrating bee file')
+                self.create_new = True
+                self._establish_connection()
 
     def _migrate(self):
         """Migrate database if necessary."""
@@ -236,16 +243,19 @@ def read(self):
     def write(self):
         if self.readonly:
             raise sqlite3.OperationalError(
-                'attempt to write a readonly database')
+                'Attempt to write to a readonly database')
         try:
             self.create_schema_on_new()
             self.write_data()
-        except sqlite3.Error:
-            if self.create_new:
-                # If writing to a new file fails, we can't recover
+        except Exception:
+            if self.retry:
+                # Trying to recover failed
                 raise
             else:
-                # Updating a file failed; try creating it from scratch instead
+                self.retry = True
+                # Try creating file from scratch and save again
+                logger.exception(
+                    f'Updating to existing file {self.filename} failed')
                 self.create_new = True
                 self._close_connection()
                 self.write()
diff --git a/tests/fileio/test_sql.py b/tests/fileio/test_sql.py
index 5f640ae..c459dde 100644
--- a/tests/fileio/test_sql.py
+++ b/tests/fileio/test_sql.py
@@ -472,6 +472,7 @@ def test_sqliteio_write_removes_nonexisting_pixmap_item(tmpfile, view):
 
 def test_sqliteio_write_update_recovers_from_borked_file(view, tmpfile):
     item = BeePixmapItem(QtGui.QImage(), filename='bee.png')
+    item.save_id = 1
     view.scene.addItem(item)
 
     with open(tmpfile, 'w') as f:
@@ -483,6 +484,16 @@ def test_sqliteio_write_update_recovers_from_borked_file(view, tmpfile):
     assert result[0] == 1
 
 
+def test_sqliteio_write_update_recovers_from_nonexisting_file(view, tmpfile):
+    item = BeePixmapItem(QtGui.QImage(), filename='bee.png')
+    item.save_id = 1
+    view.scene.addItem(item)
+    io = SQLiteIO(tmpfile, view.scene, create_new=False)
+    io.write()
+    result = io.fetchone('SELECT COUNT(*) FROM items')
+    assert result[0] == 1
+
+
 def test_sqliteio_write_updates_progress(tmpfile, view):
     worker = MagicMock(canceled=False)
     io = SQLiteIO(tmpfile, view.scene, create_new=True, worker=worker)