Add missing imports in migration

The problem showed during tests, but the fix is there was a missing
import in a couple of migrations.