HHH-19559 add hibernate.multi_tenant.set_schema

gavinking · gavinking · commit 75d60315bef9 · 2025-06-28T16:21:38.000+02:00
diff --git a/hibernate-core/src/main/java/org/hibernate/boot/internal/SessionFactoryOptionsBuilder.java b/hibernate-core/src/main/java/org/hibernate/boot/internal/SessionFactoryOptionsBuilder.java
@@ -187,6 +187,7 @@ public class SessionFactoryOptionsBuilder implements SessionFactoryOptions {
 	// multi-tenancy
 	private boolean multiTenancyEnabled;
 	private CurrentTenantIdentifierResolver<Object> currentTenantIdentifierResolver;
+	private boolean setTenantSchemaEnabled;
 
 	// Queries
 	private SqmFunctionRegistry sqmFunctionRegistry;
@@ -371,6 +372,7 @@ public SessionFactoryOptionsBuilder(StandardServiceRegistry serviceRegistry, Boo
 				null
 			);
 		}
+		setTenantSchemaEnabled = configurationService.getSetting( SET_TENANT_SCHEMA, BOOLEAN, false );
 
 		delayBatchFetchLoaderCreations =
 				configurationService.getSetting( DELAY_ENTITY_LOADER_CREATIONS, BOOLEAN, true );
@@ -1003,6 +1005,11 @@ public boolean isMultiTenancyEnabled() {
 		return multiTenancyEnabled;
 	}
 
