HHH-15858 Improvements in Session casting

hibernate-core/src/main/java/org/hibernate/engine/internal/AbstractEntityEntry.java

@@ -276,7 +276,7 @@ public abstract class AbstractEntityEntry implements Serializable, EntityEntry {
 		session
 				.getFactory()
 				.getCustomEntityDirtinessStrategy()
-				.resetDirty( entity, persister, (Session) session );
+				.resetDirty( entity, persister, session.asSessionImplementor() );
 	}
 
 	private static void clearDirtyAttributes(final SelfDirtinessTracker entity) {
@@ -370,8 +370,9 @@ public abstract class AbstractEntityEntry implements Serializable, EntityEntry {
 
 		final CustomEntityDirtinessStrategy customEntityDirtinessStrategy =
 				getPersistenceContext().getSession().getFactory().getCustomEntityDirtinessStrategy();
-		if ( customEntityDirtinessStrategy.canDirtyCheck( entity, getPersister(), (Session) getPersistenceContext().getSession() ) ) {
+		final Session session = getPersistenceContext().getSession().asSessionImplementor();
-			return ! customEntityDirtinessStrategy.isDirty( entity, getPersister(), (Session) getPersistenceContext().getSession() );
+		if ( customEntityDirtinessStrategy.canDirtyCheck( entity, getPersister(), session  ) ) {
+			return ! customEntityDirtinessStrategy.isDirty( entity, getPersister(), session );
 		}
 
 		if ( getPersister().hasMutableProperties() ) {

hibernate-core/src/main/java/org/hibernate/engine/spi/SessionImplementor.java

@@ -112,4 +112,15 @@ public interface SessionImplementor extends Session, SharedSessionContractImplem
 	 */
 	@Deprecated
 	void removeOrphanBeforeUpdates(String entityName, Object child);
+
+	@Override
+	default SessionImplementor asSessionImplementor() {
+		return this;
+	}
+
+	@Override
+	default boolean isSessionImplementor() {
+		return true;
+	}
+
 }

hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionContractImplementor.java

@@ -15,6 +15,7 @@ import org.hibernate.CacheMode;
 import org.hibernate.FlushMode;
 import org.hibernate.HibernateException;
 import org.hibernate.Interceptor;
+import org.hibernate.StatelessSession;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.query.Query;
 import org.hibernate.SharedSessionContract;
@@ -440,4 +441,30 @@ public interface SharedSessionContractImplementor
 	 */
 	void afterOperation(boolean success);
 
+	/**
+	 * If this can be casted to a @{@link SessionImplementor},
+	 * you'll get this returned after an efficient cast.
+	 * @throws ClassCastException if this is not compatible!
+	 */
+	default SessionImplementor asSessionImplementor() {
+		throw new ClassCastException();
+	}
+
+	default boolean isSessionImplementor() {
+		return false;
+	}
+
+	/**
+	 * If this can be casted to a @{@link StatelessSession},
+	 * you'll get this returned after an efficient cast.
+	 * @throws ClassCastException if this is not compatible!
+	 */
+	default StatelessSession asStatelessSession() {
+		throw new ClassCastException();
+	}
+
+	default boolean isStatelessSession() {
+		return false;
+	}
+
 }

hibernate-core/src/main/java/org/hibernate/generator/internal/CurrentTimestampGeneration.java

@@ -92,7 +92,7 @@ public class CurrentTimestampGeneration implements InMemoryGenerator, InDatabase
 
 	@Override
 	public Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {
-		return generator.generateValue( (Session) session, owner, currentValue );
+		return generator.generateValue( session.asSessionImplementor(), owner, currentValue );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/generator/internal/SourceGeneration.java

@@ -90,7 +90,7 @@ public class SourceGeneration implements InMemoryGenerator {
 
 	@Override
 	public Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {
-		return valueGenerator.generateValue( (Session) session, owner, currentValue );
+		return valueGenerator.generateValue( session.asSessionImplementor(), owner, currentValue );
 	}
 
 	public Object generateValue(Session session, Object owner) {

hibernate-core/src/main/java/org/hibernate/id/ForeignGenerator.java

@@ -10,13 +10,10 @@ import java.util.Properties;
 
 import org.hibernate.MappingException;
 import org.hibernate.Session;
-import org.hibernate.StatelessSession;
 import org.hibernate.TransientObjectException;
-import org.hibernate.engine.internal.ForeignKeys;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.id.factory.spi.StandardGenerator;
 import org.hibernate.internal.CoreMessageLogger;
-import org.hibernate.loader.PropertyPath;
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.type.EntityType;
@@ -118,18 +115,18 @@ public class ForeignGenerator implements IdentifierGenerator, StandardGenerator
 						associatedEntityName
 				);
 			}
-			if ( sessionImplementor instanceof Session ) {
+			if ( sessionImplementor.isSessionImplementor() ) {
-				id = ((Session) sessionImplementor).save( associatedEntityName, associatedObject );
+				id = sessionImplementor.asSessionImplementor().save( associatedEntityName, associatedObject );
 			}
-			else if ( sessionImplementor instanceof StatelessSession ) {
+			else if ( sessionImplementor.isStatelessSession() ) {
-				id = ((StatelessSession) sessionImplementor).insert( associatedEntityName, associatedObject );
+				id = sessionImplementor.asStatelessSession().insert( associatedEntityName, associatedObject );
 			}
 			else {
 				throw new IdentifierGenerationException("sessionImplementor is neither Session nor StatelessSession");
 			}
 		}
 
-		if ( sessionImplementor instanceof Session && ((Session) sessionImplementor).contains( entityName, object ) ) {
+		if ( sessionImplementor.isSessionImplementor() && sessionImplementor.asSessionImplementor().contains( entityName, object ) ) {
 			//abort the save (the object is already saved by a circular cascade)
 			return SHORT_CIRCUIT_INDICATOR;
 			//throw new IdentifierGenerationException("save associated object first, or disable cascade for inverse association");

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

@@ -725,4 +725,15 @@ public class StatelessSessionImpl extends AbstractSharedSessionContract implemen
 	private LockMode getNullSafeLockMode(LockMode lockMode) {
 		return lockMode == null ? LockMode.NONE : lockMode;
 	}
+
+	@Override
+	public StatelessSession asStatelessSession() {
+		return this;
+	}
+
+	@Override
+	public boolean isStatelessSession() {
+		return true;
+	}
+
 }

hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/DeferredResultSetAccess.java

@@ -119,7 +119,7 @@ public class DeferredResultSetAccess extends AbstractResultSetAccess {
 
 						executionContext.getCallback().registerAfterLoadAction(
 								(session, entity, persister) ->
-									( (SessionImplementor) session ).lock(
+									session.asSessionImplementor().lock(
 											persister.getEntityName(),
 											entity,
 											lockOptionsToUse

hibernate-core/src/main/java/org/hibernate/tuple/ValueGeneration.java

@@ -59,7 +59,7 @@ public interface ValueGeneration extends InMemoryGenerator, InDatabaseGenerator
 
 	@Override
 	default Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {
-		return getValueGenerator().generateValue( (Session) session, owner, currentValue );
+		return getValueGenerator().generateValue( session.asSessionImplementor(), owner, currentValue );
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/tuple/VmValueGeneration.java

@@ -52,6 +52,6 @@ public class VmValueGeneration implements InMemoryGenerator {
 
 	@Override
 	public Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {
-		return generator.generateValue( (Session) session, owner, currentValue );
+		return generator.generateValue( session.asSessionImplementor(), owner, currentValue );
 	}
 }