Update EntityId handling code

Change-Id: Ie7863e1cb1790aa3038978603a40dcbaed803345