Put entity id in class instead