import javax.persistence.EntityNotFoundException;
import javax.ws.rs.core.Response;
import javax.ws.rs.ext.ExceptionMapper;
import javax.ws.rs.ext.Provider;

@Provider
public class EntityNotFoundEjbExceptionMapper implements ExceptionMapper<EntityNotFoundException> {

  @Override
  public Response toResponse(EntityNotFoundException exception) {
    Throwable cause = exception.getCause();
    if (cause instanceof EntityNotFoundException) {
      EntityNotFoundException actual = (EntityNotFoundException) cause;
      return Response.status(Response.Status.NOT_FOUND).
        header("cause", "Entity not found.").
        header("additional-info", actual.getMessage()).
        build();
    }

    return Response.serverError().
      header("cause", exception.toString()).build();
  }
}