Code cleanup on migrations