Refactor for usage with new persistence container functionality