From 810842ac74ae6e58df9a2cc719cb078e393d6fa9 Mon Sep 17 00:00:00 2001
From: Marco Belladelli <marcobladel@gmail.com>
Date: Thu, 10 Oct 2024 12:06:11 +0200
Subject: [PATCH] HHH-18679 Allow `@Generated(writable=true)` with assigned
 identifiers

---
 .../internal/GeneratedGeneration.java         | 33 ++++++++++++++++++-
 .../mutation/InsertCoordinatorStandard.java   | 13 +++++---
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/generator/internal/GeneratedGeneration.java b/hibernate-core/src/main/java/org/hibernate/generator/internal/GeneratedGeneration.java
index 03a456ec76..2fc155401c 100644
--- a/hibernate-core/src/main/java/org/hibernate/generator/internal/GeneratedGeneration.java
+++ b/hibernate-core/src/main/java/org/hibernate/generator/internal/GeneratedGeneration.java
@@ -6,8 +6,11 @@ package org.hibernate.generator.internal;
 
 import org.hibernate.annotations.Generated;
 import org.hibernate.dialect.Dialect;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.generator.BeforeExecutionGenerator;
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.OnExecutionGenerator;
+import org.hibernate.persister.entity.EntityPersister;
 
 import java.util.EnumSet;
 
@@ -23,7 +26,7 @@ import static org.hibernate.internal.util.StringHelper.isEmpty;
  * @author Steve Ebersole
  * @author Gunnar Morling
  */
-public class GeneratedGeneration implements OnExecutionGenerator {
+public class GeneratedGeneration implements OnExecutionGenerator, BeforeExecutionGenerator {
 
 	private final EnumSet<EventType> eventTypes;
 	private final boolean writable;
@@ -60,4 +63,32 @@ public class GeneratedGeneration implements OnExecutionGenerator {
 	public boolean writePropertyValue() {
 		return writable && sql==null;
 	}
+
+	@Override
+	public boolean generatedOnExecution() {
+		return true;
+	}
+
+	@Override
+	public boolean generatedOnExecution(Object entity, SharedSessionContractImplementor session) {
+		if ( !writable ) {
+			return true;
+		}
+
+		// When this is the identifier generator and writable is true, allow pre-assigned identifiers
+		final EntityPersister entityPersister = session.getEntityPersister( null, entity );
+		return entityPersister.getGenerator() != this || entityPersister.getIdentifier( entity, session ) == null;
+	}
+
+	@Override
+	public Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {
+		final EntityPersister entityPersister = session.getEntityPersister( null, owner );
+		assert entityPersister.getGenerator() == this;
+		return entityPersister.getIdentifier( owner, session );
+	}
+
+	@Override
+	public boolean allowAssignedIdentifiers() {
+		return writable;
+	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/InsertCoordinatorStandard.java b/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/InsertCoordinatorStandard.java
index ddd8432d00..80ee294942 100644
--- a/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/InsertCoordinatorStandard.java
+++ b/hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/InsertCoordinatorStandard.java
@@ -24,6 +24,7 @@ import org.hibernate.generator.Generator;
 import org.hibernate.generator.OnExecutionGenerator;
 import org.hibernate.generator.values.GeneratedValues;
 import org.hibernate.generator.values.GeneratedValuesMutationDelegate;
+import org.hibernate.id.IdentifierGenerationException;
 import org.hibernate.metamodel.mapping.AttributeMapping;
 import org.hibernate.metamodel.mapping.AttributeMappingsList;
 import org.hibernate.metamodel.mapping.BasicEntityIdentifierMapping;
@@ -433,11 +434,13 @@ public class InsertCoordinatorStandard extends AbstractMutationCoordinator imple
 				if ( generator.referenceColumnsInSql( dialect ) ) {
 					final BasicEntityIdentifierMapping identifierMapping = (BasicEntityIdentifierMapping) entityPersister().getIdentifierMapping();
 					final String[] columnValues = generator.getReferencedColumnValues( dialect );
-					tableMapping.getKeyMapping().forEachKeyColumn( (i, column) -> tableInsertBuilder.addKeyColumn(
-							column.getColumnName(),
-							columnValues[i],
-							identifierMapping.getJdbcMapping()
-					) );
+					if ( columnValues != null ) {
+						tableMapping.getKeyMapping().forEachKeyColumn( (i, column) -> tableInsertBuilder.addKeyColumn(
+								column.getColumnName(),
+								columnValues[i],
+								identifierMapping.getJdbcMapping()
+						) );
+					}
 				}
 			}
 			else {