database: add AppendColumn and RemoveColumn functions
This commit is contained in:
parent
21c1c3cb1e
commit
24984164f0
|
@ -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)
|
||||
|
|
Loading…
Reference in a new issue