import logging import peewee as pw import ipaddress logger = logging.getLogger(__name__) def migrate(migrator, database, fake=False, **kwargs): """Write your migrations here.""" IPListNew = migrator.orm["iplist_new"] IPList = migrator.orm["iplist"] def iplist_select(): # FIXME: remove this function after migrating to peewee 3 try: yield from IPList.select(IPList).dicts().iterator() except RuntimeError: return try: from im360.utils.net import pack_ip_network except ImportError: pass else: with database.atomic(): for ip_obj in iplist_select(): try: ip = ipaddress.ip_network(ip_obj["ip"]) except ValueError: # malformed ip continue net, mask, version = pack_ip_network(ip) ip_obj.update( { "network_address": net, "netmask": mask, "version": version, } ) try: IPListNew.insert(ip_obj).execute() except pw.IntegrityError as e: logger.warning("Error inserting IP: %s", e) migrator.sql("DROP TABLE iplist") migrator.sql("ALTER TABLE iplist_new RENAME TO iplist") migrator.sql('CREATE INDEX "iplist_listname" ON "iplist" ("listname")') migrator.sql('CREATE INDEX "iplist_expiration" ON "iplist" ("expiration")') migrator.sql('CREATE INDEX "iplist_ip" ON "iplist" ("ip")') def rollback(migrator, database, fake=False, **kwargs): """Write your rollback migrations here.""" pass