New unwrap methods

Author: Sanne

hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySessionImpl.java

@@ -11,6 +11,7 @@
 import jakarta.persistence.EntityGraph;
 import jakarta.persistence.FlushModeType;
 import jakarta.persistence.LockModeType;
+import jakarta.persistence.PersistenceException;
 import jakarta.persistence.criteria.CriteriaDelete;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.CriteriaUpdate;
@@ -24,6 +25,7 @@
 import org.hibernate.reactive.common.AffectedEntities;
 import org.hibernate.reactive.common.Identifier;
 import org.hibernate.reactive.common.ResultSetMapping;
+import org.hibernate.reactive.engine.spi.ReactiveSharedSessionContractImplementor;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.mutiny.Mutiny;
@@ -32,6 +34,8 @@
 import org.hibernate.reactive.mutiny.Mutiny.SelectionQuery;
 import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.query.sqm.iternal.ReactiveQuerySqmImpl;
+import org.hibernate.reactive.session.ReactiveConnectionSupplier;
+import org.hibernate.reactive.session.ReactiveQueryProducer;
 import org.hibernate.reactive.session.ReactiveSession;
 
 import java.lang.invoke.MethodHandles;
@@ -538,4 +542,21 @@ public <T> EntityGraph<T> createEntityGraph(Class<T> rootType) {
 	public <T> EntityGraph<T> createEntityGraph(Class<T> rootType, String graphName) {
 		return delegate.createEntityGraph( rootType, graphName );
 	}
+
+	public <T> T unwrap(Class<T> clazz) {
+		if ( ReactiveConnectionSupplier.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveSession.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveQueryProducer.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveSharedSessionContractImplementor.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		throw new PersistenceException( "Cannot unwrap type " + clazz );
+	}
+
 }

hibernate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageSessionImpl.java

@@ -10,6 +10,7 @@
 import jakarta.persistence.EntityGraph;
 import jakarta.persistence.FlushModeType;
 import jakarta.persistence.LockModeType;
+import jakarta.persistence.PersistenceException;
 import jakarta.persistence.criteria.CriteriaDelete;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.CriteriaUpdate;
@@ -25,9 +26,12 @@
 import org.hibernate.reactive.common.Identifier;
 import org.hibernate.reactive.common.ResultSetMapping;
 import org.hibernate.reactive.engine.ReactiveActionQueue;
+import org.hibernate.reactive.engine.spi.ReactiveSharedSessionContractImplementor;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.pool.ReactiveConnection;
+import org.hibernate.reactive.session.ReactiveConnectionSupplier;
+import org.hibernate.reactive.session.ReactiveQueryProducer;
 import org.hibernate.reactive.session.ReactiveSession;
 import org.hibernate.reactive.stage.Stage;
 import org.hibernate.reactive.stage.Stage.MutationQuery;
@@ -540,4 +544,20 @@ public <R> Query<R> createNativeQuery(String queryString) {
 	public <R> Query<R> createNativeQuery(String queryString, AffectedEntities affectedEntities) {
 		return new StageQueryImpl<>( delegate.createReactiveNativeQuery( queryString, affectedEntities ) );
 	}
+
+	public <T> T unwrap(Class<T> clazz) {
+		if ( ReactiveConnectionSupplier.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveSession.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveQueryProducer.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		else if ( ReactiveSharedSessionContractImplementor.class.isAssignableFrom( clazz ) ) {
+			return clazz.cast( this.delegate );
+		}
+		throw new PersistenceException( "Cannot unwrap type " + clazz );
+	}
 }