database: add AppendColumn and RemoveColumn functions

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)