+	@Override
+	public boolean isSetTenantSchemaEnabled() {
+		return setTenantSchemaEnabled;
+	}
+
 	@Override
 	public CurrentTenantIdentifierResolver<Object> getCurrentTenantIdentifierResolver() {
 		return currentTenantIdentifierResolver;
@@ -1450,6 +1457,10 @@ public void applyCurrentTenantIdentifierResolver(CurrentTenantIdentifierResolver
 		this.currentTenantIdentifierResolver = (CurrentTenantIdentifierResolver<Object>) resolver;
 	}
 
+	public void applySetTenantSchema(boolean enabled) {
+		this.setTenantSchemaEnabled = enabled;
+	}
+
 	public void enableNamedQueryCheckingOnStartup(boolean enabled) {
 		this.namedQueryStartupCheckingEnabled = enabled;
 	}
diff --git a/hibernate-core/src/main/java/org/hibernate/boot/spi/AbstractDelegatingSessionFactoryOptions.java b/hibernate-core/src/main/java/org/hibernate/boot/spi/AbstractDelegatingSessionFactoryOptions.java
@@ -219,6 +219,11 @@ public boolean isMultiTenancyEnabled() {
 		return delegate.isMultiTenancyEnabled();
 	}
 
+	@Override
+	public boolean isSetTenantSchemaEnabled() {
+		return delegate.isSetTenantSchemaEnabled();
+	}
+
 	@Override
 	public CurrentTenantIdentifierResolver<Object> getCurrentTenantIdentifierResolver() {
 		return delegate.getCurrentTenantIdentifierResolver();
diff --git a/hibernate-core/src/main/java/org/hibernate/boot/spi/SessionFactoryOptions.java b/hibernate-core/src/main/java/org/hibernate/boot/spi/SessionFactoryOptions.java
@@ -304,6 +304,17 @@ default SessionEventListener[] buildSessionEventListeners() {
 	 */
 	boolean isMultiTenancyEnabled();
 
+	/**
+	 * Should the schema be {@linkplain java.sql.Connection#setSchema set} to the
+	 * {@linkplain CurrentTenantIdentifierResolver#schemaName schema belonging to
+	 * the current tenant} each time a connection is obtained?
+	 *
+	 * @see org.hibernate.cfg.MultiTenancySettings#SET_TENANT_SCHEMA
+	 *
+	 * @since 7.1
+	 */
+	boolean isSetTenantSchemaEnabled();
+
 	/**
 	 * Obtain a reference to the
 	 * {@linkplain CurrentTenantIdentifierResolver current tenant identifier resolver},
diff --git a/hibernate-core/src/main/java/org/hibernate/cfg/MultiTenancySettings.java b/hibernate-core/src/main/java/org/hibernate/cfg/MultiTenancySettings.java
@@ -26,7 +26,7 @@ public interface MultiTenancySettings {
 	 * either:
 	 * <ul>
 	 *     <li>an instance of {@code CurrentTenantIdentifierResolver},
-	 *     <li>a {@link Class} representing an class that implements {@code CurrentTenantIdentifierResolver}, or
+	 *     <li>a {@link Class} representing a class that implements {@code CurrentTenantIdentifierResolver}, or
 	 *     <li>the name of a class that implements {@code CurrentTenantIdentifierResolver}.
 	 * </ul>
 	 *
@@ -42,4 +42,17 @@ public interface MultiTenancySettings {
 	 * This setting configures the name of the DataSource to use for this access
 	 */
 	String TENANT_IDENTIFIER_TO_USE_FOR_ANY_KEY = "hibernate.multi_tenant.datasource.identifier_for_any";
+
+	/**
+	 * Specifies that {@link java.sql.Connection#setSchema(String)}}
+	 * should be called with the current tenant id, or with the schema name
+	 * returned by {@link CurrentTenantIdentifierResolver#schemaName}.
+	 *
+	 * @settingDefault false
+	 *
+	 * @see CurrentTenantIdentifierResolver#schemaName
+	 *
+	 * @since 7.1
+	 */
+	String SET_TENANT_SCHEMA = "hibernate.multi_tenant.set_schema";
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/context/spi/AbstractCurrentSessionContext.java b/hibernate-core/src/main/java/org/hibernate/context/spi/AbstractCurrentSessionContext.java
@@ -34,15 +34,15 @@ public SessionFactoryImplementor factory() {
 
 	protected SessionBuilder baseSessionBuilder() {
 		final SessionBuilderImplementor builder = factory.withOptions();
-		final CurrentTenantIdentifierResolver<Object> resolver = factory.getCurrentTenantIdentifierResolver();
+		final var resolver = factory.getCurrentTenantIdentifierResolver();
 		if ( resolver != null ) {
 			builder.tenantIdentifier( resolver.resolveCurrentTenantIdentifier() );
 		}
 		return builder;
 	}
 
 	protected void validateExistingSession(Session existingSession) {
-		final CurrentTenantIdentifierResolver<Object> resolver = factory.getCurrentTenantIdentifierResolver();
+		final var resolver = factory.getCurrentTenantIdentifierResolver();
 		if ( resolver != null && resolver.validateExistingCurrentSessions() ) {
 			final Object currentValue = resolver.resolveCurrentTenantIdentifier();
 			final JavaType<Object> tenantIdentifierJavaType = factory.getTenantIdentifierJavaType();
diff --git a/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java b/hibernate-core/src/main/java/org/hibernate/context/spi/CurrentTenantIdentifierResolver.java
@@ -4,6 +4,8 @@
  */
 package org.hibernate.context.spi;
 
+import org.checkerframework.checker.nullness.qual.NonNull;
+
 /**
  * A callback registered with the {@link org.hibernate.SessionFactory} that is
  * responsible for resolving the current tenant identifier.
@@ -48,4 +50,20 @@ public interface CurrentTenantIdentifierResolver<T> {
 	default boolean isRoot(T tenantId) {
 		return false;
 	}
+
+	/**
+	 * The name of the database schema for data belonging to the tenant with the
+	 * given identifier.
+	 * <p>
+	 * Called when {@value org.hibernate.cfg.MultiTenancySettings#SET_TENANT_SCHEMA}
+	 * is enabled.
+	 *
+	 * @param tenantIdentifier The tenant identifier
+	 * @return The name of the database schema belonging to that tenant
+	 *
+	 * @see org.hibernate.cfg.MultiTenancySettings#SET_TENANT_SCHEMA
+	 */
+	default @NonNull String schemaName(@NonNull T tenantIdentifier) {
+		return tenantIdentifier.toString();
+	}
 }
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
@@ -67,6 +67,8 @@
 import org.hibernate.stat.SessionStatistics;
 import org.hibernate.type.format.FormatMapper;
 
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -1247,4 +1249,14 @@ public SessionAssociationMarkers getSessionAssociationMarkers() {
 	public boolean isIdentifierRollbackEnabled() {
 		return delegate.isIdentifierRollbackEnabled();
 	}
+
+	@Override
+	public void afterObtainConnection(Connection connection) throws SQLException {
+		delegate.afterObtainConnection( connection );
+	}
+
+	@Override
+	public void beforeReleaseConnection(Connection connection) throws SQLException {
+		delegate.beforeReleaseConnection( connection );
+	}
 }
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
@@ -4,6 +4,8 @@
  */
 package org.hibernate.engine.spi;
 
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.Set;
 import java.util.UUID;
 import jakarta.persistence.TransactionRequiredException;
@@ -584,4 +586,31 @@ default boolean isStatelessSession() {
 	 */
 	@Incubating
 	SessionAssociationMarkers getSessionAssociationMarkers();
+
+	/**
+	 * Called after a non-contextual JDBC connection is obtained.
+	 * <p>
+	 * Sets the schema to the schema belonging to the current tenant if:
+	 * <ol>
+	 * <li>{@value org.hibernate.cfg.MultiTenancySettings#SET_TENANT_SCHEMA} is enabled, and
+	 * <li>this session has an active tenant id.
+	 * </ol>
+	 *
+	 * @param connection A JDBC connection which was just acquired
+	 *
+	 * @since 7.1
+	 */
+	void afterObtainConnection(Connection connection) throws SQLException;
+
+	/**
+	 * Called right before non-contextual a JDBC connection is released.
+	 * <p>
+	 * Unset the schema which was set by {@link #afterObtainConnection},
+	 * if any.
+	 * .
+	 * @param connection The JDBC connection which is about to be released
+	 *
+	 * @since 7.1
+	 */
+	void beforeReleaseConnection(Connection connection) throws SQLException;
 }
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
@@ -44,6 +44,8 @@
 import org.hibernate.resource.transaction.spi.TransactionCoordinator;
 import org.hibernate.type.format.FormatMapper;
 
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.List;
 import java.util.Set;
 import java.util.TimeZone;
@@ -703,4 +705,14 @@ public SessionAssociationMarkers getSessionAssociationMarkers() {
 	public boolean isIdentifierRollbackEnabled() {
 		return delegate.isIdentifierRollbackEnabled();
 	}
+
+	@Override
+	public void afterObtainConnection(Connection connection) throws SQLException {
+		delegate.afterObtainConnection( connection );
+	}
+
+	@Override
+	public void beforeReleaseConnection(Connection connection) throws SQLException {
+		delegate.beforeReleaseConnection(  connection );
+	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/generator/internal/TenantIdGeneration.java b/hibernate-core/src/main/java/org/hibernate/generator/internal/TenantIdGeneration.java
@@ -9,7 +9,6 @@
 
 import org.hibernate.PropertyValueException;
 import org.hibernate.annotations.TenantId;
-import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.generator.BeforeExecutionGenerator;
@@ -51,8 +50,7 @@ public Object generate(SharedSessionContractImplementor session, Object owner, O
 		final SessionFactoryImplementor sessionFactory = session.getSessionFactory();
 		final Object tenantId = session.getTenantIdentifierValue();
 		if ( currentValue != null ) {
-			final CurrentTenantIdentifierResolver<Object> resolver =
-					sessionFactory.getCurrentTenantIdentifierResolver();
+			final var resolver = sessionFactory.getCurrentTenantIdentifierResolver();
 			if ( resolver != null && resolver.isRoot( tenantId ) ) {
 				// the "root" tenant is allowed to set the tenant id explicitly
 				return currentValue;
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java b/hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java
@@ -28,7 +28,6 @@
 import org.hibernate.boot.spi.SessionFactoryOptions;
 import org.hibernate.bytecode.enhance.spi.interceptor.SessionAssociationMarkers;
 import org.hibernate.cache.spi.CacheTransactionSynchronization;
-import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
 import org.hibernate.dialect.Dialect;
 import org.hibernate.engine.internal.SessionEventListenerManagerImpl;
 import org.hibernate.engine.jdbc.LobCreator;
@@ -108,6 +107,7 @@
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.Serial;
+import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.List;
 import java.util.Locale;
@@ -245,7 +245,7 @@ protected final void setUpMultitenancy(SessionFactoryImplementor factory, LoadQu
 				throw new HibernateException( "SessionFactory configured for multi-tenancy, but no tenant identifier specified" );
 			}
 			else {
-				final CurrentTenantIdentifierResolver<Object> resolver = factory.getCurrentTenantIdentifierResolver();
+				final var resolver = factory.getCurrentTenantIdentifierResolver();
 				if ( resolver==null || !resolver.isRoot( tenantIdentifier ) ) {
 					// turn on the filter, unless this is the "root" tenant with access to all partitions
 					loadQueryInfluencers.enableFilter( TenantIdBinder.FILTER_NAME )
@@ -683,16 +683,18 @@ public boolean isConnected() {
 
 	@Override
 	public JdbcConnectionAccess getJdbcConnectionAccess() {
-		// See class-level JavaDocs for a discussion of the concurrent-access safety of this method
+		// See class-level Javadoc for a discussion of the concurrent-access safety of this method
 		if ( jdbcConnectionAccess == null ) {
 			if ( !factoryOptions.isMultiTenancyEnabled() ) {
+				// we might still be using schema-based multitenancy
 				jdbcConnectionAccess = new NonContextualJdbcConnectionAccess(
 						sessionEventsManager,
 						factory.connectionProvider,
 						this
 				);
 			}
 			else {
+				// we're using datasource-based multitenancy
 				jdbcConnectionAccess = new ContextualJdbcConnectionAccess(
 						tenantIdentifier,
 						sessionEventsManager,
@@ -704,6 +706,32 @@ public JdbcConnectionAccess getJdbcConnectionAccess() {
 		return jdbcConnectionAccess;
 	}
 
+	private boolean useSchemaBasedMultiTenancy() {
+		return tenantIdentifier != null
+			&& factory.getSessionFactoryOptions().isSetTenantSchemaEnabled();
+	}
+
+	private String tenantSchema() {
+		final var tenantIdResolver = factory.getCurrentTenantIdentifierResolver();
+		return tenantIdResolver == null
+				? (String) tenantIdentifier
+				: tenantIdResolver.schemaName( tenantIdentifier );
+	}
+
+	@Override
+	public void afterObtainConnection(Connection connection) throws SQLException {
+		if ( useSchemaBasedMultiTenancy() ) {
+			connection.setSchema( tenantSchema() );
+		}
+	}
+
+	@Override
+	public void beforeReleaseConnection(Connection connection) throws SQLException {
+		if ( useSchemaBasedMultiTenancy() ) {
+			connection.setSchema( null );
+		}
+	}
+
 	@Override
 	public EntityKey generateEntityKey(Object id, EntityPersister persister) {
 		return new EntityKey( id, persister );
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/NonContextualJdbcConnectionAccess.java b/hibernate-core/src/main/java/org/hibernate/internal/NonContextualJdbcConnectionAccess.java
@@ -15,6 +15,7 @@
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.event.monitor.spi.EventMonitor;
 import org.hibernate.event.monitor.spi.DiagnosticEvent;
+import org.jboss.logging.Logger;
 
 /**
  * @author Steve Ebersole
@@ -24,6 +25,8 @@ public class NonContextualJdbcConnectionAccess implements JdbcConnectionAccess,
 	private final ConnectionProvider connectionProvider;
 	private final SharedSessionContractImplementor session;
 
+	private static final Logger log = Logger.getLogger( NonContextualJdbcConnectionAccess.class );
+
 	public NonContextualJdbcConnectionAccess(
 			SessionEventListener listener,
 			ConnectionProvider connectionProvider,
@@ -39,9 +42,10 @@ public NonContextualJdbcConnectionAccess(
 	public Connection obtainConnection() throws SQLException {
 		final EventMonitor eventMonitor = session.getEventMonitor();
 		final DiagnosticEvent jdbcConnectionAcquisitionEvent = eventMonitor.beginJdbcConnectionAcquisitionEvent();
+		final Connection connection;
 		try {
 			listener.jdbcConnectionAcquisitionStart();
-			return connectionProvider.getConnection();
+			connection = connectionProvider.getConnection();
 		}
 		finally {
 			eventMonitor.completeJdbcConnectionAcquisitionEvent(
@@ -51,10 +55,31 @@ public Connection obtainConnection() throws SQLException {
 			);
 			listener.jdbcConnectionAcquisitionEnd();
 		}
+
+		try {
+			session.afterObtainConnection( connection );
+		}
+		catch (SQLException e) {
+			try {
+				releaseConnection( connection );
+			}
+			catch (SQLException re) {
+				e.addSuppressed( re );
+			}
+			throw e;
+		}
+		return connection;
 	}
 
 	@Override
 	public void releaseConnection(Connection connection) throws SQLException {
+		try {
+			session.beforeReleaseConnection( connection );
+		}
+		catch (SQLException e) {
+			log.warn( "Error before releasing JDBC connection", e );
+		}
+
 		final EventMonitor eventMonitor = session.getEventMonitor();
 		final DiagnosticEvent jdbcConnectionReleaseEvent = eventMonitor.beginJdbcConnectionReleaseEvent();
 		try {
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/SessionFactoryImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/SessionFactoryImpl.java
