import json import os import time import traceback from couchpotato import CPLog from couchpotato.api import addApiView from couchpotato.core.event import addEvent, fireEvent from couchpotato.core.helpers.encoding import toUnicode from couchpotato.core.helpers.variable import getImdb, tryInt log = CPLog(__name__) class Database(object): indexes = [] db = None def __init__(self): addApiView('database.list_documents', self.listDocuments) addApiView('database.reindex', self.reindex) addApiView('database.compact', self.compact) addApiView('database.document.update', self.updateDocument) addApiView('database.document.delete', self.deleteDocument) addEvent('database.setup.after', self.startup_compact) addEvent('database.setup_index', self.setupIndex) addEvent('app.migrate', self.migrate) addEvent('app.after_shutdown', self.close) def getDB(self): if not self.db: from couchpotato import get_db self.db = get_db() return self.db def close(self, **kwargs): self.getDB().close() def setupIndex(self, index_name, klass): self.indexes.append(index_name) db = self.getDB() # Category index index_instance = klass(db.path, index_name) try: db.add_index(index_instance) db.reindex_index(index_name) except: previous = db.indexes_names[index_name] previous_version = previous._version current_version = klass._version # Only edit index if versions are different if previous_version < current_version: log.debug('Index "%s" already exists, updating and reindexing', index_name) db.destroy_index(previous) db.add_index(index_instance) db.reindex_index(index_name) def deleteDocument(self, **kwargs): db = self.getDB() try: document_id = kwargs.get('_request').get_argument('id') document = db.get('id', document_id) db.delete(document) return { 'success': True } except: return { 'success': False, 'error': traceback.format_exc() } def updateDocument(self, **kwargs): db = self.getDB() try: document = json.loads(kwargs.get('_request').get_argument('document')) d = db.update(document) document.update(d) return { 'success': True, 'document': document } except: return { 'success': False, 'error': traceback.format_exc() } def listDocuments(self, **kwargs): db = self.getDB() results = { 'unknown': [] } for document in db.all('id'): key = document.get('_t', 'unknown') if kwargs.get('show') and key != kwargs.get('show'): continue if not results.get(key): results[key] = [] results[key].append(document) return results def reindex(self, **kwargs): success = True try: db = self.getDB() db.reindex() except: log.error('Failed index: %s', traceback.format_exc()) success = False return { 'success': success } def compact(self, **kwargs): success = True try: start = time.time() db = self.getDB() size = float(db.get_db_details().get('size', 0)) log.debug('Compacting database, current size: %sMB', round(size/1048576, 2)) db.compact() new_size = float(db.get_db_details().get('size', 0)) log.debug('Done compacting database in %ss, new size: %sMB, saved: %sMB', (round(time.time()-start, 2), round(new_size/1048576, 2), round((size-new_size)/1048576, 2))) except: log.error('Failed compact: %s', traceback.format_exc()) success = False return { 'success': success } # Compact on start def startup_compact(self): from couchpotato import Env db = self.getDB() size = db.get_db_details().get('size') prop_name = 'last_db_compact' last_check = int(Env.prop(prop_name, default = 0)) if size > 26214400 and last_check < time.time()-604800: # 25MB / 7 days self.compact() Env.prop(prop_name, value = int(time.time())) def migrate(self): from couchpotato import Env old_db = os.path.join(Env.get('data_dir'), 'couchpotato.db') if not os.path.isfile(old_db): return log.info('=' * 30) log.info('Migrating database, hold on..') time.sleep(1) if os.path.isfile(old_db): migrate_start = time.time() import sqlite3 conn = sqlite3.connect(old_db) migrate_list = { 'category': ['id', 'label', 'order', 'required', 'preferred', 'ignored', 'destination'], 'profile': ['id', 'label', 'order', 'core', 'hide'], 'profiletype': ['id', 'order', 'finish', 'wait_for', 'quality_id', 'profile_id'], 'quality': ['id', 'identifier', 'order', 'size_min', 'size_max'], 'movie': ['id', 'last_edit', 'library_id', 'status_id', 'profile_id', 'category_id'], 'library': ['id', 'identifier', 'info'], 'librarytitle': ['id', 'title', 'default', 'libraries_id'], 'library_files__file_library': ['library_id', 'file_id'], 'release': ['id', 'identifier', 'movie_id', 'status_id', 'quality_id', 'last_edit'], 'releaseinfo': ['id', 'identifier', 'value', 'release_id'], 'release_files__file_release': ['release_id', 'file_id'], 'status': ['id', 'identifier'], 'properties': ['id', 'identifier', 'value'], 'file': ['id', 'path', 'type_id'], 'filetype': ['identifier', 'id'] } migrate_data = {} c = conn.cursor() for ml in migrate_list: migrate_data[ml] = {} rows = migrate_list[ml] try: c.execute('SELECT %s FROM `%s`' % ('`' + '`,`'.join(rows) + '`', ml)) except: # ignore faulty destination_id database if ml == 'category': migrate_data[ml] = {} else: raise for p in c.fetchall(): columns = {} for row in migrate_list[ml]: columns[row] = p[rows.index(row)] if not migrate_data[ml].get(p[0]): migrate_data[ml][p[0]] = columns else: if not isinstance(migrate_data[ml][p[0]], list): migrate_data[ml][p[0]] = [migrate_data[ml][p[0]]] migrate_data[ml][p[0]].append(columns) conn.close() log.info('Getting data took %s', time.time() - migrate_start) db = self.getDB() # Use properties properties = migrate_data['properties'] log.info('Importing %s properties', len(properties)) for x in properties: property = properties[x] Env.prop(property.get('identifier'), property.get('value')) # Categories categories = migrate_data.get('category', []) log.info('Importing %s categories', len(categories)) category_link = {} for x in categories: c = categories[x] new_c = db.insert({ '_t': 'category', 'order': c.get('order', 999), 'label': toUnicode(c.get('label', '')), 'ignored': toUnicode(c.get('ignored', '')), 'preferred': toUnicode(c.get('preferred', '')), 'required': toUnicode(c.get('required', '')), 'destination': toUnicode(c.get('destination', '')), }) category_link[x] = new_c.get('_id') # Profiles log.info('Importing profiles') new_profiles = db.all('profile', with_doc = True) new_profiles_by_label = {} for x in new_profiles: # Remove default non core profiles if not x['doc'].get('core'): db.delete(x['doc']) else: new_profiles_by_label[x['doc']['label']] = x['_id'] profiles = migrate_data['profile'] profile_link = {} for x in profiles: p = profiles[x] exists = new_profiles_by_label.get(p.get('label')) # Update existing with order only if exists and p.get('core'): profile = db.get('id', exists) profile['order'] = tryInt(p.get('order')) profile['hide'] = p.get('hide') in [1, True, 'true', 'True'] db.update(profile) profile_link[x] = profile.get('_id') else: new_profile = { '_t': 'profile', 'label': p.get('label'), 'order': int(p.get('order', 999)), 'core': p.get('core', False), 'qualities': [], 'wait_for': [], 'finish': [] } types = migrate_data['profiletype'] for profile_type in types: p_type = types[profile_type] if types[profile_type]['profile_id'] == p['id']: if p_type['quality_id']: new_profile['finish'].append(p_type['finish']) new_profile['wait_for'].append(p_type['wait_for']) new_profile['qualities'].append(migrate_data['quality'][p_type['quality_id']]['identifier']) if len(new_profile['qualities']) > 0: new_profile.update(db.insert(new_profile)) profile_link[x] = new_profile.get('_id') else: log.error('Corrupt profile list for "%s", using default.', p.get('label')) # Qualities log.info('Importing quality sizes') new_qualities = db.all('quality', with_doc = True) new_qualities_by_identifier = {} for x in new_qualities: new_qualities_by_identifier[x['doc']['identifier']] = x['_id'] qualities = migrate_data['quality'] quality_link = {} for x in qualities: q = qualities[x] q_id = new_qualities_by_identifier[q.get('identifier')] quality = db.get('id', q_id) quality['order'] = q.get('order') quality['size_min'] = tryInt(q.get('size_min')) quality['size_max'] = tryInt(q.get('size_max')) db.update(quality) quality_link[x] = quality # Titles titles = migrate_data['librarytitle'] titles_by_library = {} for x in titles: title = titles[x] if title.get('default'): titles_by_library[title.get('libraries_id')] = title.get('title') # Releases releaseinfos = migrate_data['releaseinfo'] for x in releaseinfos: info = releaseinfos[x] # Skip if release doesn't exist for this info if not migrate_data['release'].get(info.get('release_id')): continue if not migrate_data['release'][info.get('release_id')].get('info'): migrate_data['release'][info.get('release_id')]['info'] = {} migrate_data['release'][info.get('release_id')]['info'][info.get('identifier')] = info.get('value') releases = migrate_data['release'] releases_by_media = {} for x in releases: release = releases[x] if not releases_by_media.get(release.get('movie_id')): releases_by_media[release.get('movie_id')] = [] releases_by_media[release.get('movie_id')].append(release) # Type ids types = migrate_data['filetype'] type_by_id = {} for t in types: type = types[t] type_by_id[type.get('id')] = type # Media log.info('Importing %s media items', len(migrate_data['movie'])) statuses = migrate_data['status'] libraries = migrate_data['library'] library_files = migrate_data['library_files__file_library'] releases_files = migrate_data['release_files__file_release'] all_files = migrate_data['file'] poster_type = migrate_data['filetype']['poster'] medias = migrate_data['movie'] for x in medias: m = medias[x] status = statuses.get(m['status_id']).get('identifier') l = libraries.get(m['library_id']) # Only migrate wanted movies, Skip if no identifier present if not l or not getImdb(l.get('identifier')): continue profile_id = profile_link.get(m['profile_id']) category_id = category_link.get(m['category_id']) title = titles_by_library.get(m['library_id']) releases = releases_by_media.get(x, []) info = json.loads(l.get('info', '')) files = library_files.get(m['library_id'], []) if not isinstance(files, list): files = [files] added_media = fireEvent('movie.add', { 'info': info, 'identifier': l.get('identifier'), 'profile_id': profile_id, 'category_id': category_id, 'title': title }, force_readd = False, search_after = False, update_after = False, notify_after = False, status = status, single = True) if not added_media: log.error('Failed adding media %s: %s', (l.get('identifier'), info)) continue added_media['files'] = added_media.get('files', {}) for f in files: ffile = all_files[f.get('file_id')] # Only migrate posters if ffile.get('type_id') == poster_type.get('id'): if ffile.get('path') not in added_media['files'].get('image_poster', []) and os.path.isfile(ffile.get('path')): added_media['files']['image_poster'] = [ffile.get('path')] break if 'image_poster' in added_media['files']: db.update(added_media) for rel in releases: empty_info = False if not rel.get('info'): empty_info = True rel['info'] = {} quality = quality_link.get(rel.get('quality_id')) if not quality: continue release_status = statuses.get(rel.get('status_id')).get('identifier') if rel['info'].get('download_id'): status_support = rel['info'].get('download_status_support', False) in [True, 'true', 'True'] rel['info']['download_info'] = { 'id': rel['info'].get('download_id'), 'downloader': rel['info'].get('download_downloader'), 'status_support': status_support, } # Add status to keys rel['info']['status'] = release_status if not empty_info: fireEvent('release.create_from_search', [rel['info']], added_media, quality, single = True) else: release = { '_t': 'release', 'identifier': rel.get('identifier'), 'media_id': added_media.get('_id'), 'quality': quality.get('identifier'), 'status': release_status, 'last_edit': int(time.time()), 'files': {} } # Add downloader info if provided try: release['download_info'] = rel['info']['download_info'] del rel['download_info'] except: pass # Add files release_files = releases_files.get(rel.get('id'), []) if not isinstance(release_files, list): release_files = [release_files] if len(release_files) == 0: continue for f in release_files: rfile = all_files[f.get('file_id')] file_type = type_by_id.get(rfile.get('type_id')).get('identifier') if not release['files'].get(file_type): release['files'][file_type] = [] release['files'][file_type].append(rfile.get('path')) try: rls = db.get('release_identifier', rel.get('identifier'), with_doc = True)['doc'] rls.update(release) db.update(rls) except: db.insert(release) log.info('Total migration took %s', time.time() - migrate_start) log.info('=' * 30) # rename old database log.info('Renaming old database to %s ', old_db + '.old') os.rename(old_db, old_db + '.old') if os.path.isfile(old_db + '-wal'): os.rename(old_db + '-wal', old_db + '-wal.old') if os.path.isfile(old_db + '-shm'): os.rename(old_db + '-shm', old_db + '-shm.old')