Add executeWithKey() to JPAInsertClause and HibernateInsertClause

Closes #1692

Author: 이승영

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/JpaInsertNativeHelper.java

@@ -0,0 +1,186 @@
+/*
+ * Copyright 2015, The Querydsl Team (http://www.querydsl.com/team)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.querydsl.jpa;
+
+import com.querydsl.core.types.Expression;
+import com.querydsl.core.types.ParamExpression;
+import com.querydsl.core.types.ParamNotSetException;
+import com.querydsl.core.types.Path;
+import com.querydsl.core.types.dsl.Param;
+import jakarta.persistence.Column;
+import jakarta.persistence.Table;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Helper for building native SQL INSERT statements from JPA entity metadata. Used by {@link
+ * com.querydsl.jpa.impl.JPAInsertClause} and {@link
+ * com.querydsl.jpa.hibernate.HibernateInsertClause} to support {@code executeWithKey()}.
+ *
+ * <p>This is an internal API and not intended for direct use by application code.
+ */
+public final class JpaInsertNativeHelper {
+
+  private JpaInsertNativeHelper() {}
+
+  /**
+   * Resolve the SQL table name for an entity class.
+   *
+   * @param entityClass the JPA entity class
+   * @return the SQL table name
+   */
+  public static String resolveTableName(Class<?> entityClass) {
+    if (entityClass.isAnnotationPresent(Table.class)) {
+      var table = entityClass.getAnnotation(Table.class);
+      if (!table.name().isEmpty()) {
+        var sb = new StringBuilder();
+        if (!table.schema().isEmpty()) {
+          sb.append(table.schema()).append('.');
+        }
+        sb.append(table.name());
+        return sb.toString();
+      }
+    }
+    return entityClass.getSimpleName();
+  }
+
+  /**
+   * Resolve the SQL column name for a path. Reads {@code @Column} annotation if present, otherwise
+   * falls back to the path metadata name.
+   *
+   * @param path the query path
+   * @return the SQL column name
+   */
+  public static String resolveColumnName(Path<?> path) {
+    if (path.getAnnotatedElement() != null
+        && path.getAnnotatedElement().isAnnotationPresent(Column.class)) {
+      var column = path.getAnnotatedElement().getAnnotation(Column.class);
+      if (!column.name().isEmpty()) {
+        return column.name();
+      }
+    }
+    return path.getMetadata().getName();
+  }
+
+  /**
+   * Build a native SQL INSERT statement from entity metadata and column paths.
+   *
+   * @param entityClass the entity class (for table name resolution)
+   * @param columns the columns to insert
+   * @return the native SQL INSERT string with positional parameters
+   */
+  public static String buildNativeInsertSQL(Class<?> entityClass, Collection<Path<?>> columns) {
+    var tableName = resolveTableName(entityClass);
+    var sb = new StringBuilder();
+    sb.append("INSERT INTO ").append(tableName).append(" (");
+
+    var first = true;
+    for (Path<?> col : columns) {
+      if (!first) {
+        sb.append(", ");
+      }
+      sb.append(resolveColumnName(col));
+      first = false;
+    }
+
+    sb.append(") VALUES (");
+    first = true;
+    for (int i = 0; i < columns.size(); i++) {
+      if (!first) {
+        sb.append(", ");
+      }
+      sb.append('?');
+      first = false;
+    }
+    sb.append(')');
+
+    return sb.toString();
+  }
+
+  /**
+   * Resolve constant values from the serializer, unwrapping {@link Param} expressions.
+   *
+   * @param constants the constants from the serializer
+   * @param params the parameter bindings
+   * @return resolved values ready for JDBC binding
+   */
+  public static Object[] resolveConstants(
+      List<Object> constants, Map<ParamExpression<?>, Object> params) {
+    var result = new Object[constants.size()];
+    for (var i = 0; i < constants.size(); i++) {
+      var val = constants.get(i);
+      if (val instanceof Param<?> param) {
+        val = params.get(val);
+        if (val == null) {
+          throw new ParamNotSetException(param);
+        }
+      }
+      result[i] = val;
+    }
+    return result;
+  }
+
+  /**
+   * Execute a native SQL INSERT with RETURN_GENERATED_KEYS and return the generated key.
+   *
+   * @param <T> the key type
+   * @param conn the JDBC connection (not closed by this method)
+   * @param sql the native SQL INSERT string
+   * @param params the parameter values to bind
+   * @param keyType the expected key type
+   * @return the generated key, or null if no rows were inserted
+   * @throws SQLException if a database error occurs
+   */
+  @Nullable
+  public static <T> T executeAndReturnKey(
+      java.sql.Connection conn, String sql, Object[] params, Class<T> keyType) throws SQLException {
+    try (PreparedStatement stmt = conn.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS)) {
+      for (int i = 0; i < params.length; i++) {
+        stmt.setObject(i + 1, params[i]);
+      }
+      stmt.executeUpdate();
+
+      try (ResultSet rs = stmt.getGeneratedKeys()) {
+        if (rs.next()) {
+          return rs.getObject(1, keyType);
+        }
+        return null;
+      }
+    }
+  }
+
+  /**
+   * Collect effective columns and values from either the set-style inserts map or the
+   * columns/values lists.
+   *
+   * @param inserts the set-style inserts (path to expression mapping)
+   * @param columns the columns list (from columns().values() style)
+   * @param values the values list
+   * @param serializer used to extract constant values from expressions
+   * @return the effective column paths
+   */
+  public static Collection<Path<?>> effectiveColumns(
+      Map<Path<?>, Expression<?>> inserts, List<Path<?>> columns) {
+    if (!inserts.isEmpty()) {
+      return inserts.keySet();
+    }
+    return columns;
+  }
+}

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/hibernate/DefaultSessionHolder.java

