suppport 'from new table' on DB2 for retrieving generated primary keys

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

@@ -45,8 +45,6 @@ import org.hibernate.metamodel.spi.RuntimeModelCreationContext;
 import org.hibernate.procedure.internal.DB2CallableStatementSupport;
 import org.hibernate.procedure.spi.CallableStatementSupport;
 import org.hibernate.query.spi.QueryEngine;
-import org.hibernate.procedure.internal.DB2CallableStatementSupport;
-import org.hibernate.procedure.spi.CallableStatementSupport;
 import org.hibernate.query.sqm.IntervalType;
 import org.hibernate.query.sqm.TemporalUnit;
 import org.hibernate.query.sqm.mutation.internal.cte.CteInsertStrategy;
@@ -842,6 +840,11 @@ public class DB2Dialect extends Dialect {
 		return new DB2IdentityColumnSupport();
 	}
 
+	@Override
+	public boolean supportsInsertReturning() {
+		return true;
+	}
+
 	@Override
 	public boolean supportsValuesList() {
 		return true;

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

@@ -39,6 +39,10 @@ import org.hibernate.sql.ast.tree.select.QuerySpec;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.ast.tree.update.UpdateStatement;
 import org.hibernate.sql.exec.spi.JdbcOperation;
+import org.hibernate.sql.model.internal.TableInsertStandard;
+
+import static org.hibernate.internal.util.collections.CollectionHelper.isEmpty;
+import static org.hibernate.internal.util.collections.CollectionHelper.isNotEmpty;
 
 /**
  * A SQL AST translator for DB2.
@@ -349,16 +353,15 @@ public class DB2SqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAst
 
 	protected boolean renderReturningClause(MutationStatement statement) {
 		final List<ColumnReference> returningColumns = statement.getReturningColumns();
-		final int size = returningColumns.size();
+		if ( isEmpty( returningColumns ) ) {
-		if ( size == 0 ) {
 			return false;
 		}
 		appendSql( "select " );
-		String separator = "";
+		for ( int i = 0; i < returningColumns.size(); i++ ) {
-		for ( int i = 0; i < size; i++ ) {
+			if ( i > 0 ) {
-			appendSql( separator );
+				appendSql( ", " );
+			}
 			appendSql( returningColumns.get( i ).getColumnExpression() );
-			separator = ",";
 		}
 		if ( statement instanceof DeleteStatement ) {
 			appendSql( " from old table (" );
@@ -369,6 +372,28 @@ public class DB2SqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAst
 		return true;
 	}
 
+	@Override
+	public void visitStandardTableInsert(TableInsertStandard tableInsert) {
+		final List<ColumnReference> returningColumns = tableInsert.getReturningColumns();
+		if ( isNotEmpty( returningColumns ) ) {
+			appendSql( "select " );
+
+			for ( int i = 0; i < returningColumns.size(); i++ ) {
+				if ( i > 0 ) {
+					appendSql( ", " );
+				}
+				appendSql( returningColumns.get( i ).getColumnExpression() );
+			}
+
+			appendSql( " from new table ( " ); // 'from final table' does not seem to play well with triggers
+			super.visitStandardTableInsert( tableInsert );
+			appendSql( ")" );
+		}
+		else {
+			super.visitStandardTableInsert( tableInsert );
+		}
+	}
+
 	@Override
 	protected void renderComparison(Expression lhs, ComparisonOperator operator, Expression rhs) {
 		if ( getDB2Version().isSameOrAfter( 11, 1 ) ) {

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

@@ -11,7 +11,6 @@ import java.util.List;
 import org.hibernate.LockMode;
 import org.hibernate.dialect.identity.H2IdentityColumnSupport;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
-import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.query.sqm.ComparisonOperator;
 import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.SqlAstNodeRenderingMode;
@@ -38,6 +37,8 @@ import org.hibernate.sql.ast.tree.select.SelectClause;
 import org.hibernate.sql.exec.spi.JdbcOperation;
 import org.hibernate.sql.model.internal.TableInsertStandard;
 
+import static org.hibernate.internal.util.collections.CollectionHelper.isNotEmpty;
+
 /**
  * A SQL AST translator for H2.
  *
@@ -53,7 +54,7 @@ public class H2SqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAstT
 
 	@Override
 	public void visitStandardTableInsert(TableInsertStandard tableInsert) {
-		if ( CollectionHelper.isNotEmpty( tableInsert.getReturningColumns() ) ) {
+		if ( isNotEmpty( tableInsert.getReturningColumns() ) ) {
 			visitReturningInsertStatement( tableInsert );
 		}
 		else {
@@ -238,7 +239,7 @@ public class H2SqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAstT
 		// i.e. `join ( (select ...) alias join ... )`, so we have to introduce a dummy table reference
 		if ( tableRef instanceof QueryPartTableReference || tableRef.getTableId().startsWith( "(select" ) ) {
 			final boolean realTableGroup = tableGroup.isRealTableGroup()
-					&& ( CollectionHelper.isNotEmpty( tableGroup.getTableReferenceJoins() )
+					&& ( isNotEmpty( tableGroup.getTableReferenceJoins() )
 					|| hasNestedTableGroupsToRender( tableGroup.getNestedTableGroupJoins() ) );
 			if ( realTableGroup ) {
 				appendSql( "dual cross join " );

hibernate-core/src/main/java/org/hibernate/dialect/identity/DB2IdentityColumnSupport.java

@@ -15,6 +15,11 @@ public class DB2IdentityColumnSupport extends IdentityColumnSupportImpl {
 		return true;
 	}
 
+	@Override
+	public boolean supportsInsertSelectIdentity() {
+		return true;
+	}
+
 	@Override
 	public String getIdentitySelectString(String table, String column, int type) {
 		return "values identity_val_local()";
@@ -29,4 +34,9 @@ public class DB2IdentityColumnSupport extends IdentityColumnSupportImpl {
 	public String getIdentityInsertString() {
 		return "default";
 	}
+
+	@Override
+	public String appendIdentitySelectToInsert(String identityColumnName, String insertString) {
+		return "select " + identityColumnName + " from final table ( " + insertString + " )"; //TODO: should it be 'from new table'?
+	}
 }