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

marko-bekhta · marko-bekhta · commit 48f01221ee27 · 2025-11-11T21:19:39.000+01:00
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/internal/ExtensionIntegrationServiceImpl.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/internal/ExtensionIntegrationServiceImpl.java
@@ -0,0 +1,50 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.internal;
+
+import org.hibernate.boot.registry.classloading.spi.ClassLoaderService;
+import org.hibernate.engine.extension.spi.ExtensionIntegration;
+import org.hibernate.engine.extension.spi.ExtensionIntegrationService;
+import org.jboss.logging.Logger;
+
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+public class ExtensionIntegrationServiceImpl implements ExtensionIntegrationService {
+
+	private static final Logger LOG = Logger.getLogger( ExtensionIntegrationServiceImpl.class );
+
+	private final LinkedHashSet<ExtensionIntegration<?>> integrators = new LinkedHashSet<>();
+
+	private ExtensionIntegrationServiceImpl() {
+	}
+
+	public static ExtensionIntegrationServiceImpl create(Set<ExtensionIntegration<?>> integrations, ClassLoaderService classLoaderService) {
+		ExtensionIntegrationServiceImpl instance = new ExtensionIntegrationServiceImpl();
+
+		// register provided integrators
+		for ( ExtensionIntegration<?> integration : integrations ) {
+			instance.addExtensionIntegration( integration );
+		}
+		for ( ExtensionIntegration<?> integration : classLoaderService.loadJavaServices(
+				ExtensionIntegration.class ) ) {
+			instance.addExtensionIntegration( integration );
+		}
+
+		return instance;
+	}
+
+	private void addExtensionIntegration(ExtensionIntegration<?> integration) {
+		if ( LOG.isDebugEnabled() ) {
+			LOG.debugf( "Adding extension integration for [%s]", integration.getExtensionType().getName() );
+		}
+		integrators.add( integration );
+	}
+
+	@Override
+	public Iterable<ExtensionIntegration<?>> extensionIntegrations() {
+		return integrators;
+	}
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/Extension.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/Extension.java
@@ -0,0 +1,12 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
+
+/// A marker interface for Session extensions.
+@Incubating
+public interface Extension {
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegration.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegration.java
@@ -0,0 +1,14 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
+
+@Incubating
+public interface ExtensionIntegration<E extends Extension> {
+	Class<E> getExtensionType();
+
+	E createExtension(ExtensionIntegrationContext context);
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationContext.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationContext.java
@@ -0,0 +1,14 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+
+@Incubating
+public interface ExtensionIntegrationContext {
+
+	SharedSessionContractImplementor getSession();
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationService.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationService.java
@@ -0,0 +1,18 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
+import org.hibernate.service.Service;
+
+@Incubating
+public interface ExtensionIntegrationService extends Service {
+	/**
+	 * Retrieve all extensions.
+	 *
+	 * @return All extensions.
+	 */
+	Iterable<ExtensionIntegration<?>> extensionIntegrations();
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationServiceInitiator.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/ExtensionIntegrationServiceInitiator.java
@@ -0,0 +1,29 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
+import org.hibernate.engine.extension.internal.ExtensionIntegrationServiceImpl;
+import org.hibernate.service.spi.SessionFactoryServiceInitiator;
+import org.hibernate.service.spi.SessionFactoryServiceInitiatorContext;
+
+import java.util.Set;
+
+@Incubating
+public class ExtensionIntegrationServiceInitiator
+		implements SessionFactoryServiceInitiator<ExtensionIntegrationService> {
+
+	public static final ExtensionIntegrationServiceInitiator INSTANCE = new ExtensionIntegrationServiceInitiator();
+
+	@Override
+	public ExtensionIntegrationService initiateService(SessionFactoryServiceInitiatorContext context) {
+		return ExtensionIntegrationServiceImpl.create( Set.of(), context.getSessionFactory().getClassLoaderService() );
+	}
+
+	@Override
+	public Class<ExtensionIntegrationService> getServiceInitiated() {
+		return ExtensionIntegrationService.class;
+	}
+}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/package-info.java b/hibernate-core/src/main/java/org/hibernate/engine/extension/spi/package-info.java
@@ -0,0 +1,7 @@
+/**
+ * This package contains an SPI for Session extensions.
+ */
+@Incubating
+package org.hibernate.engine.extension.spi;
+
+import org.hibernate.Incubating;
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/spi/SessionDelegatorBaseImpl.java b/hibernate-core/src/main/java/org/hibernate/engine/spi/SessionDelegatorBaseImpl.java
@@ -43,6 +43,7 @@
 import org.hibernate.bytecode.enhance.spi.interceptor.SessionAssociationMarkers;
 import org.hibernate.cache.spi.CacheTransactionSynchronization;
 import org.hibernate.collection.spi.PersistentCollection;
+import org.hibernate.engine.extension.spi.Extension;
 import org.hibernate.engine.jdbc.LobCreator;
 import org.hibernate.engine.jdbc.connections.spi.JdbcConnectionAccess;
 import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
@@ -517,6 +518,11 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
 		return delegate.getEntityGraph( graphName );
 	}
 
+	@Override
+	public <T extends Extension> T getExtension(Class<T> extension) {
+		return delegate.getExtension( extension );
+	}
+
 	@Override
 	public <T> QueryImplementor<T> createQuery(CriteriaSelect<T> selectQuery) {
 		return delegate.createQuery( selectQuery );
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionContractImplementor.java b/hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionContractImplementor.java
@@ -6,6 +6,7 @@
 
 import java.util.Set;
 import java.util.UUID;
+
 import jakarta.persistence.TransactionRequiredException;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -18,6 +19,7 @@
 import org.hibernate.StatelessSession;
 import org.hibernate.bytecode.enhance.spi.interceptor.SessionAssociationMarkers;
 import org.hibernate.dialect.Dialect;
+import org.hibernate.engine.extension.spi.Extension;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.graph.spi.RootGraphImplementor;
 import org.hibernate.query.Query;
@@ -621,4 +623,17 @@ default boolean isStatelessSession() {
 
 	@Override
 	RootGraphImplementor<?> getEntityGraph(String graphName);
+
+	/**
+	 * Allows accessing session scoped extension storages of the particular session instance.
+	 * <p>
+	 * Extensions first had to be registered with the {@link org.hibernate.SessionFactory}
+	 *
+	 * @param extension The extension storage attached to the current session.
+	 * if the current session does not yet have the particular storage type attached to this session.
+	 * @param <E> The type of the extension storage.
+	 */
+	@Incubating
+	<E extends Extension> E getExtension(Class<E> extension);
+
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionDelegatorBaseImpl.java b/hibernate-core/src/main/java/org/hibernate/engine/spi/SharedSessionDelegatorBaseImpl.java
@@ -24,6 +24,7 @@
 import org.hibernate.bytecode.enhance.spi.interceptor.SessionAssociationMarkers;
 import org.hibernate.cache.spi.CacheTransactionSynchronization;
 import org.hibernate.collection.spi.PersistentCollection;
+import org.hibernate.engine.extension.spi.Extension;
 import org.hibernate.engine.jdbc.LobCreator;
 import org.hibernate.engine.jdbc.connections.spi.JdbcConnectionAccess;
 import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
@@ -663,6 +664,11 @@ public RootGraphImplementor<?> getEntityGraph(String graphName) {
 		return delegate.getEntityGraph( graphName );
 	}
 
+	@Override
+	public <T extends Extension> T getExtension(Class<T> extension) {
+		return delegate.getExtension( extension);
+	}
+
 	@Override
 	public <T> List<EntityGraph<? super T>> getEntityGraphs(Class<T> entityClass) {
 		return delegate.getEntityGraphs( entityClass );
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java b/hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java
@@ -33,6 +33,10 @@
 import org.hibernate.engine.creation.internal.SharedSessionBuilderImpl;
 import org.hibernate.engine.creation.internal.SharedSessionCreationOptions;
 import org.hibernate.engine.creation.internal.SharedStatelessSessionBuilderImpl;
+import org.hibernate.engine.extension.spi.Extension;
+import org.hibernate.engine.extension.spi.ExtensionIntegration;
+import org.hibernate.engine.extension.spi.ExtensionIntegrationContext;
+import org.hibernate.engine.extension.spi.ExtensionIntegrationService;
 import org.hibernate.engine.internal.SessionEventListenerManagerImpl;
 import org.hibernate.engine.jdbc.LobCreator;
 import org.hibernate.engine.jdbc.connections.spi.JdbcConnectionAccess;
@@ -111,8 +115,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;
@@ -139,7 +145,7 @@
  *
  * @author Steve Ebersole
  */
-public abstract class AbstractSharedSessionContract implements SharedSessionContractImplementor {
+public abstract class AbstractSharedSessionContract implements SharedSessionContractImplementor, ExtensionIntegrationContext {
 
 	private transient SessionFactoryImpl factory;
 	private transient SessionFactoryOptions factoryOptions;
@@ -186,6 +192,8 @@ public abstract class AbstractSharedSessionContract implements SharedSessionCont
 	private transient ExceptionConverter exceptionConverter;
 	private transient SessionAssociationMarkers sessionAssociationMarkers;
 
+	private transient final Map<Class<?>, Object> extensions;
+
 	public AbstractSharedSessionContract(SessionFactoryImpl factory, SessionCreationOptions options) {
 		this.factory = factory;
 
@@ -248,6 +256,13 @@ public void onParentClose() {
 			transactionCoordinator = factory.transactionCoordinatorBuilder
 					.buildTransactionCoordinator( jdbcCoordinator, this );
 		}
+
+		extensions = new HashMap<>();
+		for ( ExtensionIntegration<?> integration : factory.getServiceRegistry()
+				.requireService( ExtensionIntegrationService.class )
+				.extensionIntegrations() ) {
+			extensions.put( integration.getExtensionType(), integration.createExtension( this ) );
+		}
 	}
 
 	final SessionFactoryOptions getSessionFactoryOptions() {
@@ -472,6 +487,11 @@ public final UUID getSessionIdentifier() {
 		return sessionIdentifier;
 	}
 
+	@Override
+	public SharedSessionContractImplementor getSession() {
+		return this;
+	}
+
 	@Override
 	public final Object getSessionToken() {
 		if ( sessionToken == null ) {
@@ -1704,6 +1724,11 @@ public SessionAssociationMarkers getSessionAssociationMarkers() {
 		return sessionAssociationMarkers;
 	}
 
+	@Override
+	public <E extends Extension> E getExtension(Class<E> extension) {
+		return extension.cast( extensions.get( extension ) );
+	}
+
 	@Serial
 	private void writeObject(ObjectOutputStream oos) throws IOException {
 		SESSION_LOGGER.serializingSession( getSessionIdentifier() );
diff --git a/hibernate-core/src/main/java/org/hibernate/service/internal/StandardSessionFactoryServiceInitiators.java b/hibernate-core/src/main/java/org/hibernate/service/internal/StandardSessionFactoryServiceInitiators.java
@@ -7,6 +7,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.hibernate.engine.extension.spi.ExtensionIntegrationServiceInitiator;
 import org.hibernate.engine.query.spi.NativeQueryInterpreterInitiator;
 import org.hibernate.engine.spi.CacheInitiator;
 import org.hibernate.service.spi.SessionFactoryServiceInitiator;
@@ -25,6 +26,7 @@ public static List<SessionFactoryServiceInitiator<?>> buildStandardServiceInitia
 		serviceInitiators.add( StatisticsInitiator.INSTANCE );
 		serviceInitiators.add( CacheInitiator.INSTANCE );
 		serviceInitiators.add( NativeQueryInterpreterInitiator.INSTANCE );
+		serviceInitiators.add( ExtensionIntegrationServiceInitiator.INSTANCE );
 		return serviceInitiators;
 	}
 
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/engine/spi/SessionExtensionTest.java b/hibernate-core/src/test/java/org/hibernate/orm/test/engine/spi/SessionExtensionTest.java
