import os import sqlite3 import time import types from . import common from . import exceptions from . import pushshift from voussoirkit import pathclass from voussoirkit import sqlhelpers from voussoirkit import vlogging log = vlogging.get_logger(__name__) # For backwards compatibility reasons, this list of format strings will help # timesearch find databases that are using the old filename style. # The final element will be used if none of the previous ones were found. DB_FORMATS_SUBREDDIT = [ '.\\{name}.db', '.\\subreddits\\{name}\\{name}.db', '.\\{name}\\{name}.db', '.\\databases\\{name}.db', '.\\subreddits\\{name}\\{name}.db', ] DB_FORMATS_USER = [ '.\\@{name}.db', '.\\users\\@{name}\\@{name}.db', '.\\@{name}\\@{name}.db', '.\\databases\\@{name}.db', '.\\users\\@{name}\\@{name}.db', ] DATABASE_VERSION = 2 DB_VERSION_PRAGMA = f''' PRAGMA user_version = {DATABASE_VERSION}; ''' DB_PRAGMAS = f''' ''' DB_INIT = f''' {DB_PRAGMAS} {DB_VERSION_PRAGMA} ---------------------------------------------------------------------------------------------------- CREATE TABLE IF NOT EXISTS config( key TEXT, value TEXT ); ---------------------------------------------------------------------------------------------------- CREATE TABLE IF NOT EXISTS submissions( idint INT, idstr TEXT, created INT, self INT, nsfw INT, author TEXT, title TEXT, url TEXT, selftext TEXT, score INT, subreddit TEXT, distinguish INT, textlen INT, num_comments INT, flair_text TEXT, flair_css_class TEXT, augmented_at INT, augmented_count INT ); CREATE INDEX IF NOT EXISTS submission_index ON submissions(idstr); ---------------------------------------------------------------------------------------------------- CREATE TABLE IF NOT EXISTS comments( idint INT, idstr TEXT, created INT, author TEXT, parent TEXT, submission TEXT, body TEXT, score INT, subreddit TEXT, distinguish TEXT, textlen INT ); CREATE INDEX IF NOT EXISTS comment_index ON comments(idstr); ---------------------------------------------------------------------------------------------------- CREATE TABLE IF NOT EXISTS submission_edits( idstr TEXT, previous_selftext TEXT, replaced_at INT ); CREATE INDEX IF NOT EXISTS submission_edits_index ON submission_edits(idstr); ---------------------------------------------------------------------------------------------------- CREATE TABLE IF NOT EXISTS comment_edits( idstr TEXT, previous_body TEXT, replaced_at INT ); CREATE INDEX IF NOT EXISTS comment_edits_index ON comment_edits(idstr); ''' DEFAULT_CONFIG = { 'store_edits': True, } SQL_SUBMISSION_COLUMNS = [ 'idint', 'idstr', 'created', 'self', 'nsfw', 'author', 'title', 'url', 'selftext', 'score', 'subreddit', 'distinguish', 'textlen', 'num_comments', 'flair_text', 'flair_css_class', 'augmented_at', 'augmented_count', ] SQL_COMMENT_COLUMNS = [ 'idint', 'idstr', 'created', 'author', 'parent', 'submission', 'body', 'score', 'subreddit', 'distinguish', 'textlen', ] SQL_EDITS_COLUMNS = [ 'idstr', 'text', 'replaced_at', ] SQL_SUBMISSION = {key:index for (index, key) in enumerate(SQL_SUBMISSION_COLUMNS)} SQL_COMMENT = {key:index for (index, key) in enumerate(SQL_COMMENT_COLUMNS)} SUBMISSION_TYPES = (common.praw.models.Submission, pushshift.DummySubmission) COMMENT_TYPES = (common.praw.models.Comment, pushshift.DummyComment) class DBEntry: ''' This class converts a tuple row from the database into an object so that you can access the attributes with dot notation. ''' def __init__(self, dbrow): if dbrow[1].startswith('t3_'): columns = SQL_SUBMISSION_COLUMNS self.object_type = 'submission' else: columns = SQL_COMMENT_COLUMNS self.object_type = 'comment' self.id = None self.idstr = None for (index, attribute) in enumerate(columns): setattr(self, attribute, dbrow[index]) def __repr__(self): return 'DBEntry(\'%s\')' % self.id class TSDB: def __init__(self, filepath, *, do_create=True, skip_version_check=False): self.filepath = pathclass.Path(filepath) if not self.filepath.is_file: if not do_create: raise exceptions.DatabaseNotFound(self.filepath) print('New database', self.filepath.relative_path) self.filepath.parent.makedirs(exist_ok=True) self.breakdown_dir = self.filepath.parent.with_child('breakdown') self.offline_reading_dir = self.filepath.parent.with_child('offline_reading') self.index_dir = self.filepath.parent.with_child('index') self.styles_dir = self.filepath.parent.with_child('styles') self.wiki_dir = self.filepath.parent.with_child('wiki') existing_database = self.filepath.exists self.sql = sqlite3.connect(self.filepath.absolute_path) self.cur = self.sql.cursor() if existing_database: if not skip_version_check: self._check_version() self._load_pragmas() else: self._first_time_setup() self.config = {} for (key, default_value) in DEFAULT_CONFIG.items(): self.cur.execute('SELECT value FROM config WHERE key == ?', [key]) existing_value = self.cur.fetchone() if existing_value is None: self.cur.execute('INSERT INTO config VALUES(?, ?)', [key, default_value]) self.config[key] = default_value else: existing_value = existing_value[0] if isinstance(default_value, int): existing_value = int(existing_value) self.config[key] = existing_value def _check_version(self): ''' Compare database's user_version against DATABASE_VERSION, raising exceptions.DatabaseOutOfDate if not correct. ''' existing = self.cur.execute('PRAGMA user_version').fetchone()[0] if existing != DATABASE_VERSION: raise exceptions.DatabaseOutOfDate( current=existing, new=DATABASE_VERSION, filepath=self.filepath, ) def _first_time_setup(self): self.sql.executescript(DB_INIT) self.sql.commit() def _load_pragmas(self): self.sql.executescript(DB_PRAGMAS) self.sql.commit() def __repr__(self): return 'TSDB(%s)' % self.filepath @staticmethod def _pick_filepath(formats, name): ''' Starting with the most specific and preferred filename format, check if there is an existing database that matches the name we're looking for, and return that path. If none of them exist, then use the most preferred filepath. ''' for form in formats: path = form.format(name=name) if os.path.isfile(path): break return pathclass.Path(path) @classmethod def _for_object_helper(cls, name, path_formats, do_create=True, fix_name=False): if name != os.path.basename(name): filepath = pathclass.Path(name) else: filepath = cls._pick_filepath(formats=path_formats, name=name) database = cls(filepath=filepath, do_create=do_create) if fix_name: return (database, name_from_path(name)) return database @classmethod def for_submission(cls, submission_id, fix_name=False, *args, **kwargs): subreddit = common.subreddit_for_submission(submission_id) database = cls.for_subreddit(subreddit, *args, **kwargs) if fix_name: return (database, subreddit.display_name) return database @classmethod def for_subreddit(cls, name, do_create=True, fix_name=False): if isinstance(name, common.praw.models.Subreddit): name = name.display_name elif not isinstance(name, str): raise TypeError(name, 'should be str or Subreddit.') return cls._for_object_helper( name, do_create=do_create, fix_name=fix_name, path_formats=DB_FORMATS_SUBREDDIT, ) @classmethod def for_user(cls, name, do_create=True, fix_name=False): if isinstance(name, common.praw.models.Redditor): name = name.name elif not isinstance(name, str): raise TypeError(name, 'should be str or Redditor.') return cls._for_object_helper( name, do_create=do_create, fix_name=fix_name, path_formats=DB_FORMATS_USER, ) def check_for_edits(self, obj, existing_entry): ''' If the item's current text doesn't match the stored text, decide what to do. Firstly, make sure to ignore deleted comments. Then, if the database is configured to store edited text, do so. Finally, return the body that we want to store in the main table. ''' if isinstance(obj, SUBMISSION_TYPES): existing_body = existing_entry[SQL_SUBMISSION['selftext']] body = obj.selftext else: existing_body = existing_entry[SQL_COMMENT['body']] body = obj.body if body != existing_body: if should_keep_existing_text(obj): body = existing_body elif self.config['store_edits']: self.insert_edited(obj, old_text=existing_body) return body def insert(self, objects, commit=True): if not isinstance(objects, (list, tuple, types.GeneratorType)): objects = [objects] log.debug('Trying to insert %d objects.', len(objects)) new_values = { 'tsdb': self, 'new_submissions': 0, 'new_comments': 0, } methods = { common.praw.models.Submission: (self.insert_submission, 'new_submissions'), common.praw.models.Comment: (self.insert_comment, 'new_comments'), } methods[pushshift.DummySubmission] = methods[common.praw.models.Submission] methods[pushshift.DummyComment] = methods[common.praw.models.Comment] for obj in objects: (method, key) = methods.get(type(obj), (None, None)) if method is None: raise TypeError('Unsupported', type(obj), obj) status = method(obj) new_values[key] += status if commit: log.debug('Committing insert.') self.sql.commit() log.debug('Done inserting.') return new_values def insert_edited(self, obj, old_text): ''' Having already detected that the item has been edited, add a record to the appropriate *_edits table containing the text that is being replaced. ''' if isinstance(obj, SUBMISSION_TYPES): table = 'submission_edits' else: table = 'comment_edits' if obj.edited is False: replaced_at = int(time.time()) else: replaced_at = int(obj.edited) postdata = { 'idstr': obj.fullname, 'text': old_text, 'replaced_at': replaced_at, } cur = self.sql.cursor() (qmarks, bindings) = sqlhelpers.insert_filler(SQL_EDITS_COLUMNS, postdata) query = 'INSERT INTO %s VALUES(%s)' % (table, qmarks) cur.execute(query, bindings) def insert_submission(self, submission): cur = self.sql.cursor() cur.execute('SELECT * FROM submissions WHERE idstr == ?', [submission.fullname]) existing_entry = cur.fetchone() if submission.author is None: author = '[DELETED]' else: author = submission.author.name if not existing_entry: if submission.is_self: # Selfpost's URL leads back to itself, so just ignore it. url = None else: url = submission.url postdata = { 'idint': common.b36(submission.id), 'idstr': submission.fullname, 'created': submission.created_utc, 'self': submission.is_self, 'nsfw': submission.over_18, 'author': author, 'title': submission.title, 'url': url, 'selftext': submission.selftext, 'score': submission.score, 'subreddit': submission.subreddit.display_name, 'distinguish': submission.distinguished, 'textlen': len(submission.selftext), 'num_comments': submission.num_comments, 'flair_text': submission.link_flair_text, 'flair_css_class': submission.link_flair_css_class, 'augmented_at': None, 'augmented_count': None, } (qmarks, bindings) = sqlhelpers.insert_filler(SQL_SUBMISSION_COLUMNS, postdata) query = 'INSERT INTO submissions VALUES(%s)' % qmarks cur.execute(query, bindings) else: selftext = self.check_for_edits(submission, existing_entry=existing_entry) query = ''' UPDATE submissions SET nsfw = coalesce(?, nsfw), score = coalesce(?, score), selftext = coalesce(?, selftext), distinguish = coalesce(?, distinguish), num_comments = coalesce(?, num_comments), flair_text = coalesce(?, flair_text), flair_css_class = coalesce(?, flair_css_class) WHERE idstr == ? ''' bindings = [ submission.over_18, submission.score, selftext, submission.distinguished, submission.num_comments, submission.link_flair_text, submission.link_flair_css_class, submission.fullname ] cur.execute(query, bindings) return existing_entry is None def insert_comment(self, comment): cur = self.sql.cursor() cur.execute('SELECT * FROM comments WHERE idstr == ?', [comment.fullname]) existing_entry = cur.fetchone() if comment.author is None: author = '[DELETED]' else: author = comment.author.name if not existing_entry: postdata = { 'idint': common.b36(comment.id), 'idstr': comment.fullname, 'created': comment.created_utc, 'author': author, 'parent': comment.parent_id, 'submission': comment.link_id, 'body': comment.body, 'score': comment.score, 'subreddit': comment.subreddit.display_name, 'distinguish': comment.distinguished, 'textlen': len(comment.body), } (qmarks, bindings) = sqlhelpers.insert_filler(SQL_COMMENT_COLUMNS, postdata) query = 'INSERT INTO comments VALUES(%s)' % qmarks cur.execute(query, bindings) else: body = self.check_for_edits(comment, existing_entry=existing_entry) query = ''' UPDATE comments SET score = coalesce(?, score), body = coalesce(?, body), distinguish = coalesce(?, distinguish) WHERE idstr == ? ''' bindings = [ comment.score, body, comment.distinguished, comment.fullname ] cur.execute(query, bindings) return existing_entry is None def name_from_path(filepath): ''' In order to support usage like > timesearch livestream -r D:\\some\\other\\filepath\\learnpython.db this function extracts the subreddit name / username based on the given path, so that we can pass it into `r.subreddit` / `r.redditor` properly. ''' if isinstance(filepath, pathclass.Path): filepath = filepath.basename else: filepath = os.path.basename(filepath) name = os.path.splitext(filepath)[0] name = name.strip('@') return name def should_keep_existing_text(obj): ''' Under certain conditions we do not want to update the entry in the db with the most recent copy of the text. For example, if the post has been deleted and the text now shows '[deleted]' we would prefer to keep whatever we already have. This function puts away the work I would otherwise have to duplicate for both submissions and comments. ''' body = obj.selftext if isinstance(obj, SUBMISSION_TYPES) else obj.body if obj.author is None and body in ['[removed]', '[deleted]']: return True greasy = ['has been overwritten', 'pastebin.com/64GuVi2F'] if any(grease in body for grease in greasy): return True return False