--wip-- last tweaks and fixes

hibernate-core/src/main/java/org/hibernate/action/internal/AbstractEntityInsertAction.java

@@ -146,6 +146,10 @@ public final void makeEntityManaged() {
 		);
 		if ( isEarlyInsert() ) {
 			addCollectionsByKeyToPersistenceContext( persistenceContextInternal, getState() );
+			final Object rowId = getRowId();
+			if ( rowId != null ) {
+				persistenceContextInternal.replaceEntityEntryRowId( getInstance(), rowId );
+			}
 		}
 	}
 
@@ -230,6 +234,8 @@ protected void markExecuted() {
 	 */
 	protected abstract EntityKey getEntityKey();
 
+	protected abstract Object getRowId();
+
 	@Override
 	public void afterDeserialize(EventSource session) {
 		super.afterDeserialize( session );

hibernate-core/src/main/java/org/hibernate/action/internal/EntityIdentityInsertAction.java

@@ -8,9 +8,11 @@
 
 import org.hibernate.AssertionFailure;
 import org.hibernate.HibernateException;
+import org.hibernate.LockMode;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.engine.spi.Status;
 import org.hibernate.event.service.spi.EventListenerGroup;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.event.spi.PostCommitInsertEventListener;
@@ -22,6 +24,8 @@
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.stat.spi.StatisticsImplementor;
 
+import static org.hibernate.engine.internal.Versioning.getVersion;
+
 /**
  * The action for performing entity insertions when entity is using IDENTITY column identifier generation
  *
@@ -33,6 +37,7 @@ public class EntityIdentityInsertAction extends AbstractEntityInsertAction  {
 	private final EntityKey delayedEntityKey;
 	private EntityKey entityKey;
 	private Object generatedId;
+	private Object rowId;
 
 	/**
 	 * Constructs an EntityIdentityInsertAction
@@ -84,8 +89,8 @@ public void execute() throws HibernateException {
 			generatedId = generatedValues.getGeneratedValue( persister.getIdentifierMapping() );
 			// Process row-id values when available early by replacing the entity entry
 			if ( persister.getRowIdMapping() != null ) {
-				final Object rowId = generatedValues.getGeneratedValue( persister.getRowIdMapping() );
-				if ( rowId != null ) {
+				rowId = generatedValues.getGeneratedValue( persister.getRowIdMapping() );
+				if ( rowId != null && isDelayed ) {
 					persistenceContext.replaceEntityEntryRowId( getInstance(), rowId );
 				}
 			}
@@ -221,6 +226,11 @@ protected EntityKey getEntityKey() {
 		return entityKey != null ? entityKey : delayedEntityKey;
 	}
 
+	@Override
+	public Object getRowId() {
+		return rowId;
+	}
+
 	protected void setEntityKey(EntityKey entityKey) {
 		this.entityKey = entityKey;
 	}

hibernate-core/src/main/java/org/hibernate/action/internal/EntityInsertAction.java

@@ -84,6 +84,11 @@ public boolean isEarlyInsert() {
 		return false;
 	}
 
+	@Override
+	protected Object getRowId() {
+		return null;
+	}
+
 	@Override
 	protected EntityKey getEntityKey() {
 		return getSession().generateEntityKey( getId(), getPersister() );

hibernate-core/src/main/java/org/hibernate/dialect/Dialect.java

@@ -4042,7 +4042,6 @@ public boolean supportsSubqueryInSelect() {
 	 *         {@code false} if {@code InsertReturningDelegate} does not work, or only
 	 *         works for specialized identity/"autoincrement" columns
 	 *
-	 * @see org.hibernate.generator.OnExecutionGenerator#getGeneratedIdentifierDelegate
 	 * @see org.hibernate.id.insert.InsertReturningDelegate
 	 *
 	 * @since 6.2

hibernate-core/src/main/java/org/hibernate/generator/values/GeneratedValueBasicResultBuilder.java

@@ -76,7 +76,7 @@ public BasicResult<?> buildResult(
 		);
 
 		final int position = valuesArrayPosition == null ?
-				jdbcPositionToValuesArrayPosition( jdbcResultsMetadata.resolveColumnPosition( modelPart.getSelectionExpression() ) ) :
+				jdbcPositionToValuesArrayPosition( jdbcResultsMetadata.resolveColumnPosition( getColumnName( modelPart ) ) ) :
 				valuesArrayPosition;
 		final SqlSelection sqlSelection = creationStateImpl.resolveSqlSelection(
 				ResultsHelper.resolveSqlExpression(
@@ -96,4 +96,13 @@ public BasicResult<?> buildResult(
 				modelPart.getJdbcMapping()
 		);
 	}
+
+	private static String getColumnName(BasicValuedModelPart modelPart) {
+		return modelPart.isEntityIdentifierMapping() ?
+				( (BasicValuedModelPart) modelPart.findContainingEntityMapping()
+						.getRootEntityDescriptor()
+						.getIdentifierMapping() )
+						.getSelectionExpression()
+				: modelPart.getSelectionExpression();
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/generated/delegate/MutationDelegateIdentityTest.java

@@ -20,6 +20,7 @@
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.values.GeneratedValuesMutationDelegate;
 import org.hibernate.id.insert.UniqueKeySelectingDelegate;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.persister.entity.mutation.EntityMutationTarget;
 import org.hibernate.sql.model.MutationType;
 
@@ -224,8 +225,9 @@ private static GeneratedValuesMutationDelegate getDelegate(
 			SessionFactoryScope scope,
 			Class<?> entityClass,
 			MutationType mutationType) {
-		final EntityMutationTarget entityDescriptor = (EntityMutationTarget) scope.getSessionFactory()
-				.getMappingMetamodel().findEntityDescriptor( entityClass );
+		final EntityPersister entityDescriptor = scope.getSessionFactory()
+				.getMappingMetamodel()
+				.findEntityDescriptor( entityClass );
 		return entityDescriptor.getMutationDelegate( mutationType );
 	}
 

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/generated/delegate/MutationDelegateJoinedInheritanceTest.java

@@ -14,6 +14,7 @@
 import org.hibernate.annotations.UpdateTimestamp;
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.values.GeneratedValuesMutationDelegate;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.persister.entity.mutation.EntityMutationTarget;
 import org.hibernate.sql.model.MutationType;
 
@@ -161,10 +162,11 @@ public void testUpdateChildEntity(SessionFactoryScope scope) {
 
 	private static GeneratedValuesMutationDelegate getDelegate(
 			SessionFactoryScope scope,
-			Class<?> entityClass,
+			@SuppressWarnings( "SameParameterValue" ) Class<?> entityClass,
 			MutationType mutationType) {
-		final EntityMutationTarget entityDescriptor = (EntityMutationTarget) scope.getSessionFactory()
-				.getMappingMetamodel().findEntityDescriptor( entityClass );
+		final EntityPersister entityDescriptor = scope.getSessionFactory()
+				.getMappingMetamodel()
+				.findEntityDescriptor( entityClass );
 		return entityDescriptor.getMutationDelegate( mutationType );
 	}
 

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/generated/delegate/MutationDelegateTest.java

@@ -20,6 +20,7 @@
 import org.hibernate.generator.EventType;
 import org.hibernate.generator.values.GeneratedValuesMutationDelegate;
 import org.hibernate.id.insert.UniqueKeySelectingDelegate;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.persister.entity.mutation.EntityMutationTarget;
 import org.hibernate.sql.model.MutationType;
 
@@ -174,8 +175,9 @@ private static GeneratedValuesMutationDelegate getDelegate(
 			SessionFactoryScope scope,
 			Class<?> entityClass,
 			MutationType mutationType) {
-		final EntityMutationTarget entityDescriptor = (EntityMutationTarget) scope.getSessionFactory()
-				.getMappingMetamodel().findEntityDescriptor( entityClass );
+		final EntityPersister entityDescriptor = scope.getSessionFactory()
+				.getMappingMetamodel()
+				.findEntityDescriptor( entityClass );
 		return entityDescriptor.getMutationDelegate( mutationType );
 	}
 

hibernate-core/src/test/java/org/hibernate/orm/test/rowid/RowIdUpdateAndDeleteTest.java

@@ -8,6 +8,11 @@
 
 import org.hibernate.annotations.RowId;
 import org.hibernate.dialect.Dialect;
+import org.hibernate.generator.values.GeneratedValuesMutationDelegate;
+import org.hibernate.orm.test.mapping.generated.delegate.MutationDelegateTest;
+import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.persister.entity.mutation.EntityMutationTarget;
+import org.hibernate.sql.model.MutationType;
 
 import org.hibernate.testing.jdbc.SQLStatementInspector;
 import org.hibernate.testing.orm.junit.DomainModel;
@@ -67,8 +72,8 @@ public void testSimpleUpdateSameTransaction(SessionFactoryScope scope) {
 			simpleEntity.setStatus( "new_status" );
 			inspector.clear();
 		} );
-		// the update should have used the primary key, as the row-id value is not available
-		checkUpdateQuery( inspector, true );
+		// the update should have used the primary key when the row-id value is not available
+		checkUpdateQuery( inspector, scope, true );
 		scope.inTransaction( session -> assertThat(
 				session.find( SimpleEntity.class, 3L ).getStatus()
 		).isEqualTo( "new_status" ) );
@@ -85,8 +90,8 @@ public void testSimpleDeleteSameTransaction(SessionFactoryScope scope) {
 			session.remove( simpleEntity );
 			inspector.clear();
 		} );
-		// the update should have used the primary key, as the row-id value is not available
-		checkUpdateQuery( inspector, true );
+		// the update should have used the primary key when the row-id value is not available
+		checkUpdateQuery( inspector, scope, true );
 		scope.inTransaction( session -> assertThat( session.find( SimpleEntity.class, 13L ) ).isNull() );
 	}
 
@@ -107,8 +112,8 @@ public void testRelatedUpdateSameTransaction(SessionFactoryScope scope) {
 			parent.getChild().setStatus( "new_status" );
 			inspector.clear();
 		} );
-		// the update should have used the primary key, as the row-id value is not available
-		checkUpdateQuery( inspector, true );
+		// the update should have used the primary key when the row-id value is not available
+		checkUpdateQuery( inspector, scope, true );
 		scope.inTransaction( session -> assertThat(
 				session.find( SimpleEntity.class, 4L ).getStatus()
 		).isEqualTo( "new_status" ) );
@@ -122,8 +127,7 @@ public void testSimpleUpdateDifferentTransaction(SessionFactoryScope scope) {
 			simpleEntity.setStatus( "new_status" );
 			inspector.clear();
 		} );
-		final Dialect dialect = scope.getSessionFactory().getJdbcServices().getDialect();
-		checkUpdateQuery( inspector, dialect.rowId( "" ) == null );
+		checkUpdateQuery( inspector, scope, false );
 		scope.inTransaction( session -> assertThat(
 				session.find( SimpleEntity.class, 1L ).getStatus()
 		).isEqualTo( "new_status" ) );
@@ -137,8 +141,7 @@ public void testSimpleDeleteDifferentTransaction(SessionFactoryScope scope) {
 			session.remove( simpleEntity );
 			inspector.clear();
 		} );
-		final Dialect dialect = scope.getSessionFactory().getJdbcServices().getDialect();
-		checkUpdateQuery( inspector, dialect.rowId( "" ) == null );
+		checkUpdateQuery( inspector, scope, false );
 		scope.inTransaction( session -> assertThat( session.find( SimpleEntity.class, 11L ) ).isNull() );
 	}
 
@@ -150,14 +153,30 @@ public void testRelatedUpdateRelatedDifferentTransaction(SessionFactoryScope sco
 			parent.getChild().setStatus( "new_status" );
 			inspector.clear();
 		} );
-		final Dialect dialect = scope.getSessionFactory().getJdbcServices().getDialect();
-		checkUpdateQuery( inspector, dialect.rowId( "" ) == null );
+		checkUpdateQuery( inspector, scope, false );
 		scope.inTransaction( session -> assertThat(
 				session.find( SimpleEntity.class, 2L ).getStatus()
 		).isEqualTo( "new_status" ) );
 	}
 
-	private void checkUpdateQuery(SQLStatementInspector inspector, boolean shouldUsePrimaryKey) {
+	private void checkUpdateQuery(SQLStatementInspector inspector, SessionFactoryScope scope, boolean sameTransaction) {
+		final Dialect dialect = scope.getSessionFactory().getJdbcServices().getDialect();
+		final boolean shouldUsePrimaryKey;
+		if ( dialect.rowId( "" ) == null ) {
+			shouldUsePrimaryKey = true;
+		}
+		else {
+			if ( sameTransaction ) {
+				final EntityPersister persister = scope.getSessionFactory()
+						.getMappingMetamodel()
+						.findEntityDescriptor( SimpleEntity.class );
+				final GeneratedValuesMutationDelegate delegate = persister.getMutationDelegate( MutationType.INSERT );
+				shouldUsePrimaryKey = delegate == null || !delegate.supportsRowId();
+			}
+			else {
+				shouldUsePrimaryKey = false;
+			}
+		}
 		inspector.assertNumberOfOccurrenceInQueryNoSpace( 0, "primary_key", shouldUsePrimaryKey ? 1 : 0 );
 	}