restore.py 1.55 KB
import csv


class Restore:
    def __init__(self, db_session, table, filename):
        self.db_session = db_session
        self.table = table
        self.filename = filename

    def open_file(self):
        return open(self.filename)

    def run(self):
        q = self.db_session.query(self.table)
        if q.first():
            return
        with self.open_file() as f:
            reader = csv.DictReader(f)
            for cf in reader:
                row = self.table()
                for fieldname in cf:
                    val = cf[fieldname]
                    if not val:
                        continue
                    setattr(row, fieldname, val)
                self.db_session.add(row)
        return True


class Append(Restore):
    def __init__(self, db_session, table, filename, keys):
        Restore.__init__(self, db_session, table, filename)
        self.keys = keys
        self.run()

    def run(self):
        with self.open_file() as f:
            reader = csv.DictReader(f)
            filter_ = dict()
            for cf in reader:
                for key in self.keys:
                    filter_[key] = cf[key]
                q = self.db_session.query(self.table).filter_by(**filter_)
                found = q.first()
                if found:
                    continue
                row = self.table()
                for fieldname in cf:
                    val = cf[fieldname]
                    if not val:
                        continue
                    setattr(row, fieldname, val)
                self.db_session.add(row)