Revert "Revert "Update entity module""

This reverts commit 36606273f6963eef35a04b63114593922912fda4.