Automatic DB restores now done if upgrade fails due to corrupt DB or incorrect db numbering

sickbeard/databases/mainDB.py

@@ -150,7 +150,6 @@ def backupDatabase(version):
     else:
         logger.log(u"Proceeding with upgrade")
 
-
 # ======================
 # = Main DB Migrations =
 # ======================

sickbeard/db.py

@@ -29,8 +29,6 @@ import sickbeard
 from sickbeard import encodingKludge as ek
 from sickbeard import logger
 from sickbeard.exceptions import ex
-from sickbeard.common import cpu_presets
-from itertools import ifilter
 
 db_lock = threading.Lock()
 
@@ -243,6 +241,8 @@ class DBConnection:
     def hasTable(self, tableName):
         return len(self.action("SELECT 1 FROM sqlite_master WHERE name = ?;", (tableName, )).fetchall()) > 0
 
+    def close(self):
+        self.connection.close()
 
 def sanityCheckDatabase(connection, sanity_check):
     sanity_check(connection).check()
@@ -268,6 +268,13 @@ def upgradeDatabase(connection, schema):
 def prettyName(class_name):
     return ' '.join([x.group() for x in re.finditer("([A-Z])([a-z0-9]+)", class_name)])
 
+def restoreDatabase(version):
+    logger.log(u"Restoring database before trying upgrade again")
+    if not sickbeard.helpers.restoreVersionedFile(dbFilename(suffix='v'+ str(version)), version):
+        logger.log_error_and_exit(u"Database restore failed, abort upgrading database")
+        return False
+    else:
+        return True
 
 def _processUpgrade(connection, upgradeClass):
     instance = upgradeClass(connection)
@@ -277,8 +284,23 @@ def _processUpgrade(connection, upgradeClass):
         try:
             instance.execute()
         except sqlite3.DatabaseError, e:
-            print "Error in " + str(upgradeClass.__name__) + ": " + ex(e)
-            raise
+            # attemping to restore previous DB backup and perform upgrade
+            try:
+                instance.execute()
+            except:
+                restored = False
+                result = connection.select("SELECT db_version FROM db_version")
+                if result:
+                    version = int(result[0]["db_version"])
+                    connection.close()
+                    if restoreDatabase(version):
+                        # initialize the main SB database
+                        upgradeDatabase(DBConnection(), sickbeard.mainDB.InitialSchema)
+                        restored = True
+
+                if not restored:
+                    print "Error in " + str(upgradeClass.__name__) + ": " + ex(e)
+                    raise
         logger.log(upgradeClass.__name__ + " upgrade completed", logger.DEBUG)
     else:
         logger.log(upgradeClass.__name__ + " upgrade not required", logger.DEBUG)

sickbeard/helpers.py

@@ -886,6 +886,44 @@ def backupVersionedFile(old_file, version):
 
     return True
 
+def restoreVersionedFile(backup_file, version):
+    numTries = 0
+
+    new_file, backup_version = os.path.splitext(backup_file)
+    restore_file = new_file + '.' + 'v' + str(version)
+
+    if not ek.ek(os.path.isfile, new_file):
+        logger.log(u"Not restoring, " + new_file + " doesn't exist", logger.DEBUG)
+        return False
+
+    try:
+        logger.log(u"Trying to backup " + new_file + " to " + new_file + '.rtmp before restoring backup', logger.DEBUG)
+        shutil.move(new_file, new_file + '.rtmp')
+    except Exception, e:
+        logger.log(u"Error while trying to backup DB file " + restore_file + " before proceeding with restore: " + ex(e), logger.WARNING)
+        return False
+
+    while not ek.ek(os.path.isfile, new_file):
+        if not ek.ek(os.path.isfile, restore_file):
+            logger.log(u"Not restoring, " + restore_file + " doesn't exist", logger.DEBUG)
+            break
+
+        try:
+            logger.log(u"Trying to restore " + restore_file + " to " + new_file, logger.DEBUG)
+            shutil.copy(restore_file, new_file)
+            logger.log(u"Restore done", logger.DEBUG)
+            break
+        except Exception, e:
+            logger.log(u"Error while trying to restore " + restore_file + ": " + ex(e), logger.WARNING)
+            numTries += 1
+            time.sleep(1)
+            logger.log(u"Trying again.", logger.DEBUG)
+
+        if numTries >= 10:
+            logger.log(u"Unable to restore " + restore_file + " to " + new_file + " please do it manually.", logger.ERROR)
+            return False
+
+    return True
 
 # try to convert to int, if it fails the default will be returned
 def tryInt(s, s_default=0):