Improve unit of work

This PR closes #31.