"""Peewee migrations -- 068_MAIN.py."""
import peewee as pw
from peewee_migrate import Migrator

try:
    import playhouse.postgres_ext as pw_pext
except ImportError:
    pass

SQL = pw.SQL


def migrate(migrator: Migrator, database, fake=False, **kwargs):
    migrator.remove_model('pumpsurcharge')


def rollback(migrator: Migrator, database, fake=False, **kwargs):
    @migrator.create_model
    class PumpSurcharge(pw.Model):
        id = pw.AutoField()
        type = pw.IntegerField()
        pump = pw.ForeignKeyField(backref='surcharges', column_name='pump_id', field='id', model=migrator.orm['pump'], on_delete='CASCADE')
        price = pw.DoubleField(null=True)

        class Meta:
            table_name = "pumpsurcharge"
            indexes = [(('type', 'pump'), True)]

