Refactor, added schema