import logging

from peewee import BooleanField

logger = logging.getLogger(__name__)


def migrate(migrator, database, fake=False, **kwargs):
    if not fake:
        MalwareHit = migrator.orm["malware_hits"]
        migrator.remove_fields(MalwareHit, "restored")


def rollback(migrator, database, fake=False, **kwargs):
    MalwareHit = migrator.orm["malware_hits"]
    migrator.add_fields(MalwareHit, restored=BooleanField(default=False))