From 18e3f9464b8328d40307c5a18de5133a7a3cb337 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 7 Sep 2013 12:55:11 +0530
Subject: [PATCH] Fix dump and restore on memory limited systems

---
 src/calibre/db/backend.py       | 42 +++++++++++++++++----------------
 src/calibre/db/tests/writing.py |  2 ++
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/calibre/db/backend.py b/src/calibre/db/backend.py
index 1dab3cb2f5..513ed92a6c 100644
--- a/src/calibre/db/backend.py
+++ b/src/calibre/db/backend.py
@@ -982,33 +982,35 @@ class DB(object):
         self.conn
 
     def dump_and_restore(self, callback=None, sql=None):
-        from io import StringIO
+        import codecs
+        from calibre.utils.apsw_shell import Shell
         from contextlib import closing
         if callback is None:
             callback = lambda x: x
         uv = int(self.user_version)
 
-        if sql is None:
-            from calibre.utils.apsw_shell import Shell
-            callback(_('Dumping database to SQL') + '...')
-            buf = StringIO()
-            shell = Shell(db=self.conn, stdout=buf)
-            shell.process_command('.dump')
-            sql = buf.getvalue()
-            del shell
-            del buf
+        with TemporaryFile(suffix='.sql') as fname:
+            if sql is None:
+                callback(_('Dumping database to SQL') + '...')
+                with codecs.open(fname, 'wb', encoding='utf-8') as buf:
+                    shell = Shell(db=self.conn, stdout=buf)
+                    shell.process_command('.dump')
+            else:
+                with open(fname, 'wb') as buf:
+                    buf.write(sql if isinstance(sql, bytes) else sql.encode('utf-8'))
 
-        with TemporaryFile(suffix='_tmpdb.db', dir=os.path.dirname(self.dbpath)) as tmpdb:
-            callback(_('Restoring database from SQL') + '...')
-            with closing(Connection(tmpdb)) as conn:
-                conn.execute(sql)
-                conn.execute('PRAGMA user_version=%d;'%uv)
+            with TemporaryFile(suffix='_tmpdb.db', dir=os.path.dirname(self.dbpath)) as tmpdb:
+                callback(_('Restoring database from SQL') + '...')
+                with closing(Connection(tmpdb)) as conn:
+                    shell = Shell(db=conn, encoding='utf-8')
+                    shell.process_command('.read ' + fname)
+                    conn.execute('PRAGMA user_version=%d;'%uv)
 
-            self.close()
-            try:
-                atomic_rename(tmpdb, self.dbpath)
-            finally:
-                self.reopen()
+                self.close()
+                try:
+                    atomic_rename(tmpdb, self.dbpath)
+                finally:
+                    self.reopen()
 
     def vacuum(self):
         self.conn.execute('VACUUM')
diff --git a/src/calibre/db/tests/writing.py b/src/calibre/db/tests/writing.py
index 762517a294..a42d47e67f 100644
--- a/src/calibre/db/tests/writing.py
+++ b/src/calibre/db/tests/writing.py
@@ -606,11 +606,13 @@ class WritingTest(BaseTest):
     def test_dump_and_restore(self):  # {{{
         ' Test roundtripping the db through SQL '
         cache = self.init_cache()
+        uv = int(cache.backend.user_version)
         all_ids = cache.all_book_ids()
         cache.dump_and_restore()
         self.assertEqual(cache.set_field('title', {1:'nt'}), set([1]), 'database connection broken')
         cache = self.init_cache()
         self.assertEqual(cache.all_book_ids(), all_ids, 'dump and restore broke database')
+        self.assertEqual(int(cache.backend.user_version), uv)
     # }}}
 
     def test_set_author_data(self):  # {{{