import logging

import peewee as pw

logger = logging.getLogger(__name__)


def migrate(migrator, database, fake=False, **kwargs):
    """Write your migrations here."""
    MalwareHits = migrator.orm["malware_hits"]
    migrator.add_fields(
        MalwareHits, size=pw.CharField(null=True), hash=pw.CharField(null=True)
    )


def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    MalwareHits = migrator.orm["malware_hits"]
    migrator.remove_fields(MalwareHits, "hash", "size")
