How to return the newly persisted entity using the reactive extensions for hibernate

I am trying to use mutiny in order to persist an entity. The add method should return a Uni<Entity> referencing the newly persisted (or merged) entity (I am using the isPersistent flag to determine whether an entity was already persisted previously). I also need a reference to the entity’s updated id if it was generated by hibernate.

    @Override
    public Uni<Entity> add(Entity entity) {
        if(entity.isPersistent()){
            return sessionFactory.withSession(s ->s.merge(entity));
        }else{
            entity.markAsPersistent();
            return sessionFactory.withSession(s ->s.persist(entity)); // Error!
        }
    }

However, s.persist() returns a Uni<Void>.
I tried to modify the code as follows (which results in a detached entity):

return sessionFactory.withSession(s ->s.persist(entity).chain(s::flush).replaceWith(entity));

How should I proceed to map the Uni<Void> to a corresponding Uni<Entity>, which is not in a detached state?

>Solution :

You need to "replace" the return value of the second block:

@Override
public Uni<Entity> add(Entity entity) {
    if (entity.isPersistent()) {
        return sessionFactory.withSession(s -> s.merge(entity));
    } else {
        entity.markAsPersistent();
        return sessionFactory.withSession(s -> s.persist(entity))
                             .replaceWith(() -> entity);
    }
}