[7.0] Backport HHH-19531 + HHH-19532 Add unwrap() method to StatelessSession/SharedSessionContract and use it instead of casts in Jakarta Data implementation

hibernate-core/src/main/java/org/hibernate/Session.java

@@ -1494,4 +1494,9 @@ public interface Session extends SharedSessionContract, EntityManager {
 	 */
 	@Override @Deprecated(since = "6.0") @SuppressWarnings("rawtypes")
 	Query createQuery(CriteriaUpdate updateQuery);
+
+	@Override
+	default <T> T unwrap(Class<T> type) {
+		return SharedSessionContract.super.unwrap(type);
+	}
 }

hibernate-core/src/main/java/org/hibernate/SharedSessionContract.java

@@ -10,6 +10,7 @@
 import java.util.function.Function;
 
 import jakarta.persistence.EntityGraph;
+import jakarta.persistence.PersistenceException;
 import org.hibernate.graph.RootGraph;
 import org.hibernate.jdbc.ReturningWork;
 import org.hibernate.jdbc.Work;
@@ -472,4 +473,28 @@ default <R> R fromTransaction(Function<? super Transaction,R> action) {
 		final Transaction transaction = beginTransaction();
 		return manageTransaction( transaction, transaction, action );
 	}
+
+	/**
+	 * Return an object of the specified type to allow access to
+	 * a provider-specific API.
+	 *
+	 * @param type the class of the object to be returned.
+	 * This is usually either the underlying class
+	 * implementing {@code SharedSessionContract} or an
+	 * interface it implements.
+	 * @return an instance of the specified class
+	 * @throws PersistenceException if the provider does not
+	 * support the given type
+	 */
+	default <T> T unwrap(Class<T> type) {
+		// Not checking type.isInstance(...) because some implementations
+		// might want to hide that they implement some types.
+		// Implementations wanting a more liberal behavior need to override this method.
+		if ( type.isAssignableFrom( SharedSessionContract.class ) ) {
+			return type.cast( this );
+		}
+
+		throw new PersistenceException(
+				"Hibernate cannot unwrap '" + getClass().getName() + "' as '" + type.getName() + "'" );
+	}
 }

hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java

@@ -2948,7 +2948,8 @@ public <T> T unwrap(Class<T> type) {
 			return type.cast( persistenceContext );
 		}
 
-		throw new PersistenceException( "Hibernate cannot unwrap EntityManager as '" + type.getName() + "'" );
+		throw new PersistenceException(
+				"Hibernate cannot unwrap '" + getClass().getName() + "' as '" + type.getName() + "'" );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/internal/StatelessSessionImpl.java

@@ -9,6 +9,7 @@
 import java.util.Set;
 import java.util.function.BiConsumer;
 
+import jakarta.persistence.PersistenceException;
 import org.hibernate.AssertionFailure;
 import org.hibernate.FlushMode;
 import org.hibernate.HibernateException;
@@ -1427,6 +1428,18 @@ public Object loadFromSecondLevelCache(EntityPersister persister, EntityKey enti
 		return CacheLoadHelper.loadFromSecondLevelCache( this, instanceToLoad, lockMode, persister, entityKey );
 	}
 
+	@Override
+	public <T> T unwrap(Class<T> type) {
+		checkOpen();
+
+		if ( type.isInstance( this ) ) {
+			return type.cast( this );
+		}
+
+		throw new PersistenceException(
+				"Hibernate cannot unwrap '" + getClass().getName() + "' as '" + type.getName() + "'" );
+	}
+
 	private static final class MultiLoadOptions implements MultiIdLoadOptions {
 		private final  LockOptions lockOptions;
 

hibernate-core/src/main/java/org/hibernate/query/specification/internal/MutationSpecificationImpl.java

@@ -122,7 +122,7 @@ public MutationQuery createQuery(StatelessSession session) {
 	}
 
 	public MutationQuery createQuery(SharedSessionContract session) {
-		final var sessionImpl = (SharedSessionContractImplementor) session;
+		final var sessionImpl = session.unwrap(SharedSessionContractImplementor.class);
 		final SqmDeleteOrUpdateStatement<T> sqmStatement = build( sessionImpl.getFactory().getQueryEngine() );
 		return new QuerySqmImpl<>( sqmStatement, true, null, sessionImpl );
 	}

hibernate-core/src/main/java/org/hibernate/query/specification/internal/SelectionSpecificationImpl.java

@@ -164,7 +164,7 @@ public SelectionQuery<T> createQuery(StatelessSession session) {
 	}
 
 	public SelectionQuery<T> createQuery(SharedSessionContract session) {
-		final var sessionImpl = (SharedSessionContractImplementor) session;
+		final var sessionImpl = session.unwrap(SharedSessionContractImplementor.class);
 		final SqmSelectStatement<T> sqmStatement = build( sessionImpl.getFactory().getQueryEngine() );
 		return new SqmSelectionQueryImpl<>( sqmStatement, true, resultType, sessionImpl );
 	}

hibernate-core/src/test/java/org/hibernate/orm/test/query/dynamic/SimpleQuerySpecificationTests.java

@@ -8,6 +8,8 @@
 import jakarta.persistence.criteria.CriteriaQuery;
 
 import org.hibernate.SessionFactory;
+import org.hibernate.engine.spi.SessionLazyDelegator;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.query.IllegalMutationQueryException;
 import org.hibernate.query.IllegalSelectQueryException;
 import org.hibernate.query.Order;
@@ -19,6 +21,7 @@
 import org.hibernate.query.restriction.Restriction;
 import org.hibernate.testing.jdbc.SQLStatementInspector;
 import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.JiraKey;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
 import org.junit.jupiter.api.Test;
 
@@ -160,6 +163,46 @@ void testSimpleMutationRestriction(SessionFactoryScope factoryScope) {
 		assertThat( sqlCollector.getSqlQueries().get( 0 ) ).contains( " where be1_0.position between ? and ?" );
 	}
 
+	@Test
+	@JiraKey("HHH-19531")
+	void testSelectionOnSessionProxy(SessionFactoryScope factoryScope) {
+		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();
+
+		factoryScope.inTransaction( (session) -> {
+			var sessionProxy = new SessionLazyDelegator( () -> session );
+			// The test only makes sense if this is true. It currently is, but who knows what the future has in store for us.
+			//noinspection ConstantValue
+			assert !(sessionProxy instanceof SharedSessionContractImplementor);
+
+			sqlCollector.clear();
+			SelectionSpecification.create( BasicEntity.class, "from BasicEntity" )
+					.createQuery( sessionProxy )
+					.list();
+		} );
+
+		assertThat( sqlCollector.getSqlQueries() ).hasSize( 1 );
+	}
+
+	@Test
+	@JiraKey("HHH-19531")
+	void testMutationOnSessionProxy(SessionFactoryScope factoryScope) {
+		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();
+
+		factoryScope.inTransaction( (session) -> {
+			var sessionProxy = new SessionLazyDelegator( () -> session );
+			// The test only makes sense if this is true. It currently is, but who knows what the future has in store for us.
+			//noinspection ConstantValue
+			assert !(sessionProxy instanceof SharedSessionContractImplementor);
+
+			sqlCollector.clear();
+			MutationSpecification.create( BasicEntity.class, "delete BasicEntity" )
+					.createQuery( sessionProxy )
+					.executeUpdate();
+		} );
+
+		assertThat( sqlCollector.getSqlQueries() ).hasSize( 1 );
+	}
+
 	@Test
 	void testSimpleMutationRestrictionAsReference(SessionFactoryScope factoryScope) {
 		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();