__init__.py 12.5 KB
Newer Older
1
# Author: echel0n <[email protected]>
echel0n's avatar
echel0n committed
2
# URL: https://sickrage.ca
3
#
echel0n's avatar
echel0n committed
4
# This file is part of SickRage.
5
#
echel0n's avatar
echel0n committed
6
# SickRage is free software: you can redistribute it and/or modify
7
8
9
10
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
echel0n's avatar
echel0n committed
11
# SickRage is distributed in the hope that it will be useful,
12
13
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14
#  GNU General Public License for more details.
15
16
#
# You should have received a copy of the GNU General Public License
echel0n's avatar
echel0n committed
17
# along with SickRage.  If not, see <http://www.gnu.org/licenses/>.
18

19
from __future__ import unicode_literals
20

21
import os
22
import re
23
24
25
26
import shutil
import tarfile
import time
import traceback
echel0n's avatar
echel0n committed
27
from sqlite3 import OperationalError
echel0n's avatar
echel0n committed
28

29
from CodernityDB.database_super_thread_safe import SuperThreadSafeDatabase
echel0n's avatar
v8.7.9    
echel0n committed
30
from CodernityDB.index import IndexNotFoundException, IndexConflict, IndexException
31
from CodernityDB.storage import IU_Storage
32
33

import sickrage
echel0n's avatar
echel0n committed
34
from sickrage.core.helpers import randomString
35

echel0n's avatar
echel0n committed
36

37
38
39
40
41
42
def Custom_IU_Storage_get(self, start, size, status='c'):
    if status == 'd':
        return None
    else:
        self._f.seek(start)
        return self.data_from(self._f.read(size))
43

echel0n's avatar
echel0n committed
44

45
class srDatabase(object):
echel0n's avatar
echel0n committed
46
    _indexes = {}
echel0n's avatar
echel0n committed
47
    _migrate_list = {}
48

49
    def __init__(self, name=''):
echel0n's avatar
echel0n committed
50
51
52
        self.name = name
        self.old_db_path = ''

53
        self.db_path = os.path.join(sickrage.app.data_dir, 'database', self.name)
54
        self.db = SuperThreadSafeDatabase(self.db_path)
echel0n's avatar
echel0n committed
55

56
57
    def initialize(self):
        # Remove database folder if both exists
echel0n's avatar
echel0n committed
58
        if self.db.exists() and os.path.isfile(self.old_db_path):
echel0n's avatar
v8.7.9    
echel0n committed
59
            self.db.open()
60
            self.db.destroy()
61

62
        if self.db.exists():
63
            self.backup()
echel0n's avatar
v8.7.9    
echel0n committed
64
            self.db.open()
65
66
67
68
        else:
            self.db.create()

        # setup database indexes
echel0n's avatar
echel0n committed
69
        self.setup_indexes()
70

71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    def backup(self):
        # Backup before start and cleanup old backups
        backup_path = os.path.join(sickrage.app.data_dir, 'db_backup', self.name)
        backup_count = 5
        existing_backups = []

        if not os.path.isdir(backup_path):
            os.makedirs(backup_path)

        for root, dirs, files in os.walk(backup_path):
            # Only consider files being a direct child of the backup_path
            if root == backup_path:
                for backup_file in sorted(files):
                    ints = re.findall('\d+', backup_file)

                    # Delete non zip files
                    if len(ints) != 1:
                        try:
                            os.remove(os.path.join(root, backup_file))
                        except:
                            pass
                    else:
                        existing_backups.append((int(ints[0]), backup_file))
            else:
                # Delete stray directories.
                shutil.rmtree(root)

        # Remove all but the last 5
        for eb in existing_backups[:-backup_count]:
            os.remove(os.path.join(backup_path, eb[1]))

        # Create new backup
        new_backup = os.path.join(backup_path, '%s.tar.gz' % int(time.time()))
        with tarfile.open(new_backup, 'w:gz') as zipf:
            for root, dirs, files in os.walk(self.db_path):
                for zfilename in files:
                    zipf.add(os.path.join(root, zfilename),
                             arcname='database/%s/%s' % (
                                 self.name,
                                 os.path.join(root[len(self.db_path) + 1:], zfilename))
                             )

113
114
115
116
117
118
119
120
121
122
    def compact(self, try_repair=True, **kwargs):
        # Removing left over compact files
        for f in os.listdir(self.db.path):
            for x in ['_compact_buck', '_compact_stor']:
                if f[-len(x):] == x:
                    os.unlink(os.path.join(self.db.path, f))

        try:
            start = time.time()
            size = float(self.db.get_db_details().get('size', 0))
