HHH-19849 Add an SPI that allows attaching session-scoped "extensions" to the session/statelesssession implementors

Author: marko-bekhta

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

@@ -517,6 +517,16 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
 		return delegate.getEntityGraph( graphName );
 	}
 
+	@Override
+	public void attachExtension(String extensionName, Object extension) {
+		delegate.attachExtension( extensionName, extension );
+	}
+
+	@Override
+	public <T> T retrieveExtension(String extensionName, Class<T> extensionType) {
+		return delegate.retrieveExtension( extensionName, extensionType );
+	}
+
 	@Override
 	public <T> QueryImplementor<T> createQuery(CriteriaSelect<T> selectQuery) {
 		return delegate.createQuery( selectQuery );

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

@@ -621,4 +621,26 @@ default boolean isStatelessSession() {
 
 	@Override
 	RootGraphImplementor<?> getEntityGraph(String graphName);
+
+	/**
+	 * Allows attaching session scoped extensions to the particular session instance they are based on.
+	 *
+	 * @param extensionName The name of the extension serves as a "key" for its retrival from a session instance.
+	 * @param extension The extension to attach to the current session.
+	 */
+	@Incubating
+	void attachExtension(String extensionName, Object extension);
+
+	/**
+	 * Returns the extensions attached to the current session.
+	 *
+	 * @param extensionName The name of the extension to retrieve.
+	 * @param extensionType The type of the extension to retrieve.
+	 * @param <T> The type of the extension to retrieve.
+	 * @return The extension instance attached to the current session,
+	 * or {@code null} if there is no extension with the requested name attached to the session.
+	 * @throws ClassCastException if the requested extension cannot be cast to the requested type.
+	 */
+	@Incubating
+	<T> T retrieveExtension(String extensionName, Class<T> extensionType);
 }

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

@@ -663,6 +663,16 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
 		return delegate.getEntityGraph( graphName );
 	}
 
+	@Override
+	public void attachExtension(String extensionName, Object extension) {
+		delegate.attachExtension( extensionName, extension );
+	}
+
+	@Override
+	public <T> T retrieveExtension(String extensionName, Class<T> extensionType) {
+		return delegate.retrieveExtension( extensionName, extensionType );
+	}
+
 	@Override
 	public <T> List<EntityGraph<? super T>> getEntityGraphs(Class<T> entityClass) {
 		return delegate.getEntityGraphs( entityClass );

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

@@ -111,8 +111,10 @@
 import java.io.Serial;
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
 import java.util.TimeZone;
 import java.util.UUID;
@@ -186,6 +188,8 @@ public abstract class AbstractSharedSessionContract implements SharedSessionCont
 	private transient ExceptionConverter exceptionConverter;
 	private transient SessionAssociationMarkers sessionAssociationMarkers;
 
+	private transient Map<String, Object> extensions;
+
 	public AbstractSharedSessionContract(SessionFactoryImpl factory, SessionCreationOptions options) {
 		this.factory = factory;
 
@@ -1704,6 +1708,23 @@ public SessionAssociationMarkers getSessionAssociationMarkers() {
 		return sessionAssociationMarkers;
 	}
 
+	@Override
+	public <T> T retrieveExtension(String extensionName, Class<T> extensionType) {
+		if ( extensions != null ) {
+			Object extension = extensions.get( extensionName );
+			return extension == null ? null : extensionType.cast( extension );
+		}
+		return null;
+	}
+
+	@Override
+	public void attachExtension(String extensionName, Object extension) {
+		if ( extensions == null ) {
+			extensions = new HashMap<>();
+		}
+		extensions.put( extensionName, extension );
+	}
+
 	@Serial
 	private void writeObject(ObjectOutputStream oos) throws IOException {
 		SESSION_LOGGER.serializingSession( getSessionIdentifier() );

hibernate-core/src/test/java/org/hibernate/orm/test/engine/spi/SessionExtensionTest.java

@@ -0,0 +1,69 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.engine.spi;
+
+import jakarta.persistence.Id;
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+@DomainModel(annotatedClasses = {
+		SessionExtensionTest.UselessEntity.class,
+})
+@SessionFactory
+public class SessionExtensionTest {
+
+	@Test
+	public void smoke(SessionFactoryScope scope) {
+		final String extensionName = "my-extension-key";
+		scope.inSession( sessionImplementor -> {
+			sessionImplementor.attachExtension( extensionName, new Extension( 1 ) );
+
+			assertThat( sessionImplementor.retrieveExtension( extensionName, Extension.class ) )
+					.isNotNull()
+					.isEqualTo( new Extension( 1 ) );
+		} );
+
+		scope.inStatelessSession( sessionImplementor -> {
+			sessionImplementor.attachExtension( extensionName, new Extension( 1 ) );
+
+			assertThat( sessionImplementor.retrieveExtension( extensionName, Extension.class ) )
+					.isNotNull()
+					.isEqualTo( new Extension( 1 ) );
+		} );
+	}
+
+	@Test
+	public void cast(SessionFactoryScope scope) {
+		final String extensionName = "my-extension-key";
+		scope.inSession( sessionImplementor -> {
+			sessionImplementor.attachExtension( extensionName, new Extension( 1 ) );
+
+			assertThatThrownBy(
+					() -> sessionImplementor.retrieveExtension( extensionName, SessionExtensionTest.class ) )
+					.isInstanceOf( ClassCastException.class );
+		} );
+
+		scope.inStatelessSession( sessionImplementor -> {
+			sessionImplementor.attachExtension( extensionName, new Extension( 1 ) );
+
+			assertThatThrownBy(
+					() -> sessionImplementor.retrieveExtension( extensionName, SessionExtensionTest.class ) )
+					.isInstanceOf( ClassCastException.class );
+		} );
+	}
+
+	private record Extension(int number) {
+	}
+
+	static class UselessEntity {
+		@Id
+		Long id;
+	}
+}