HHH-16626 - JPA hint for Session (EntityManager) level tenant-id

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

@@ -128,6 +128,7 @@ import static org.hibernate.cfg.AvailableSettings.JAKARTA_VALIDATION_FACTORY;
 import static org.hibernate.cfg.AvailableSettings.JPA_VALIDATION_FACTORY;
 import static org.hibernate.internal.FetchProfileHelper.getFetchProfiles;
 import static org.hibernate.internal.util.config.ConfigurationHelper.getBoolean;
+import static org.hibernate.jpa.HibernateHints.HINT_TENANT_ID;
 import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
 import static org.hibernate.resource.jdbc.spi.PhysicalConnectionHandlingMode.DELAYED_ACQUISITION_AND_RELEASE_AFTER_STATEMENT;
 
@@ -701,12 +702,23 @@ public class SessionFactoryImpl extends QueryParameterBindingTypeResolverImpl im
 		SessionBuilderImplementor builder = withOptions();
 		builder.autoJoinTransactions( synchronizationType == SYNCHRONIZED );
 
+		if ( map != null ) {
+			//noinspection SuspiciousMethodCalls
+			final String tenantIdHint = (String) map.get( HINT_TENANT_ID );
+			if ( tenantIdHint != null ) {
+				builder = (SessionBuilderImplementor) builder.tenantIdentifier( tenantIdHint );
+			}
+		}
+
 		final Session session = builder.openSession();
 		if ( map != null ) {
 			for ( Map.Entry<K, V> o : map.entrySet() ) {
 				final K key = o.getKey();
 				if ( key instanceof String ) {
 					final String sKey = (String) key;
+					if ( HINT_TENANT_ID.equals( sKey ) ) {
+						continue;
+					}
 					session.setProperty( sKey, o.getValue() );
 				}
 			}

hibernate-core/src/main/java/org/hibernate/jpa/HibernateHints.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.jpa;
 
+import jakarta.persistence.EntityManagerFactory;
+
 /**
  * List of Hibernate-specific (extension) hints available to query,
  * load, and lock scenarios.
@@ -156,4 +158,12 @@ public interface HibernateHints {
 	 * to a function rather than a call to a procedure.
 	 */
 	String HINT_CALLABLE_FUNCTION = "org.hibernate.callableFunction";
+
+	/**
+	 * Hint for specifying the tenant-id to use when creating an
+	 * {@linkplain EntityManagerFactory#createEntityManager(java.util.Map) EntityManager}.
+	 *
+	 * @see org.hibernate.SessionBuilder#tenantIdentifier
+	 */
+	String HINT_TENANT_ID = "org.hibernate.tenantId";
 }

hibernate-core/src/test/java/org/hibernate/orm/test/tenantid/TenantIdTest.java

@@ -8,6 +8,7 @@ package org.hibernate.orm.test.tenantid;
 
 import org.hibernate.HibernateError;
 import org.hibernate.PropertyValueException;
+import org.hibernate.Session;
 import org.hibernate.boot.SessionFactoryBuilder;
 import org.hibernate.boot.spi.MetadataImplementor;
 import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
@@ -20,13 +21,19 @@ import org.hibernate.testing.orm.junit.SessionFactoryProducer;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
 import org.hibernate.testing.orm.junit.Setting;
 import org.hibernate.binder.internal.TenantIdBinder;
+import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.type.descriptor.DateTimeUtils;
 
 import org.hibernate.testing.orm.junit.SkipForDialect;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
+import jakarta.persistence.EntityManager;
+import jakarta.persistence.EntityManagerFactory;
+
 import static org.hibernate.cfg.AvailableSettings.JAKARTA_HBM2DDL_DATABASE_ACTION;
+import static org.hibernate.internal.util.collections.CollectionHelper.toMap;
+import static org.hibernate.jpa.HibernateHints.HINT_TENANT_ID;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -161,6 +168,58 @@ public class TenantIdTest implements SessionFactoryProducer {
         } );
     }
 
+    @Test
+    public void testEntityManagerHint(SessionFactoryScope scope) {
+        currentTenant = "mine";
+        Record record = new Record();
+        scope.inTransaction( s -> s.persist( record ) );
+        assertEquals( "mine", record.state.tenantId );
+        assertNotNull( record.state.updated );
+
+        currentTenant = "yours";
+        Record record2 = new Record();
+        scope.inTransaction( s -> s.persist( record2 ) );
+        assertEquals( "yours", record2.state.tenantId );
+        assertNotNull( record2.state.updated );
+
+        currentTenant = null;
+        final EntityManagerFactory emf = scope.getSessionFactory();
+        try (EntityManager em = emf.createEntityManager( toMap( HINT_TENANT_ID, "mine" ) ) ) {
+            Record r = em.find( Record.class, record.id );
+            assertEquals( "mine", r.state.tenantId );
+
+            // Session seems to not apply tenant-id on #find
+            Record yours = em.find( Record.class, record2.id );
+            assertEquals( "yours", yours.state.tenantId );
+
+
+            em.createQuery( "from Record where id = :id", Record.class )
+                    .setParameter( "id", record.id )
+                    .getSingleResult();
+            assertEquals( "mine", r.state.tenantId );
+
+            // However, Session does seem to apply tenant-id on queries
+            try {
+                em.createQuery( "from Record where id = :id", Record.class )
+                        .setParameter( "id", record2.id )
+                        .getSingleResult();
+                fail( "Expecting an exception" );
+            }
+            catch (Exception expected) {
+            }
+        }
+        catch (RuntimeException e) {
+            currentTenant = "yours";
+            scope.inTransaction( (s) -> s.createMutationQuery( "delete Record" ) );
+
+            throw e;
+        }
+        finally {
+            // for cleanup
+            currentTenant = "mine";
+        }
+    }
+
     private static void waitALittle() {
         try {
             Thread.sleep( 10 );