123
            sickrage.app.log.info(
echel0n's avatar
echel0n committed
124
                'Compacting {} database, current size: {}MB'.format(self.name, round(size / 1048576, 2)))
125
126

            self.db.compact()
127

128
            new_size = float(self.db.get_db_details().get('size', 0))
129
            sickrage.app.log.info(
130
131
132
133
                'Done compacting {} database in {}s, new size: {}MB, saved: {}MB'.format(
                    self.name, round(time.time() - start, 2),
                    round(new_size / 1048576, 2), round((size - new_size) / 1048576, 2))
            )
134
        except (IndexException, AttributeError, TypeError) as e:
135
            if try_repair:
136
                sickrage.app.log.debug('Something wrong with indexes, trying repair')
echel0n's avatar
echel0n committed
137
138
139
140
141
142
143
144
145

                # Remove all indexes
                old_indexes = self._indexes.keys()
                for index_name in old_indexes:
                    try:
                        self.db.destroy_index(index_name)
                    except IndexNotFoundException:
                        pass
                    except:
146
                        sickrage.app.log.debug('Failed removing old index %s', index_name)
echel0n's avatar
echel0n committed
147
148
149
150
151
152
153
154
155

                # Add them again
                for index_name in self._indexes:
                    try:
                        self.db.add_index(self._indexes[index_name](self.db.path, index_name))
                        self.db.reindex_index(index_name)
                    except IndexConflict:
                        pass
                    except:
156
                        sickrage.app.log.debug('Failed adding index %s', index_name)
echel0n's avatar
echel0n committed
157
158
                        raise

159
160
                self.compact(try_repair=False)
            else:
161
                sickrage.app.log.debug('Failed compact: {}'.format(traceback.format_exc()))
162
        except:
163
            sickrage.app.log.debug('Failed compact: {}'.format(traceback.format_exc()))
164

echel0n's avatar
echel0n committed
165
    def setup_indexes(self):
166
167
168
        # setup database indexes
        for index_name in self._indexes:
            try:
echel0n's avatar
echel0n committed
169
170
171
172
173
174
175
176
177
178
179
180
181
                # Make sure store and bucket don't exist
                exists = []
                for x in ['buck', 'stor']:
                    full_path = os.path.join(self.db.path, '%s_%s' % (index_name, x))
                    if os.path.exists(full_path):
                        exists.append(full_path)

                if index_name not in self.db.indexes_names:
                    # Remove existing buckets if index isn't there
                    for x in exists:
                        os.unlink(x)

                    self.db.add_index(self._indexes[index_name](self.db.path, index_name))
182
                    self.db.reindex_index(index_name)
echel0n's avatar
echel0n committed
183
184
185
186
187
                else:
                    # Previous info
                    previous_version = self.db.indexes_names[index_name]._version
                    current_version = self._indexes[index_name]._version

188
                    self.check_versions(index_name, current_version, previous_version)
189
            except:
190
                sickrage.app.log.debug('Failed adding index {}'.format(index_name))
191

192
193
194
195
196
197
198
    def check_versions(self, index_name, current_version, previous_version):
        # Only edit index if versions are different
        if previous_version < current_version:
            self.db.destroy_index(self.db.indexes_names[index_name])
            self.db.add_index(self._indexes[index_name](self.db.path, index_name))
            self.db.reindex_index(index_name)

echel0n's avatar
echel0n committed
199
200
201
    def close(self):
        self.db.close()

202
    def upgrade(self):
203
204
        pass

205
206
207
    def cleanup(self):
        pass

208
209
210
211
212
213
214
215
216
217
218
219
220
221
    @property
    def version(self):
        try:
            dbData = list(self.all('version'))[-1]
        except IndexError:
            dbData = {
                '_t': 'version',
                'database_version': 1
            }

            dbData.update(self.insert(dbData))

        return dbData['database_version']

222
223
224
225
    @property
    def opened(self):
        return self.db.opened

echel0n's avatar
echel0n committed
226
227
228
229
230
    def check_integrity(self):
        for index_name in self._indexes:
            try:
                for x in self.db.all(index_name):
                    try:
231
                        self.get('id', x.get('_id'))
echel0n's avatar
echel0n committed
232
                    except (ValueError, TypeError) as e:
233
                        self.delete(self.get(index_name, x.get('key')))
