class DatabaseHandling: def __init__(self, **kwargs): for key, value in kwargs.items(): setattr(self, key, value) self.db_name = config.FEEDGASM_DATABASE self.session = Database(db=self.db_name) def create_table(self): log.info("%s: Creating table %s.", self.db_name, self.table) columns = 'id INTEGER PRIMARY KEY' for col, datatype in self.sql_columns.items(): columns += f', {col} {datatype}' sql = f'CREATE TABLE if not exists "{self.table}" ({columns});' '' self.session.alter(sql) if self.row_limit: self.create_row_trigger() def create_row_trigger(self): log.debug('%s: Row limit set to %d, creating trigger.', self.db_name, self.row_limit) sql = f"""CREATE TRIGGER row_limit AFTER INSERT ON '{self.table}' BEGIN DELETE FROM '{self.table}' WHERE id <= ( SELECT id FROM '{self.table}' ORDER BY id DESC LIMIT {str(self.row_limit)}, 1); END; """ self.session.alter(sql) def insert(self): code = "INSERT INTO '{}' ({}) VALUES ({})" cols = [] for c in self.sql_columns: cols.append(c) columns = ', '.join(cols) question_marks = ','.join('?' * len(cols)) pubs = self.publications.copy() for n, cols in pubs.items(): values = [] for data in cols.values(): values.append(data) code = code.format(self.table, columns, question_marks) ret = self.session.alter(code, values) if ret: if 'SQLite error' in ret: self.publications.pop(n) log.warning( '%s: A publication was ignored because of SQLite insertion error.', self.table) def iterate_columns(self): sql_select = "SELECT {0} FROM '{1}' WHERE {0} = '{2}'" pubs = self.publications.copy() for key, columns in pubs.items(): sql = sql_select.format(self.comparison_item, self.table, columns[self.comparison_item]) ret = self.session.retrieve(sql) if ret: if 'no such table' in ret: self.create_table() else: # Item already in database. self.publications.pop(key) # Insert new items, if any. if self.publications: self.insert() self.session.close() return self.publications
class DatabaseHandling(Main): def __init__(self, client=None, db=None, regex_function=False): if client: self.client = client self.db = Database(db=db) self.columns = { 'author': 'TEXT', 'author.id': 'TEXT', 'author.name': 'TEXT', 'author.nick': 'TEXT', 'clean_content': 'TEXT', 'channel.name': 'TEXT', 'message_id': 'TEXT', 'attachment_urls': 'TEXT', 'pinned': 'TEXT', 'reactions': 'TEXT', 'raw_mentions': 'TEXT', 'raw_channel_mentions': 'TEXT', 'raw_role_mentions': 'TEXT', 'created_at': 'TEXT', 'edited_at': 'TEXT', 'jump_url': 'TEXT', } self.db_name = db if regex_function: self.db.conn.create_function('REGEXP', 2, self.regexp) def regexp(self, expr, item): reg = re.compile(expr.lower()) return reg.search(item) is not None def drop_table(self, table): log.info("%s: Dropping table (if exists) %s.", self.db_name, table) sql = f'DROP TABLE if exists "{table}"' self.db.alter(sql) def create_table(self, table): log.info("%s: Attempting to create table %s.", self.db_name, table) columns = 'id INTEGER PRIMARY KEY' for col, datatype in self.columns.items(): columns += f', {col} {datatype}' sql = f'CREATE TABLE if not exists "{table}" ({columns})'.replace( '.', '_') self.db.alter(sql) def insert_message(self, message=None, table=None): code = "INSERT INTO '{}' ({}) VALUES ({})" columns = ', '.join([c for c in self.columns]) question_marks = ','.join('?' * len(self.columns)) values = [] for col, datatype in self.columns.items(): if col == 'message_id': val = message.id elif col == 'attachment_urls': attachments = message.attachments val = [] if attachments: for attachment in attachments: val.append(attachment.url) val = ','.join(val) else: val = None elif '.' in col: c = col.split('.') val = getattr(message, c[0]) try: val = str(getattr(val, c[1])) except: val = None else: if col == 'created_at' or col == 'edited_at' and getattr( message, col): # Convert datetime try: strptime = datetime.strptime( str(getattr(message, col)), '%Y-%m-%d %H:%M:%S.%f') except ValueError: strptime = datetime.strptime( str(getattr(message, col)), '%Y-%m-%d %H:%M:%S') val = Main.utc_to_local('this.selfie', strptime) else: val = str(getattr(message, col)) values.append(val) code = code.format(table, columns, question_marks).replace('.', '_') ret = self.db.alter(code, values) if ret: if 'no such table' in ret: self.create_table(table) ret = self.db.alter(code, values) elif 'has no column named' in ret: while ret: column = re.search('column named (.*)', ret)[1] col_type = self.columns[column] sql = f"ALTER TABLE '{table}' ADD COLUMN {column} {col_type}" log.info( 'Altering the table by adding the missing column.') ret = self.db.alter(sql) def retrieve(self, sql, params=None): return (self.db.retrieve(sql, params)) def close(self): self.db.close()