#1261 - ENH: Support use of getGeneratedKeys with SqlUpdate

src/main/java/io/ebean/SqlUpdate.java

@@ -54,6 +54,25 @@ public interface SqlUpdate {
    */
   int execute();
 
+  /**
+   * Return the generated key value.
+   */
+  Object getGeneratedKey();
+
+  /**
+   * Execute and return the generated key. This is effectively a short cut for:
+   *
+   * <pre>{@code
+   *
+   *   sqlUpdate.execute();
+   *   Object key = sqlUpdate.getGeneratedKey();
+   *
+   * }</pre>
+   *
+   * @return The generated key value
+   */
+  Object executeGetKey();
+
   /**
    * Return true if eBean should automatically deduce the table modification
    * information and process it.
@@ -88,6 +107,11 @@ public interface SqlUpdate {
    */
   SqlUpdate setLabel(String label);
 
+  /**
+   * Set to true when we want to use getGeneratedKeys with this statement.
+   */
+  SqlUpdate setGetGeneratedKeys(boolean getGeneratedKeys);
+
   /**
    * Return the sql statement.
    */

src/main/java/io/ebeaninternal/api/SpiSqlUpdate.java

@@ -4,7 +4,23 @@ import io.ebean.SqlUpdate;
 
 public interface SpiSqlUpdate extends SqlUpdate {
 
+  /**
+   * Return the Bind parameters.
+   */
   BindParams getBindParams();
 
+  /**
+   * Set the final sql being executed with named parameters replaced etc.
+   */
   void setGeneratedSql(String sql);
+
+  /**
+   * Return true if we are using getGeneratedKeys.
+   */
+  boolean isGetGeneratedKeys();
+
+  /**
+   * Set the generated key value.
+   */
+  void setGeneratedKey(Object idValue);
 }

src/main/java/io/ebeaninternal/server/core/DefaultSqlUpdate.java

@@ -65,6 +65,10 @@ public final class DefaultSqlUpdate implements Serializable, SpiSqlUpdate {
    */
   private int addPos;
 
+  private boolean getGeneratedKeys;
+
+  private Object generatedKey;
+
   /**
    * Create with server sql and bindParams object.
    * <p>
@@ -93,6 +97,12 @@ public final class DefaultSqlUpdate implements Serializable, SpiSqlUpdate {
     this(null, sql, new BindParams());
   }
 
+  @Override
+  public Object executeGetKey() {
+    execute();
+    return getGeneratedKey();
+  }
+
   @Override
   public int execute() {
     if (server != null) {
@@ -103,6 +113,16 @@ public final class DefaultSqlUpdate implements Serializable, SpiSqlUpdate {
     }
   }
 
+  @Override
+  public Object getGeneratedKey() {
+    return generatedKey;
+  }
+
+  @Override
+  public void setGeneratedKey(Object idValue) {
+    this.generatedKey = idValue;
+  }
+
   @Override
   public boolean isAutoTableMod() {
     return isAutoTableMod;
@@ -125,6 +145,17 @@ public final class DefaultSqlUpdate implements Serializable, SpiSqlUpdate {
     return this;
   }
 
+  @Override
+  public boolean isGetGeneratedKeys() {
+    return getGeneratedKeys;
+  }
+
+  @Override
+  public SqlUpdate setGetGeneratedKeys(boolean getGeneratedKeys) {
+    this.getGeneratedKeys = getGeneratedKeys;
+    return this;
+  }
+
   @Override
   public String getGeneratedSql() {
     return generatedSql;

src/main/java/io/ebeaninternal/server/core/PersistRequestUpdateSql.java

@@ -73,6 +73,11 @@ public final class PersistRequestUpdateSql extends PersistRequest {
    */
   @Override
   public void setGeneratedKey(Object idValue) {
+    updateSql.setGeneratedKey(idValue);
+  }
+
+  public boolean isGetGeneratedKeys() {
+    return updateSql.isGetGeneratedKeys();
   }
 
   /**

src/main/java/io/ebeaninternal/server/persist/ExeOrmUpdate.java

@@ -16,7 +16,7 @@ import java.sql.SQLException;
 /**
  * Executes the UpdateSql requests.
  */
-public class ExeOrmUpdate {
+class ExeOrmUpdate {
 
   private static final Logger logger = LoggerFactory.getLogger(ExeOrmUpdate.class);
 
@@ -27,7 +27,7 @@ public class ExeOrmUpdate {
   /**
    * Create with a given binder.
    */
-  public ExeOrmUpdate(Binder binder) {
+  ExeOrmUpdate(Binder binder) {
     this.pstmtFactory = new PstmtFactory();
     this.binder = binder;
   }
@@ -107,7 +107,7 @@ public class ExeOrmUpdate {
       if (logSql) {
         t.logSql(sql);
       }
-      pstmt = pstmtFactory.getPstmt(t, sql);
+      pstmt = pstmtFactory.getPstmt(t, sql, false);
     }
 
     String bindLog = null;

src/main/java/io/ebeaninternal/server/persist/ExeUpdateSql.java

@@ -6,17 +6,19 @@ import io.ebeaninternal.api.SpiTransaction;
 import io.ebeaninternal.server.core.PersistRequestUpdateSql;
 import io.ebeaninternal.server.core.PersistRequestUpdateSql.SqlType;
 import io.ebeaninternal.server.util.BindParamsParser;
+import io.ebeaninternal.util.JdbcClose;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.persistence.PersistenceException;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 
 /**
  * Executes the UpdateSql requests.
  */
-public class ExeUpdateSql {
+class ExeUpdateSql {
 
   private static final Logger logger = LoggerFactory.getLogger(ExeUpdateSql.class);
 
@@ -27,7 +29,7 @@ public class ExeUpdateSql {
   /**
    * Create with a given binder.
    */
-  public ExeUpdateSql(Binder binder) {
+  ExeUpdateSql(Binder binder) {
     this.binder = binder;
     this.pstmtFactory = new PstmtFactory();
   }
@@ -51,6 +53,9 @@ public class ExeUpdateSql {
       } else {
         int rowCount = pstmt.executeUpdate();
         request.checkRowCount(rowCount);
+        if (request.isGetGeneratedKeys()) {
+          readGeneratedKeys(pstmt, request);
+        }
         request.postExecute();
         return rowCount;
       }
@@ -68,6 +73,22 @@ public class ExeUpdateSql {
     }
   }
 
+  private void readGeneratedKeys(PreparedStatement stmt, PersistRequestUpdateSql request) {
+
+    ResultSet resultSet = null;
+      try {
+        resultSet = stmt.getGeneratedKeys();
+        if (resultSet.next()) {
+          request.setGeneratedKey(resultSet.getObject(1));
+        }
+      } catch (SQLException ex) {
+        throw new PersistenceException(ex);
+
+      } finally {
+        JdbcClose.close(resultSet);
+      }
+  }
+
   private PreparedStatement bindStmt(PersistRequestUpdateSql request, boolean batchThisRequest) throws SQLException {
 
     SpiSqlUpdate updateSql = request.getUpdateSql();
@@ -90,7 +111,7 @@ public class ExeUpdateSql {
       if (logSql) {
         t.logSql(TrimLogSql.trim(sql));
       }
-      pstmt = pstmtFactory.getPstmt(t, sql);
+      pstmt = pstmtFactory.getPstmt(t, sql, request.isGetGeneratedKeys());
     }
 
     if (updateSql.getTimeout() > 0) {

src/main/java/io/ebeaninternal/server/persist/PstmtFactory.java

@@ -6,6 +6,7 @@ import java.sql.CallableStatement;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.sql.Statement;
 
 /**
  * Factory for creating Statements.
@@ -14,9 +15,9 @@ import java.sql.SQLException;
  * getGeneratedKeys.
  * </p>
  */
-public class PstmtFactory {
+class PstmtFactory {
 
-  public PstmtFactory() {
+  PstmtFactory() {
   }
 
   /**
@@ -30,9 +31,13 @@ public class PstmtFactory {
   /**
    * Get a prepared statement without any batching.
    */
-  public PreparedStatement getPstmt(SpiTransaction t, String sql) throws SQLException {
+  public PreparedStatement getPstmt(SpiTransaction t, String sql, boolean getGeneratedKeys) throws SQLException {
     Connection conn = t.getInternalConnection();
-    return conn.prepareStatement(sql);
+    if (getGeneratedKeys) {
+      return conn.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
+    } else {
+      return conn.prepareStatement(sql);
+    }
   }
 
   /**

src/main/java/io/ebeaninternal/server/persist/dml/DmlHandler.java

@@ -24,6 +24,8 @@ public abstract class DmlHandler implements PersistHandler, BindableRequest {
 
   private static final Logger logger = LoggerFactory.getLogger(DmlHandler.class);
 
+  private static final int[] GENERATED_KEY_COLUMNS = new int[]{1};
+
   /**
    * The originating request.
    */
@@ -260,8 +262,7 @@ public abstract class DmlHandler implements PersistHandler, BindableRequest {
       // the Id generated is always the first column
       // Required to stop Oracle10 giving us Oracle rowId??
       // Other jdbc drivers seem fine without this hint.
-      int[] columns = {1};
-      return conn.prepareStatement(sql, columns);
+      return conn.prepareStatement(sql, GENERATED_KEY_COLUMNS);
 
     } else {
       return conn.prepareStatement(sql);

src/main/java/io/ebeaninternal/util/JdbcClose.java

@@ -4,6 +4,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 
 /**
@@ -13,6 +14,19 @@ public class JdbcClose {
 
   private static final Logger logger = LoggerFactory.getLogger(JdbcClose.class);
 
+  /**
+   * Close the resultSet logging if an error occurs.
+   */
+  public static void close(ResultSet resultSet) {
+    try {
+      if (resultSet != null) {
+        resultSet.close();
+      }
+    } catch (SQLException e) {
+      logger.warn("Error closing resultSet", e);
+    }
+  }
+
   /**
    * Close the connection logging if an error occurs.
    */

src/test/java/org/tests/update/EPersonOnline.java

@@ -0,0 +1,57 @@
+package org.tests.update;
+
+import io.ebean.annotation.Index;
+import io.ebean.annotation.WhenModified;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.Table;
+import java.time.Instant;
+
+@Entity
+@Table(name = "e_person_online")
+public class EPersonOnline {
+
+  @Id
+  Long id;
+
+  @Index(unique = true)
+  String email;
+
+  boolean online;
+
+  @WhenModified
+  Instant whenUpdated;
+
+  public Long getId() {
+    return id;
+  }
+
+  public void setId(Long id) {
+    this.id = id;
+  }
+
+  public String getEmail() {
+    return email;
+  }
+
+  public void setEmail(String email) {
+    this.email = email;
+  }
+
+  public boolean isOnline() {
+    return online;
+  }
+
+  public void setOnline(boolean online) {
+    this.online = online;
+  }
+
+  public Instant getWhenUpdated() {
+    return whenUpdated;
+  }
+
+  public void setWhenUpdated(Instant whenUpdated) {
+    this.whenUpdated = whenUpdated;
+  }
+}

src/test/java/org/tests/update/TestSqlUpdateUpsert.java

@@ -0,0 +1,128 @@
+package org.tests.update;
+
+import io.ebean.BaseTestCase;
+import io.ebean.Ebean;
+import io.ebean.SqlUpdate;
+import io.ebean.annotation.ForPlatform;
+import io.ebean.annotation.Platform;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TestSqlUpdateUpsert extends BaseTestCase {
+
+  @ForPlatform(Platform.H2)
+  @Test
+  public void h2Merge() {
+
+    String sql = "merge into e_person_online (email, online, when_updated) key(email) values (?, ?, now())";
+
+    String email = "baz@one.com";
+
+    Object key = Ebean.createSqlUpdate(sql)
+      .setGetGeneratedKeys(true)
+      .setParameter(1, email)
+      .setParameter(2, true)
+      .executeGetKey();
+
+    EPersonOnline found = Ebean.find(EPersonOnline.class, key);
+    assertThat(found).isNotNull();
+    assertThat(found.getEmail()).isEqualTo(email);
+    assertThat(found.isOnline()).isTrue();
+
+    String sqlNamed = "merge into e_person_online (email, online, when_updated) key(email) values (:email, :online, now())";
+
+    SqlUpdate sqlUpdate2 = Ebean.createSqlUpdate(sqlNamed)
+      .setGetGeneratedKeys(true)
+      .setParameter("email", email)
+      .setParameter("online", false);
+
+    Object key2 = sqlUpdate2.executeGetKey();
+    assertThat(key2).isNull();
+
+
+    EPersonOnline found2 = Ebean.find(EPersonOnline.class).where().eq("email", email).findOne();
+    assertThat(found2).isNotNull();
+    assertThat(found2.getId()).isEqualTo(key);
+    assertThat(found2.getEmail()).isEqualTo(email);
+    assertThat(found2.isOnline()).isFalse();
+    assertThat(found2.getWhenUpdated()).isGreaterThan(found.getWhenUpdated());
+  }
+
+  @ForPlatform(Platform.POSTGRES)
+  @Test
+  public void postgresUpsert() {
+
+    String sql = "insert into e_person_online (email, online, when_updated) values (?, ?, now()) on conflict (email) do update set when_updated=now(), online = ?";
+
+    String email = "foo@one.com";
+
+    Object key = Ebean.createSqlUpdate(sql)
+      .setGetGeneratedKeys(true)
+      .setParameter(1, email)
+      .setParameter(2, true)
+      .setParameter(3, true)
+      .executeGetKey();
+
+    EPersonOnline found = Ebean.find(EPersonOnline.class, key);
+    assertThat(found).isNotNull();
+    assertThat(found.getEmail()).isEqualTo("foo@one.com");
+    assertThat(found.isOnline()).isTrue();
+
+
+    String sqlNamed = "insert into e_person_online (email, online, when_updated) values (:email, :online, now()) on conflict (email) do update set when_updated=now(), online = :online";
+    SqlUpdate sqlUpdate2 = Ebean.createSqlUpdate(sqlNamed)
+      .setGetGeneratedKeys(true)
+      .setParameter("email", email)
+      .setParameter("online", false);
+
+    Object key2 = sqlUpdate2.executeGetKey();
+
+    EPersonOnline found2 = Ebean.find(EPersonOnline.class, key2);
+    assertThat(found2).isNotNull();
+    assertThat(found2.getId()).isEqualTo(key);
+    assertThat(found2.getEmail()).isEqualTo("foo@one.com");
+    assertThat(found2.isOnline()).isFalse();
+    assertThat(found2.getWhenUpdated()).isGreaterThan(found.getWhenUpdated());
+
+  }
+
+  @ForPlatform(Platform.MYSQL)
+  @Test
+  public void mySqlUpsert() {
+
+    String email = "bar@one.com";
+
+    String sql = "insert into e_person_online (email, online, when_updated) values (?, ?, current_time) on duplicate key update when_updated=current_time, online = ?";
+    SqlUpdate sqlUpdate = Ebean.createSqlUpdate(sql)
+      .setGetGeneratedKeys(true)
+      .setParameter(1, email)
+      .setParameter(2, true)
+      .setParameter(3, true);
+
+    Object key = sqlUpdate.executeGetKey();
+    assertThat(key).isNotNull();
+
+    EPersonOnline found = Ebean.find(EPersonOnline.class, key);
+    assertThat(found).isNotNull();
+    assertThat(found.getEmail()).isEqualTo("bar@one.com");
+    assertThat(found.isOnline()).isTrue();
+
+
+    String sqlNamed = "insert into e_person_online (email, online, when_updated) values (:email, :online, current_time) on duplicate key update when_updated=current_time, online = :online";
+    SqlUpdate sqlUpdate2 = Ebean.createSqlUpdate(sqlNamed)
+      .setGetGeneratedKeys(true)
+      .setParameter("email", email)
+      .setParameter("online", false);
+
+    sqlUpdate2.execute();
+    Object key2 = sqlUpdate2.getGeneratedKey();
+
+    EPersonOnline found2 = Ebean.find(EPersonOnline.class, key2);
+    assertThat(found2).isNotNull();
+    assertThat(found2.getId()).isEqualTo(key);
+    assertThat(found2.getEmail()).isEqualTo("bar@one.com");
+    assertThat(found2.isOnline()).isFalse();
+
+  }
+}

src/test/resources/ebean.properties

@@ -103,7 +103,7 @@ datasource.hsqldb.databaseDriver=org.hsqldb.jdbcDriver
 
 datasource.mysql.username=test_ebean
 datasource.mysql.password=test
-datasource.mysql.databaseUrl=jdbc:mysql://127.0.0.1:3306/test_ebean
+datasource.mysql.databaseUrl=jdbc:mysql://127.0.0.1:4306/test_ebean
 datasource.mysql.databaseDriver=com.mysql.jdbc.Driver