@@ -14,6 +14,7 @@
 package com.querydsl.jpa.hibernate;
 
 import org.hibernate.Session;
+import org.hibernate.jdbc.ReturningWork;
 import org.hibernate.query.NativeQuery;
 import org.hibernate.query.Query;
 
@@ -39,4 +40,9 @@ public Query<?> createQuery(String queryString) {
   public NativeQuery<?> createSQLQuery(String queryString) {
     return session.createNativeQuery(queryString);
   }
+
+  @Override
+  public <T> T doReturningWork(ReturningWork<T> work) {
+    return session.doReturningWork(work);
+  }
 }

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/hibernate/HibernateInsertClause.java

@@ -14,6 +14,7 @@
 package com.querydsl.jpa.hibernate;
 
 import com.querydsl.core.JoinType;
+import com.querydsl.core.QueryException;
 import com.querydsl.core.dml.InsertClause;
 import com.querydsl.core.support.QueryMixin;
 import com.querydsl.core.types.EntityPath;
@@ -25,6 +26,8 @@
 import com.querydsl.jpa.JPAQueryMixin;
 import com.querydsl.jpa.JPQLSerializer;
 import com.querydsl.jpa.JPQLTemplates;
+import com.querydsl.jpa.JpaInsertNativeHelper;
+import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -35,6 +38,7 @@
 import org.hibernate.Session;
 import org.hibernate.StatelessSession;
 import org.hibernate.query.Query;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * UpdateClause implementation for Hibernate
@@ -97,6 +101,70 @@ public long execute() {
     return query.executeUpdate();
   }
 
