UnitOfWork refactoring