diff --git a/src/calibre/db/backend.py b/src/calibre/db/backend.py index c6aa2e646f..97dc673ecc 100644 --- a/src/calibre/db/backend.py +++ b/src/calibre/db/backend.py @@ -1059,6 +1059,24 @@ class DB(object): if wam is not None: wam.close_handles() + def add_format(self, book_id, fmt, stream, title, author, path): + fname = self.construct_file_name(book_id, title, author) + path = os.path.join(self.library_path, path) + fmt = ('.' + fmt.lower()) if fmt else '' + dest = os.path.join(path, fname + fmt) + if not os.path.exists(path): + os.makedirs(path) + size = 0 + + if (not getattr(stream, 'name', False) or not samefile(dest, stream.name)): + with lopen(dest, 'wb') as f: + shutil.copyfileobj(stream, f) + size = f.tell() + elif os.path.exists(dest): + size = os.path.getsize(dest) + + return size, fname + def update_path(self, book_id, title, author, path_field, formats_field): path = self.construct_path_name(book_id, title, author) current_path = path_field.for_book(book_id) diff --git a/src/calibre/db/cache.py b/src/calibre/db/cache.py index 88f06b43ba..4f7de11269 100644 --- a/src/calibre/db/cache.py +++ b/src/calibre/db/cache.py @@ -7,12 +7,13 @@ __license__ = 'GPL v3' __copyright__ = '2011, Kovid Goyal ' __docformat__ = 'restructuredtext en' -import os, traceback, random +import os, traceback, random, shutil from io import BytesIO from collections import defaultdict from functools import wraps, partial from calibre.constants import iswindows +from calibre.customize.ui import run_plugins_on_import, run_plugins_on_postimport from calibre.db import SPOOL_SIZE from calibre.db.categories import get_categories from calibre.db.locking import create_locks @@ -22,6 +23,7 @@ from calibre.db.search import Search from calibre.db.tables import VirtualTable from calibre.db.write import get_series_values from calibre.db.lazy import FormatMetadata, FormatsList +from calibre.ebooks import check_ebook_format from calibre.ebooks.metadata import string_to_authors from calibre.ebooks.metadata.book.base import Metadata from calibre.ebooks.metadata.opf2 import metadata_to_opf @@ -51,6 +53,18 @@ def wrap_simple(lock, func): return func(*args, **kwargs) return ans +def run_import_plugins(path_or_stream, fmt): + fmt = fmt.lower() + if hasattr(path_or_stream, 'seek'): + path_or_stream.seek(0) + pt = PersistentTemporaryFile('_import_plugin.'+fmt) + shutil.copyfileobj(path_or_stream, pt, 1024**2) + pt.close() + path = pt.name + else: + path = path_or_stream + return run_plugins_on_import(path, fmt) + class Cache(object): @@ -943,6 +957,43 @@ class Cache(object): if extra is not None or force_changes: protected_set_field(idx, extra) + @write_api + def add_format(self, book_id, fmt, stream_or_path, replace=True, run_hooks=True, dbapi=None): + if run_hooks: + # Run import plugins + npath = run_import_plugins(stream_or_path, fmt) + fmt = os.path.splitext(npath)[-1].lower().replace('.', '').upper() + stream_or_path = lopen(npath, 'rb') + fmt = check_ebook_format(stream_or_path, fmt) + + fmt = (fmt or '').upper() + self.format_metadata_cache[book_id].pop(fmt, None) + try: + name = self.fields['formats'].format_fname(book_id, fmt) + except: + name = None + + if name and not replace: + return False + + path = self._field_for('path', book_id).replace('/', os.sep) + title = self._field_for('title', book_id, default_value=_('Unknown')) + author = self._field_for('authors', book_id, default_value=(_('Unknown'),))[0] + stream = stream_or_path if hasattr(stream_or_path, 'read') else lopen(stream_or_path, 'rb') + size, fname = self.backend.add_format(book_id, fmt, stream, title, author, path) + del stream + + max_size = self.fields['formats'].table.update_fmt(book_id, fmt, fname, size, self.backend) + self.fields['size'].table.update_size(book_id, max_size) + self._update_last_modified((book_id,)) + + if run_hooks: + # Run post import plugins + run_plugins_on_postimport(dbapi or self, book_id, fmt) + stream_or_path.close() + + return True + # }}} class SortKey(object): # {{{ @@ -959,3 +1010,4 @@ class SortKey(object): # {{{ return 0 # }}} + diff --git a/src/calibre/db/tables.py b/src/calibre/db/tables.py index 83d4b23712..4dba10abff 100644 --- a/src/calibre/db/tables.py +++ b/src/calibre/db/tables.py @@ -8,6 +8,7 @@ __copyright__ = '2011, Kovid Goyal ' __docformat__ = 'restructuredtext en' from datetime import datetime +from collections import defaultdict from dateutil.tz import tzoffset @@ -98,6 +99,9 @@ class SizeTable(OneToOneTable): 'WHERE data.book=books.id) FROM books'): self.book_col_map[row[0]] = self.unserialize(row[1]) + def update_size(self, book_id, size): + self.book_col_map[book_id] = size + class UUIDTable(OneToOneTable): def read(self, db): @@ -194,8 +198,9 @@ class FormatsTable(ManyToManyTable): pass def read_maps(self, db): - self.fname_map = {} - for row in db.conn.execute('SELECT book, format, name FROM data'): + self.fname_map = defaultdict(dict) + self.size_map = defaultdict(dict) + for row in db.conn.execute('SELECT book, format, name, uncompressed_size FROM data'): if row[1] is not None: fmt = row[1].upper() if fmt not in self.col_book_map: @@ -204,9 +209,8 @@ class FormatsTable(ManyToManyTable): if row[0] not in self.book_col_map: self.book_col_map[row[0]] = [] self.book_col_map[row[0]].append(fmt) - if row[0] not in self.fname_map: - self.fname_map[row[0]] = {} self.fname_map[row[0]][fmt] = row[2] + self.size_map[row[0]][fmt] = row[3] for key in tuple(self.book_col_map.iterkeys()): self.book_col_map[key] = tuple(sorted(self.book_col_map[key])) @@ -216,6 +220,13 @@ class FormatsTable(ManyToManyTable): db.conn.execute('UPDATE data SET name=? WHERE book=? AND format=?', (fname, book_id, fmt)) + def update_fmt(self, book_id, fmt, fname, size, db): + self.fname_map[book_id][fmt] = fname + self.size_map[book_id][fmt] = size + db.conn.execute('INSERT OR REPLACE INTO data (book,format,uncompressed_size,name) VALUES (?,?,?,?)', + (book_id, fmt, size, fname)) + return max(self.size_map[book_id].itervalues()) + class IdentifiersTable(ManyToManyTable): def read_id_maps(self, db):