HHH-17117 allow @TenantID to form part of composite key

beikov · beikov · commit d6c1f696ef8c · 2025-07-31T11:03:51.000+02:00
diff --git a/hibernate-core/src/main/java/org/hibernate/boot/internal/InFlightMetadataCollectorImpl.java b/hibernate-core/src/main/java/org/hibernate/boot/internal/InFlightMetadataCollectorImpl.java
@@ -2133,7 +2133,8 @@ private void processExportableProducers() {
 			handleIdentifierValueBinding(
 					entityBinding.getIdentifier(),
 					dialect,
-					(RootClass) entityBinding
+					(RootClass) entityBinding,
+					entityBinding.getIdentifierProperty()
 			);
 
 		}
@@ -2146,15 +2147,14 @@ private void processExportableProducers() {
 			handleIdentifierValueBinding(
 					( (IdentifierCollection) collection ).getIdentifier(),
 					dialect,
+					null,
 					null
 			);
 		}
 	}
 
 	private void handleIdentifierValueBinding(
-			KeyValue identifierValueBinding,
-			Dialect dialect,
-			RootClass entityBinding) {
+			KeyValue identifierValueBinding, Dialect dialect, RootClass entityBinding, Property identifierProperty) {
 		// todo : store this result (back into the entity or into the KeyValue, maybe?)
 		// 		This process of instantiating the id-generator is called multiple times.
 		//		It was done this way in the old code too, so no "regression" here; but
@@ -2163,7 +2163,8 @@ private void handleIdentifierValueBinding(
 			final Generator generator = identifierValueBinding.createGenerator(
 					bootstrapContext.getIdentifierGeneratorFactory(),
 					dialect,
-					entityBinding
+					entityBinding,
+					identifierProperty
 			);
 
 			if ( generator instanceof ExportableProducer ) {
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/SessionFactoryImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/SessionFactoryImpl.java
@@ -467,7 +467,8 @@ private static Map<String, Generator> createGenerators(
 					final Generator generator = id.createGenerator(
 							bootstrapContext.getIdentifierGeneratorFactory(),
 							jdbcServices.getJdbcEnvironment().getDialect(),
-							(RootClass) model
+							(RootClass) model,
+							model.getIdentifierProperty()
 					);
 					if ( generator instanceof Configurable ) {
 						final Configurable identifierGenerator = (Configurable) generator;
diff --git a/hibernate-core/src/main/java/org/hibernate/mapping/Component.java b/hibernate-core/src/main/java/org/hibernate/mapping/Component.java
@@ -671,12 +671,19 @@ public Generator createGenerator(
 			IdentifierGeneratorFactory identifierGeneratorFactory,
 			Dialect dialect,
 			RootClass rootClass) throws MappingException {
+		return createGenerator( identifierGeneratorFactory, dialect, rootClass, rootClass.getIdentifierProperty() );
+	}
+
+	@Override
+	public Generator createGenerator(
+			IdentifierGeneratorFactory identifierGeneratorFactory,
+			Dialect dialect,
+			RootClass rootClass,
+			Property property) throws MappingException {
 		if ( builtIdentifierGenerator == null ) {
-			builtIdentifierGenerator = buildIdentifierGenerator(
-					identifierGeneratorFactory,
-					dialect,
-					rootClass
-			);
+			builtIdentifierGenerator = DEFAULT_ID_GEN_STRATEGY.equals( getIdentifierGeneratorStrategy() )
+					? buildIdentifierGenerator( identifierGeneratorFactory, dialect, rootClass )
+					: super.createGenerator( identifierGeneratorFactory, dialect, rootClass, property );
 		}
 		return builtIdentifierGenerator;
 	}
@@ -685,32 +692,10 @@ private Generator buildIdentifierGenerator(
 			IdentifierGeneratorFactory identifierGeneratorFactory,
 			Dialect dialect,
 			RootClass rootClass) throws MappingException {
-		final boolean hasCustomGenerator = ! DEFAULT_ID_GEN_STRATEGY.equals( getIdentifierGeneratorStrategy() );
-		if ( hasCustomGenerator ) {
-			return super.createGenerator( identifierGeneratorFactory, dialect, rootClass );
-		}
-
-		final Class<?> entityClass = rootClass.getMappedClass();
-		final Class<?> attributeDeclarer; // what class is the declarer of the composite pk attributes
-		// IMPL NOTE : See the javadoc discussion on CompositeNestedGeneratedValueGenerator wrt the
-		//		various scenarios for which we need to account here
-		if ( rootClass.getIdentifierMapper() != null ) {
-			// we have the @IdClass / <composite-id mapped="true"/> case
-			attributeDeclarer = resolveComponentClass();
-		}
-		else if ( rootClass.getIdentifierProperty() != null ) {
-			// we have the "@EmbeddedId" / <composite-id name="idName"/> case
-			attributeDeclarer = resolveComponentClass();
-		}
-		else {
-			// we have the "straight up" embedded (again the Hibernate term) component identifier
-			attributeDeclarer = entityClass;
-		}
-
-		final CompositeNestedGeneratedValueGenerator.GenerationContextLocator locator =
-				new StandardGenerationContextLocator( rootClass.getEntityName() );
-		final CompositeNestedGeneratedValueGenerator generator =
-				new CompositeNestedGeneratedValueGenerator( locator, getType() );
+		final CompositeNestedGeneratedValueGenerator generator = new CompositeNestedGeneratedValueGenerator(
+				new StandardGenerationContextLocator( rootClass.getEntityName() ),
+				getType()
+		);
 
 		final List<Property> properties = getProperties();
 		for ( int i = 0; i < properties.size(); i++ ) {
@@ -723,15 +708,37 @@ else if ( rootClass.getIdentifierProperty() != null ) {
 					// skip any 'assigned' generators, they would have been handled by
 					// the StandardGenerationContextLocator
 					generator.addGeneratedValuePlan( new ValueGenerationPlan(
-							value.createGenerator( identifierGeneratorFactory, dialect, rootClass ),
-							getType().isMutable() ? injector( property, attributeDeclarer ) : null,
+							value.createGenerator( identifierGeneratorFactory, dialect, rootClass, property ),
+							getType().isMutable() ? injector( property, getAttributeDeclarer( rootClass ) ) : null,
 							i
 					) );
 				}
 			}
 		}
 		return generator;
 	}
+	/**
+	 * Return the class that declares the composite pk attributes,
+	 * which might be an {@code @IdClass}, an {@code @EmbeddedId},
+	 * of the entity class itself.
+	 */
+	private Class<?> getAttributeDeclarer(RootClass rootClass) {
+		// See the javadoc discussion on CompositeNestedGeneratedValueGenerator
+		// for the various scenarios we need to account for here
+		if ( rootClass.getIdentifierMapper() != null ) {
+			// we have the @IdClass / <composite-id mapped="true"/> case
+			return resolveComponentClass();
+		}
+		else if ( rootClass.getIdentifierProperty() != null ) {
+			// we have the "@EmbeddedId" / <composite-id name="idName"/> case
+			return resolveComponentClass();
+		}
+		else {
+			// we have the "straight up" embedded (again the Hibernate term)
+			// component identifier: the entity class itself is the id class
+			return rootClass.getMappedClass();
+		}
+	}
 
 	private Setter injector(Property property, Class<?> attributeDeclarer) {
 		return property.getPropertyAccessStrategy( attributeDeclarer )
diff --git a/hibernate-core/src/main/java/org/hibernate/mapping/KeyValue.java b/hibernate-core/src/main/java/org/hibernate/mapping/KeyValue.java
@@ -34,6 +34,14 @@ Generator createGenerator(
 			Dialect dialect,
 			RootClass rootClass);
 
+	default Generator createGenerator(
+			IdentifierGeneratorFactory identifierGeneratorFactory,
+			Dialect dialect,
+			RootClass rootClass,
+			Property property) {
+		return createGenerator( identifierGeneratorFactory, dialect, rootClass );
+	}
+
 	/**
 	 * @deprecated Use {@link #createGenerator(IdentifierGeneratorFactory, Dialect, RootClass)} instead.
 	 *             No longer used except in legacy tests.
diff --git a/hibernate-core/src/main/java/org/hibernate/mapping/SimpleValue.java b/hibernate-core/src/main/java/org/hibernate/mapping/SimpleValue.java
@@ -392,10 +392,19 @@ public Generator createGenerator(
 			IdentifierGeneratorFactory identifierGeneratorFactory,
 			Dialect dialect,
 			RootClass rootClass) throws MappingException {
+		return createGenerator( identifierGeneratorFactory, dialect, rootClass, rootClass.getIdentifierProperty() );
+	}
+
+	@Override
+	public Generator createGenerator(
+			IdentifierGeneratorFactory identifierGeneratorFactory,
+			Dialect dialect,
+			RootClass rootClass,
+			Property property) throws MappingException {
 		if ( generator == null ) {
 			if ( customIdGeneratorCreator != null ) {
 				generator = customIdGeneratorCreator.createGenerator(
-						new IdGeneratorCreationContext( identifierGeneratorFactory, null, null, rootClass )
+						new IdGeneratorCreationContext( identifierGeneratorFactory, null, null, rootClass, property)
 				);
 			}
 			else {
@@ -1080,12 +1089,19 @@ private class IdGeneratorCreationContext implements CustomIdGeneratorCreationCon
 		private final String defaultCatalog;
 		private final String defaultSchema;
 		private final RootClass rootClass;
-
-		public IdGeneratorCreationContext(IdentifierGeneratorFactory identifierGeneratorFactory, String defaultCatalog, String defaultSchema, RootClass rootClass) {
+		private final Property property;
+
+		public IdGeneratorCreationContext(
+				IdentifierGeneratorFactory identifierGeneratorFactory,
+				String defaultCatalog,
+				String defaultSchema,
+				RootClass rootClass,
+				Property property) {
 			this.identifierGeneratorFactory = identifierGeneratorFactory;
 			this.defaultCatalog = defaultCatalog;
 			this.defaultSchema = defaultSchema;
 			this.rootClass = rootClass;
+			this.property = property;
 		}
 
 		@Override
@@ -1125,7 +1141,7 @@ public PersistentClass getPersistentClass() {
 
 		@Override
 		public Property getProperty() {
-			return rootClass.getIdentifierProperty();
+			return property;
 		}
 
 		// we could add these if it helps integrate old infrastructure
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/Account.java b/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/Account.java
@@ -0,0 +1,32 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.orm.test.tenantidpk;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+import jakarta.persistence.ManyToOne;
+import org.hibernate.annotations.TenantId;
+
+import java.util.UUID;
+
+@Entity
+public class Account {
+
+    @Id @GeneratedValue Long id;
+
+    @Id @TenantId UUID tenantId;
+
+    @ManyToOne(optional = false)
+    Client client;
+
+    public Account(Client client) {
+        this.client = client;
+    }
+
+    Account() {}
+}
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/Client.java b/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/Client.java
@@ -0,0 +1,39 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.orm.test.tenantidpk;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+import jakarta.persistence.OneToMany;
+import org.hibernate.annotations.TenantId;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.UUID;
+
+@Entity
+public class Client {
+    @Id
+    @GeneratedValue
+    Long id;
+
+    String name;
+
+    @Id @TenantId
+    UUID tenantId;
+
+    @OneToMany(mappedBy = "client")
+    Set<Account> accounts = new HashSet<>();
+
+    public Client(String name) {
+        this.name = name;
+    }
+
+    Client() {}
+}
+
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/TenantPkTest.java b/hibernate-core/src/test/java/org/hibernate/orm/test/tenantidpk/TenantPkTest.java
