from peewee import CharField, Check, IntegerField, Model from defence360agent.contracts.config import IContactMessageType class IContactThrottle(Model): class Meta: db_table = "icontact_throttle" message_type = CharField( primary_key=True, constraints=[ Check( "message_type in {}".format( ( str(IContactMessageType.MALWARE_FOUND), str(IContactMessageType.SCAN_NOT_SCHEDULED), ) ) ) ], ) timestamp = IntegerField(default=0) def migrate(migrator, database, fake=False, **kwargs): migrator.create_model(IContactThrottle) def rollback(migrator, database, fake=False, **kwargs): IContactThrottle = migrator.orm["icontact_throttle"] migrator.remove_model(IContactThrottle)