import peewee class MalwareScan(peewee.Model): class Meta: db_table = "malware_scans" scanid = peewee.CharField(primary_key=True) class MalwareHit(peewee.Model): class Meta: db_table = "malware_hits" id = peewee.PrimaryKeyField() scanid = peewee.ForeignKeyField( MalwareScan, null=False, related_name="hits", on_delete="CASCADE" ) user = peewee.CharField(null=False) orig_file = peewee.BlobField(null=False) type = peewee.CharField(null=False) malicious = peewee.BooleanField(null=False, default=False) vendor = peewee.CharField(null=False, default="ai-bolit") hash = peewee.CharField(null=True) size = peewee.CharField(null=True) timestamp = peewee.FloatField(null=True) status = peewee.CharField(default="found") cleaned_at = peewee.FloatField(null=True) @classmethod def get_field_names(cls): return map(lambda field: field.column_name, cls._meta.sorted_fields) def migrate(migrator, database, fake=False, **kwargs): migrator.sql("ALTER TABLE malware_hits RENAME TO malware_hits_old;") migrator.create_model(MalwareHit) # specify fields order directly, because '*' doesnt guarantee order malware_hit_fields = ",".join(MalwareHit.get_field_names()) migrator.sql( "INSERT INTO malware_hits ({0}) SELECT {0} FROM malware_hits_old;" .format(malware_hit_fields) ) migrator.sql("DROP TABLE malware_hits_old;") def rollback(migrator, database, fake=False, **kwargs): pass