+  /**
+   * Execute the clause and return the generated key with the type of the given path. If no rows
+   * were created, null is returned, otherwise the key of the first row is returned.
+   *
+   * <p>This method bypasses JPQL and executes a native SQL INSERT via JDBC to retrieve the
+   * generated key using Hibernate's {@code Session.doReturningWork()}.
+   *
+   * <p>Note: {@code INSERT ... SELECT} subqueries are not supported by this method.
+   *
+   * @param <T> key type
+   * @param path path for key (used to determine return type)
+   * @return generated key, or null if no rows were created
+   * @throws QueryException if a database error occurs
+   */
+  @SuppressWarnings("unchecked")
+  @Nullable
+  public <T> T executeWithKey(Path<T> path) {
+    return executeWithKey((Class<T>) path.getType());
+  }
+
+  /**
+   * Execute the clause and return the generated key cast to the given type. If no rows were
+   * created, null is returned, otherwise the key of the first row is returned.
+   *
+   * @param <T> key type
+   * @param type class of the key type
+   * @return generated key, or null if no rows were created
+   * @throws QueryException if a database error occurs
+   */
+  @Nullable
+  public <T> T executeWithKey(Class<T> type) {
+    if (subQuery != null) {
+      throw new UnsupportedOperationException(
+          "executeWithKey is not supported for INSERT ... SELECT subqueries");
+    }
+
+    var effectiveColumns = JpaInsertNativeHelper.effectiveColumns(inserts, columns);
+    if (effectiveColumns.isEmpty()) {
+      throw new IllegalStateException("No columns specified for insert");
+    }
+
+    var entityClass = queryMixin.getMetadata().getJoins().get(0).getTarget().getType();
+
+    // Serialize to collect constant values
+    var serializer = new JPQLSerializer(templates, null);
+    serializer.serializeForInsert(
+        queryMixin.getMetadata(), effectiveColumns, values, subQuery, inserts);
+
+    var params =
+        JpaInsertNativeHelper.resolveConstants(
+            serializer.getConstants(), queryMixin.getMetadata().getParams());
+
+    var sql = JpaInsertNativeHelper.buildNativeInsertSQL(entityClass, effectiveColumns);
+
+    return session.doReturningWork(
+        connection -> {
+          try {
+            return JpaInsertNativeHelper.executeAndReturnKey(connection, sql, params, type);
+          } catch (SQLException e) {
+            throw new QueryException("Failed to execute insert with generated key", e);
+          }
+        });
+  }
+
   @Override
   public HibernateInsertClause columns(Path<?>... columns) {
     this.columns.addAll(Arrays.asList(columns));

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/hibernate/NoSessionHolder.java

@@ -13,6 +13,7 @@
  */
 package com.querydsl.jpa.hibernate;
 
+import org.hibernate.jdbc.ReturningWork;
 import org.hibernate.query.NativeQuery;
 import org.hibernate.query.Query;
 
@@ -36,4 +37,9 @@ public Query<?> createQuery(String queryString) {
   public NativeQuery<?> createSQLQuery(String queryString) {
     throw new UnsupportedOperationException("No session in detached Query available");
   }
+
+  @Override
+  public <T> T doReturningWork(ReturningWork<T> work) {
+    throw new UnsupportedOperationException("No session in detached Query available");
+  }
 }

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/hibernate/SessionHolder.java

@@ -13,6 +13,7 @@
  */
 package com.querydsl.jpa.hibernate;
 
+import org.hibernate.jdbc.ReturningWork;
 import org.hibernate.query.NativeQuery;
 import org.hibernate.query.Query;
 
@@ -38,4 +39,13 @@ public interface SessionHolder {
    * @return query
    */
   NativeQuery<?> createSQLQuery(String queryString);
+
+  /**
+   * Execute a {@link ReturningWork} within the session's transaction context.
+   *
+   * @param <T> the return type
+   * @param work the work to execute with a JDBC Connection
+   * @return the result of the work
+   */
+  <T> T doReturningWork(ReturningWork<T> work);
 }

querydsl-libraries/querydsl-jpa/src/main/java/com/querydsl/jpa/hibernate/StatelessSessionHolder.java

@@ -14,6 +14,7 @@
 package com.querydsl.jpa.hibernate;
 
 import org.hibernate.StatelessSession;
+import org.hibernate.jdbc.ReturningWork;
 import org.hibernate.query.NativeQuery;
 import org.hibernate.query.Query;
 
@@ -41,4 +42,9 @@ public Query<?> createQuery(String queryString) {
   public NativeQuery<?> createSQLQuery(String queryString) {
     return session.createNativeQuery(queryString);
   }
+
+  @Override
+  public <T> T doReturningWork(ReturningWork<T> work) {
+    return session.doReturningWork(work);
+  }
 }