echel0n's avatar
echel0n committed
234
235
236
            except Exception as e:
                if index_name in self.db.indexes_names:
                    self.db.destroy_index(self.db.indexes_names[index_name])
237
238
                self.db.add_index(self._indexes[index_name](self.db.path, index_name))
                #self.db.reindex_index(index_name)
echel0n's avatar
echel0n committed
239

240
    def migrate(self):
echel0n's avatar
echel0n committed
241
        if os.path.isfile(self.old_db_path):
242
243
            sickrage.app.log.info('=' * 30)
            sickrage.app.log.info('Migrating %s database, please wait...', self.name)
echel0n's avatar
echel0n committed
244
245
246
247
            migrate_start = time.time()

            import sqlite3
            conn = sqlite3.connect(self.old_db_path)
248
            conn.text_factory = lambda x: (x.decode('utf-8', 'ignore'))
echel0n's avatar
echel0n committed
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278

            migrate_data = {}
            rename_old = False

            try:
                c = conn.cursor()

                for ml in self._migrate_list:
                    migrate_data[ml] = {}
                    rows = self._migrate_list[ml]

                    try:
                        c.execute('SELECT {} FROM `{}`'.format('`' + '`,`'.join(rows) + '`', ml))
                    except:
                        # ignore faulty destination_id database
                        rename_old = True
                        raise

                    for p in c.fetchall():
                        columns = {}
                        for row in self._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)

279
                sickrage.app.log.info('Getting data took %s', (time.time() - migrate_start))
echel0n's avatar
echel0n committed
280
281
282
283
284
285

                if not self.db.opened:
                    return

                for t_name in migrate_data:
                    t_data = migrate_data.get(t_name, {})
286
                    sickrage.app.log.info('Importing %s %s' % (len(t_data), t_name))
echel0n's avatar
echel0n committed
287
288
289
290
                    for k, v in t_data.items():
                        if isinstance(v, list):
                            for d in v:
                                d.update({'_t': t_name})
291
                                self.insert(d)
echel0n's avatar
echel0n committed
292
293
                        else:
                            v.update({'_t': t_name})
294
                            self.insert(v)
echel0n's avatar
echel0n committed
295

296
297
                sickrage.app.log.info('Total migration took %s', (time.time() - migrate_start))
                sickrage.app.log.info('=' * 30)
echel0n's avatar
echel0n committed
298
299

                rename_old = True
300
            except OperationalError:
301
                sickrage.app.log.debug('Migrating from unsupported/corrupt %s database version', self.name)
echel0n's avatar
echel0n committed
302
303
                rename_old = True
            except:
304
                sickrage.app.log.debug('Migration of %s database failed', self.name)
305
306
            finally:
                conn.close()
echel0n's avatar
echel0n committed
307
308

            # rename old database
echel0n's avatar
echel0n committed
309
            if rename_old:
echel0n's avatar
echel0n committed
310
                random = randomString()
311
                sickrage.app.log.info('Renaming old database to %s.%s_old' % (self.old_db_path, random))
echel0n's avatar
echel0n committed
312
313
314
315
316
317
                os.rename(self.old_db_path, '{}.{}_old'.format(self.old_db_path, random))

                if os.path.isfile(self.old_db_path + '-wal'):
                    os.rename(self.old_db_path + '-wal', '{}-wal.{}_old'.format(self.old_db_path, random))
                if os.path.isfile(self.old_db_path + '-shm'):
                    os.rename(self.old_db_path + '-shm', '{}-shm.{}_old'.format(self.old_db_path, random))
318

319
    def all(self, *args, **kwargs):
320
        return (x['doc'] for x in self.db.all(with_doc=True, *args, **kwargs))
321

322
    def get_many(self, *args, **kwargs):
323
        return (x['doc'] for x in self.db.get_many(with_doc=True, *args, **kwargs))
324
325

    def get(self, *args, **kwargs):
326
327
        x = self.db.get(with_doc=True, *args, **kwargs)
        return x.get('doc', x)
328

echel0n's avatar
echel0n committed
329
    def delete(self, *args):
330
        return self.db.delete(*args)
echel0n's avatar
echel0n committed
331

332
    def update(self, *args):
333
        return self.db.update(*args)
334
335

    def insert(self, *args):
336
        return self.db.insert(*args)
echel0n's avatar
echel0n committed
337

338

339
# Monkey-Patch storage to suppress logging messages
echel0n's avatar
echel0n committed
340
IU_Storage.get = Custom_IU_Storage_get