Fix dump and restore on memory limited systems

This commit is contained in:
Kovid Goyal 2013-09-07 12:55:11 +05:30
parent 00bd5a658e
commit 18e3f9464b
2 changed files with 24 additions and 20 deletions

View File

@ -982,33 +982,35 @@ class DB(object):
self.conn self.conn
def dump_and_restore(self, callback=None, sql=None): 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 from contextlib import closing
if callback is None: if callback is None:
callback = lambda x: x callback = lambda x: x
uv = int(self.user_version) uv = int(self.user_version)
if sql is None: with TemporaryFile(suffix='.sql') as fname:
from calibre.utils.apsw_shell import Shell if sql is None:
callback(_('Dumping database to SQL') + '...') callback(_('Dumping database to SQL') + '...')
buf = StringIO() with codecs.open(fname, 'wb', encoding='utf-8') as buf:
shell = Shell(db=self.conn, stdout=buf) shell = Shell(db=self.conn, stdout=buf)
shell.process_command('.dump') shell.process_command('.dump')
sql = buf.getvalue() else:
del shell with open(fname, 'wb') as buf:
del 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: with TemporaryFile(suffix='_tmpdb.db', dir=os.path.dirname(self.dbpath)) as tmpdb:
callback(_('Restoring database from SQL') + '...') callback(_('Restoring database from SQL') + '...')
with closing(Connection(tmpdb)) as conn: with closing(Connection(tmpdb)) as conn:
conn.execute(sql) shell = Shell(db=conn, encoding='utf-8')
conn.execute('PRAGMA user_version=%d;'%uv) shell.process_command('.read ' + fname)
conn.execute('PRAGMA user_version=%d;'%uv)
self.close() self.close()
try: try:
atomic_rename(tmpdb, self.dbpath) atomic_rename(tmpdb, self.dbpath)
finally: finally:
self.reopen() self.reopen()
def vacuum(self): def vacuum(self):
self.conn.execute('VACUUM') self.conn.execute('VACUUM')

View File

@ -606,11 +606,13 @@ class WritingTest(BaseTest):
def test_dump_and_restore(self): # {{{ def test_dump_and_restore(self): # {{{
' Test roundtripping the db through SQL ' ' Test roundtripping the db through SQL '
cache = self.init_cache() cache = self.init_cache()
uv = int(cache.backend.user_version)
all_ids = cache.all_book_ids() all_ids = cache.all_book_ids()
cache.dump_and_restore() cache.dump_and_restore()
self.assertEqual(cache.set_field('title', {1:'nt'}), set([1]), 'database connection broken') self.assertEqual(cache.set_field('title', {1:'nt'}), set([1]), 'database connection broken')
cache = self.init_cache() cache = self.init_cache()
self.assertEqual(cache.all_book_ids(), all_ids, 'dump and restore broke database') 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): # {{{ def test_set_author_data(self): # {{{