diff --git a/IzzyLib/database.py b/IzzyLib/database.py index 978deba..e9d4344 100644 --- a/IzzyLib/database.py +++ b/IzzyLib/database.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import sessionmaker from . import logging from .cache import LRUCache -from .misc import DotDict, RandomGen, NfsCheck +from .misc import DotDict, RandomGen, NfsCheck, PrintMethods SqlTypes = DotDict({t.lower(): getattr(Types, t) for t in dir(Types) if not t.startswith('_')}) @@ -225,6 +225,53 @@ class Session(object): return [row[0] for row in rows] + def AppendColumn(self, tbl, col): + table = self.table[tbl] + + try: + column = getattr(table.c, col) + + except AttributeError: + logging.error(f'Table "{tbl}" does not have column "{col}"') + return + + columns = [row[1] for row in self.execute(f'PRAGMA table_info({tbl})')] + + if col in columns: + logging.info(f'Column "{col}" already exists') + return + + sql = f'ALTER TABLE {tbl} ADD COLUMN {col} {column.type}' + + if not column.nullable: + sql += ' NOT NULL' + + if column.primary_key: + sql += ' PRIMARY KEY' + + if column.unique: + sql += ' UNIQUE' + + self.execute(sql) + + + def RemoveColumn(self, tbl, col): + table = self.table[tbl] + column = getattr(table, col, None) + columns = [row[1] for row in self.execute(f'PRAGMA table_info({tbl})')] + + if col not in columns: + logging.info(f'Column "{col}" already exists') + return + + columns.remove(col) + coltext = ', '.join(columns) + + self.execute(f'CREATE TABLE {tbl}_temp AS SELECT {coltext} FROM {tbl}') + self.execute(f'DROP TABLE {tbl}') + self.execute(f'ALTER TABLE {tbl}_temp RENAME TO {tbl}') + + class CustomRows(object): def get(self, name): return getattr(self, name, self.Row)