HHH-19004 Fix regression of TenantId not being allowed on @EmbeddedId

hibernate-core/src/main/java/org/hibernate/annotations/TenantId.java

@@ -26,6 +26,7 @@
  * @author Gavin King
  */
 @ValueGenerationType(generatedBy = TenantIdGeneration.class)
+@IdGeneratorType(TenantIdGeneration.class)
 @AttributeBinderType(binder = TenantIdBinder.class)
 @Target({METHOD, FIELD})
 @Retention(RUNTIME)

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 ) {

hibernate-core/src/main/java/org/hibernate/generator/internal/TenantIdGeneration.java

@@ -18,6 +18,7 @@
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.EventTypeSets;
 import org.hibernate.generator.GeneratorCreationContext;
+import org.hibernate.id.factory.spi.CustomIdGeneratorCreationContext;
 import org.hibernate.type.descriptor.java.JavaType;
 
 import static org.hibernate.generator.EventTypeSets.INSERT_ONLY;
@@ -33,6 +34,10 @@ public class TenantIdGeneration implements BeforeExecutionGenerator {
 	private final String entityName;
 	private final String propertyName;
 
+	public TenantIdGeneration(TenantId annotation, Member member, CustomIdGeneratorCreationContext context) {
+		this(annotation, member, (GeneratorCreationContext) context);
+	}
+
 	public TenantIdGeneration(TenantId annotation, Member member, GeneratorCreationContext context) {
 		entityName = context.getPersistentClass() == null
 				? member.getDeclaringClass().getName() //it's an attribute of an embeddable

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

@@ -463,11 +463,12 @@
 		bootMetamodel.getEntityBindings().stream()
 				.filter( model -> !model.isInherited() )
 				.forEach( model -> {
 					final KeyValue id = model.getIdentifier();
 					final Generator generator = id.createGenerator(
 							bootstrapContext.getIdentifierGeneratorFactory(),
 							jdbcServices.getJdbcEnvironment().getDialect(),
-							(RootClass) model
+							(RootClass) model,
+							model == null ? null : model.getIdentifierProperty()
 					);
 					if ( generator instanceof Configurable ) {
 						final Configurable identifierGenerator = (Configurable) generator;

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 == null ? null : 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 )

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.

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 == null ? null : 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

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

@@ -0,0 +1,120 @@
+/*
+ * 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.tenantid;
+
+import org.hibernate.annotations.TenantId;
+import org.hibernate.boot.SessionFactoryBuilder;
+import org.hibernate.boot.spi.MetadataImplementor;
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.Jira;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryProducer;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Embeddable;
+import jakarta.persistence.EmbeddedId;
+import jakarta.persistence.Entity;
+import jakarta.persistence.EnumType;
+import jakarta.persistence.Enumerated;
+
+import static org.hibernate.cfg.SchemaToolingSettings.JAKARTA_HBM2DDL_DATABASE_ACTION;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+@SessionFactory
+@DomainModel(annotatedClasses = { TenantIdGeneratedTest.DisplayIdBE.class })
+@ServiceRegistry(
+        settings = {
+                @Setting(name = JAKARTA_HBM2DDL_DATABASE_ACTION, value = "create-drop")
+        }
+)
+@Jira("https://hibernate.atlassian.net/browse/HHH-19004")
+public class TenantIdGeneratedTest implements SessionFactoryProducer {
+
+    @AfterEach
+    public void cleanup(SessionFactoryScope scope) {
+        scope.getSessionFactory().getSchemaManager().truncateMappedObjects();
+    }
+
+    @Override
+    public SessionFactoryImplementor produceSessionFactory(MetadataImplementor model) {
+        final SessionFactoryBuilder sessionFactoryBuilder = model.getSessionFactoryBuilder();
+        sessionFactoryBuilder.applyCurrentTenantIdentifierResolver( new CurrentTenantIdentifierResolver<Long>() {
+            @Override
+            public Long resolveCurrentTenantIdentifier() {
+                return 0L;
+            }
+            @Override
+            public boolean validateExistingCurrentSessions() {
+                return false;
+            }
+        } );
+        return (SessionFactoryImplementor) sessionFactoryBuilder.build();
+    }
+
+    @Test
+    public void test(SessionFactoryScope scope) {
+        scope.inTransaction( session -> {
+            session.persist(new DisplayIdBE( new DisplayIdKeyBE( null, DisplayIdType.TYPE1 ), 1L ));
+        } );
+        scope.inTransaction( session -> {
+            assertNotNull( session.find( DisplayIdBE.class, new DisplayIdKeyBE( 0L, DisplayIdType.TYPE1 ) ) );
+            assertEquals( 1, session.createQuery("from DisplayIdBE", DisplayIdBE.class).getResultList().size() );
+        } );
+    }
+
+    @Entity(name = "DisplayIdBE")
+    public static class DisplayIdBE {
+
+        @EmbeddedId
+        private DisplayIdKeyBE id;
+
+        @Column(name = "display_id_value", nullable = false)
+        private long displayIdValue;
+
+        protected DisplayIdBE() {
+        }
+
+        public DisplayIdBE(DisplayIdKeyBE id, long displayIdValue) {
+            this.id = id;
+            this.displayIdValue = displayIdValue;
+        }
+    }
+
+    public enum DisplayIdType {
+        TYPE1,
+        TYPE2
+    }
+
+    @Embeddable
+    public static class DisplayIdKeyBE {
+
+        @TenantId
+        @Column(name = "tenant_id", nullable = false)
+        private Long tenantId;
+
+        @Column(name = "type", nullable = false)
+        @Enumerated(EnumType.STRING)
+        private DisplayIdType type;
+
+        protected DisplayIdKeyBE() {}
+
+        public DisplayIdKeyBE(Long tenantId, DisplayIdType type) {
+            this.tenantId = tenantId;
+            this.type = type;
+        }
+    }
+
+}

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() {}
+}