Added entity class. Doesn't work right now.