import logging
import os

import peewee as pw


logger = logging.getLogger(__name__)


def migrate(migrator, database, fake=False, **kwargs):
    MalwareHit = migrator.orm["malware_hits"]
    migrator.add_fields(
        MalwareHit,
        uid=pw.IntegerField(null=True),
        gid=pw.IntegerField(null=True),
    )


def rollback(migrator, database, fake=False, **kwargs):
    MalwareHit = migrator.orm["malware_hits"]
    migrator.remove_fields(MalwareHit, "uid", "gid")