--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
@@ -4053,17 +4052,31 @@ public boolean supportsInsertReturning() {
 
 	/**
 	 * Does this dialect supports returning the {@link org.hibernate.annotations.RowId} column
-	 *  after execution of an {@code insert} statement, using native SQL syntax?
-	 *
+	 * after execution of an {@code insert} statement, using native SQL syntax?
 	 *
 	 * @return {@code true} is the dialect supports returning the rowid column
 	 *
 	 * @see #supportsInsertReturning()
+	 * @since 7.0
 	 */
 	public boolean supportsInsertReturningRowId() {
 		return supportsInsertReturning();
 	}
 
+	/**
+	 * Does this dialect fully support returning arbitrary generated column values
+	 * after execution of an {@code update} statement, using native SQL syntax?
+	 * <p>
+	 * Defaults to the value of {@link #supportsInsertReturning()} but can be overridden
+	 * to explicitly disable this for updates.
+	 *
+	 * @see #supportsInsertReturning()
+	 * @since 7.0
+	 */
+	public boolean supportsUpdateReturning() {
+		return supportsInsertReturning();
+	}
+
 	/**
 	 * Does this dialect fully support returning arbitrary generated column values
 	 * after execution of an {@code insert} statement, using the JDBC method

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

@@ -248,6 +248,11 @@ public boolean supportsInsertReturning() {
 		return getVersion().isSameOrAfter( 10, 5 );
 	}
 
+	@Override
+	public boolean supportsUpdateReturning() {
+		return false;
+	}
+
 	@Override
 	public FunctionalDependencyAnalysisSupport getFunctionalDependencyAnalysisSupport() {
 		return FunctionalDependencyAnalysisSupportImpl.TABLE_GROUP_AND_CONSTANTS;

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

@@ -32,7 +32,7 @@ public EventType getTiming() {
 	}
 
 	protected JdbcValuesMappingProducer getMappingProducer(Consumer<String> columnNameConsumer) {
-		return getMappingProducer( columnNameConsumer, false );
+		return getMappingProducer( columnNameConsumer, true );
 	}
 
 	protected JdbcValuesMappingProducer getMappingProducer(Consumer<String> columnNameConsumer, boolean useIndex) {

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

@@ -67,7 +67,7 @@ public BasicResult<?> buildResult(
 
 		final TableGroup tableGroup = creationStateImpl.getFromClauseAccess().resolveTableGroup(
 				navigablePath.getParent(),
-				(p) -> this.tableGroup
+				path -> this.tableGroup
 		);
 		final TableReference tableReference = tableGroup.resolveTableReference(
 				navigablePath,
@@ -76,7 +76,7 @@ public BasicResult<?> buildResult(
 		);
 
 		final int position = valuesArrayPosition == null ?
-				jdbcPositionToValuesArrayPosition( jdbcResultsMetadata.resolveColumnPosition( modelPart.getSelectionExpression() ) ) :
+				columnIndex( jdbcResultsMetadata, modelPart ) :
 				valuesArrayPosition;
 		final SqlSelection sqlSelection = creationStateImpl.resolveSqlSelection(
 				ResultsHelper.resolveSqlExpression(
@@ -96,4 +96,24 @@ public BasicResult<?> buildResult(
 				modelPart.getJdbcMapping()
 		);
 	}
+
+	private static int columnIndex(JdbcValuesMetadata jdbcResultsMetadata, BasicValuedModelPart modelPart) {
+		final BasicValuedModelPart resultModelPart = modelPart.isEntityIdentifierMapping() ?
+				( (BasicValuedModelPart) modelPart.findContainingEntityMapping()
+						.getRootEntityDescriptor()
+						.getIdentifierMapping() )
+				: modelPart;
+		try {
+			return jdbcPositionToValuesArrayPosition( jdbcResultsMetadata.resolveColumnPosition( resultModelPart.getSelectionExpression() ) );
+		}
+		catch (Exception e) {
+			if ( modelPart.isEntityIdentifierMapping() ) {
+				// Default to the first position for entity identifiers
+				return 0;
+			}
+			else {
+				throw e;
+			}
+		}
+	}
 }

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

@@ -335,7 +335,7 @@ && noCustomSql( persister, timing ) ) {
 		if ( dialect.supportsInsertReturningGeneratedKeys() ) {
 			return new GetGeneratedKeysDelegate( persister, dialect, false, timing );
 		}
-		else if ( dialect.supportsInsertReturning() && noCustomSql( persister, timing ) ) {
+		else if ( supportsReturning( dialect, timing ) && noCustomSql( persister, timing ) ) {
 			return new InsertReturningDelegate( persister, dialect, timing );
 		}
 		else if ( timing == EventType.INSERT && persister.getNaturalIdentifierProperties() != null
@@ -350,6 +350,10 @@ else if ( timing == EventType.INSERT && persister.getNaturalIdentifierProperties
 		return null;
 	}
 
+	private static boolean supportsReturning(Dialect dialect, EventType timing) {
+		return timing == EventType.INSERT ? dialect.supportsInsertReturning() : dialect.supportsUpdateReturning();
+	}
+
 	public static boolean noCustomSql(EntityPersister persister, EventType timing) {
 		final EntityTableMapping identifierTable = persister.getIdentifierTableMapping();
 		final TableMapping.MutationDetails mutationDetails = timing == EventType.INSERT ?

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 );
 	}