HHH-18642 Enable getGeneratedKeys() for arbitrary attributes on DB2

Also revert to Statement.RETURN_GENERATED_KEYS when selecting identity only

Author: mbellade

hibernate-community-dialects/src/main/java/org/hibernate/community/dialect/DB2LegacyDialect.java

@@ -1060,6 +1060,24 @@ public IdentityColumnSupport getIdentityColumnSupport() {
 		return DB2IdentityColumnSupport.INSTANCE;
 	}
 
+	@Override
+	public boolean supportsInsertReturningGeneratedKeys() {
+		return true;
+	}
+
+	/**
+	 * @return {@code true} because we can use {@code select ... from new table (insert .... )}
+	 */
+	@Override
+	public boolean supportsInsertReturning() {
+		return true;
+	}
+
+	@Override
+	public boolean supportsInsertReturningRowId() {
+		return false;
+	}
+
 	@Override
 	public boolean supportsValuesList() {
 		return true;

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

@@ -1127,6 +1127,11 @@ public IdentityColumnSupport getIdentityColumnSupport() {
 		return DB2IdentityColumnSupport.INSTANCE;
 	}
 
+	@Override
+	public boolean supportsInsertReturningGeneratedKeys() {
+		return true;
+	}
+
 	/**
 	 * @return {@code true} because we can use {@code select ... from new table (insert .... )}
 	 */

hibernate-core/src/main/java/org/hibernate/id/IdentityGenerator.java

@@ -4,15 +4,17 @@
  */
 package org.hibernate.id;
 
+import java.util.List;
+
 import org.hibernate.boot.spi.SessionFactoryOptions;
 import org.hibernate.dialect.Dialect;
-import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.generator.OnExecutionGenerator;
 import org.hibernate.id.insert.BasicSelectingDelegate;
 import org.hibernate.id.insert.GetGeneratedKeysDelegate;
 import org.hibernate.id.insert.InsertGeneratedIdentifierDelegate;
 import org.hibernate.id.insert.InsertReturningDelegate;
 import org.hibernate.id.insert.UniqueKeySelectingDelegate;
+import org.hibernate.metamodel.mapping.ModelPart;
 import org.hibernate.persister.entity.EntityPersister;
 
 import static org.hibernate.generator.EventType.INSERT;
@@ -52,11 +54,16 @@ public String[] getReferencedColumnValues(Dialect dialect) {
 
 	@Override
 	public InsertGeneratedIdentifierDelegate getGeneratedIdentifierDelegate(EntityPersister persister) {
-		final SessionFactoryImplementor factory = persister.getFactory();
-		final Dialect dialect = factory.getJdbcServices().getDialect();
+		final Dialect dialect = persister.getFactory().getJdbcServices().getDialect();
+		final SessionFactoryOptions sessionFactoryOptions = persister.getFactory().getSessionFactoryOptions();
+		final List<? extends ModelPart> generatedProperties = persister.getGeneratedProperties( INSERT );
+		if ( generatedProperties.size() == 1 && sessionFactoryOptions.isGetGeneratedKeysEnabled() ) {
+			// Use Connection#prepareStatement(sql, Statement.RETURN_GENERATED_KEYS) when only retrieving identity
+			assert generatedProperties.get( 0 ).isEntityIdentifierMapping();
+			return dialect.getIdentityColumnSupport().buildGetGeneratedKeysDelegate( persister );
+		}
 		// Try to use generic delegates if the dialects supports them
-		final SessionFactoryOptions sessionFactoryOptions = factory.getSessionFactoryOptions();
-		if ( dialect.supportsInsertReturningGeneratedKeys() && sessionFactoryOptions.isGetGeneratedKeysEnabled() ) {
+		else if ( dialect.supportsInsertReturningGeneratedKeys() && sessionFactoryOptions.isGetGeneratedKeysEnabled() ) {
 			return new GetGeneratedKeysDelegate( persister, false, INSERT );
 		}
 		else if ( dialect.supportsInsertReturning() && noCustomSql( persister, INSERT ) ) {