Fix missing import statements for schema code

Closes https://github.com/simolus3/drift/issues/3582