Add support for Postgres lock types (no key, share, key share) with FOR UPDATE

ebean-api/src/main/java/io/ebean/ExpressionList.java

@@ -174,6 +174,11 @@ public interface ExpressionList<T> {
    */
   Query<T> forUpdate();
 
+  /**
+   * Execute using "for update" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdate(Query.LockType lockType);
+
   /**
    * Execute using "for update" clause with No Wait option.
    * <p>
@@ -182,6 +187,11 @@ public interface ExpressionList<T> {
    */
   Query<T> forUpdateNoWait();
 
+  /**
+   * Execute using "for update nowait" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdateNoWait(Query.LockType lockType);
+
   /**
    * Execute using "for update" clause with Skip Locked option.
    * <p>
@@ -190,6 +200,11 @@ public interface ExpressionList<T> {
    */
   Query<T> forUpdateSkipLocked();
 
+  /**
+   * Execute using "for update skip locked" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdateSkipLocked(Query.LockType lockType);
+
   /**
    * Execute the query including soft deleted rows.
    */

ebean-api/src/main/java/io/ebean/Query.java

@@ -179,6 +179,36 @@ import java.util.stream.Stream;
  */
 public interface Query<T> {
 
+  /**
+   * The lock type (strength) to use with query FOR UPDATE row locking.
+   */
+  enum LockType {
+    /**
+     * The default lock type - See PlatformConfig.lockWithKey option.
+     */
+    Default,
+
+    /**
+     * FOR UPDATE.
+     */
+    Update,
+
+    /**
+     * FOR NO KEY UPDATE.
+     */
+    NoKeyUpdate,
+
+    /**
+     * FOR SHARE.
+     */
+    Share,
+
+    /**
+     * FOR KEY SHARE.
+     */
+    KeyShare
+  }
+
   /**
    * For update mode.
    */
@@ -1618,6 +1648,11 @@ public interface Query<T> {
    */
   Query<T> forUpdate();
 
+  /**
+   * Execute using "for update" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdate(LockType lockType);
+
   /**
    * Execute using "for update" clause with "no wait" option.
    * <p>
@@ -1626,6 +1661,11 @@ public interface Query<T> {
    */
   Query<T> forUpdateNoWait();
 
+  /**
+   * Execute using "for update nowait" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdateNoWait(LockType lockType);
+
   /**
    * Execute using "for update" clause with "skip locked" option.
    * <p>
@@ -1634,6 +1674,11 @@ public interface Query<T> {
    */
   Query<T> forUpdateSkipLocked();
 
+  /**
+   * Execute using "for update skip locked" with given lock type (currently Postgres only).
+   */
+  Query<T> forUpdateSkipLocked(LockType lockType);
+
   /**
    * Return true if this query has forUpdate set.
    */
@@ -1644,6 +1689,11 @@ public interface Query<T> {
    */
   ForUpdate getForUpdateMode();
 
+  /**
+   * Return the lock type (strength) to use with "for update".
+   */
+  LockType getForUpdateLockType();
+
   /**
    * Set root table alias.
    */

ebean-api/src/main/java/io/ebean/config/dbplatform/DatabasePlatform.java

@@ -679,9 +679,8 @@ public class DatabasePlatform {
 
   public String completeSql(String sql, Query<?> query) {
     if (query.isForUpdate()) {
-      sql = withForUpdate(sql, query.getForUpdateMode());
+      sql = withForUpdate(sql, query.getForUpdateMode(), query.getForUpdateLockType());
     }
-
     return sql;
   }
 
@@ -693,7 +692,7 @@ public class DatabasePlatform {
     return null;
   }
 
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     // silently assume the database does not support the "for update" clause.
     logger.info("it seems your database does not support the 'for update' clause");
     return sql;

ebean-api/src/main/java/io/ebean/config/dbplatform/h2/H2Platform.java

@@ -51,7 +51,7 @@ public class H2Platform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     // NOWAIT and SKIP LOCKED currently not supported with H2
     return sql + " for update";
   }

ebean-api/src/main/java/io/ebean/config/dbplatform/hana/HanaPlatform.java

@@ -1,5 +1,6 @@
 package io.ebean.config.dbplatform.hana;
 
+import io.ebean.Query;
 import io.ebean.Query.ForUpdate;
 import io.ebean.annotation.PersistBatch;
 import io.ebean.annotation.Platform;
@@ -68,7 +69,7 @@ public class HanaPlatform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, ForUpdate forUpdateMode, Query.LockType lockType) {
     switch (forUpdateMode) {
       case BASE:
         return sql + " for update";

ebean-api/src/main/java/io/ebean/config/dbplatform/mysql/BaseMySqlPlatform.java

@@ -60,7 +60,7 @@ public abstract class BaseMySqlPlatform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     // NOWAIT and SKIP LOCKED currently not supported with MySQL
     return sql + " for update";
   }

ebean-api/src/main/java/io/ebean/config/dbplatform/nuodb/NuoDbPlatform.java

@@ -48,7 +48,7 @@ public class NuoDbPlatform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     switch (forUpdateMode) {
       case NOWAIT:
         return sql + " for update nowait";

ebean-api/src/main/java/io/ebean/config/dbplatform/oracle/OraclePlatform.java

@@ -79,7 +79,7 @@ public class OraclePlatform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     switch (forUpdateMode) {
       case SKIPLOCKED:
         return sql + " for update skip locked";

ebean-api/src/main/java/io/ebean/config/dbplatform/postgres/PostgresPlatform.java

@@ -26,10 +26,14 @@ import java.sql.Types;
  */
 public class PostgresPlatform extends DatabasePlatform {
 
-  // by default using NO KEY option with FOR UPDATE clauses
-  private String forUpdateSkipLocked = " for no key update skip locked";
-  private String forUpdateNowait = " for no key update nowait";
-  private String forUpdate = " for no key update";
+  private static final String SKIP_LOCKED = " skip locked";
+  private static final String NO_WAIT = " nowait";
+  private static final String FOR_UPDATE = " for update";
+  private static final String FOR_NO_KEY_UPDATE = " for no key update";
+  private static final String FOR_SHARE = " for share";
+  private static final String FOR_KEY_SHARE = " for key share";
+
+  private boolean defaultLockWithKey = true;
 
   public PostgresPlatform() {
     super();
@@ -90,11 +94,7 @@ public class PostgresPlatform extends DatabasePlatform {
   @Override
   public void configure(PlatformConfig config) {
     super.configure(config);
-    if (config.isLockWithKey()) {
-      this.forUpdateSkipLocked = " for update skip locked";
-      this.forUpdateNowait = " for update nowait";
-      this.forUpdate = " for update";
-    }
+    defaultLockWithKey = config.isLockWithKey();
   }
 
   @Override
@@ -124,22 +124,32 @@ public class PostgresPlatform extends DatabasePlatform {
    */
   @Override
   public PlatformIdGenerator createSequenceIdGenerator(BackgroundExecutor be, DataSource ds, int stepSize, String seqName) {
-
     return new PostgresSequenceIdGenerator(be, ds, seqName, sequenceBatchSize);
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     switch (forUpdateMode) {
       case SKIPLOCKED:
-        return sql + forUpdateSkipLocked;
+        return sql + lock(lockType) + SKIP_LOCKED;
       case NOWAIT:
-        return sql + forUpdateNowait;
+        return sql + lock(lockType) + NO_WAIT;
       default:
-        return sql + forUpdate;
+        return sql + lock(lockType);
     }
   }
 
+  private String lock(Query.LockType lockType) {
+    switch (lockType) {
+      case Update: return FOR_UPDATE;
+      case NoKeyUpdate: return FOR_NO_KEY_UPDATE;
+      case Share: return FOR_SHARE;
+      case KeyShare: return FOR_KEY_SHARE;
+      case Default: return defaultLockWithKey ? FOR_UPDATE : FOR_NO_KEY_UPDATE;
+    }
+    return FOR_UPDATE;
+  }
+
   @Override
   public boolean tablePartitionsExist(Connection connection, String table) throws SQLException {
     try (PreparedStatement statement = connection.prepareStatement("select count(*) from pg_inherits i WHERE  i.inhparent = ?::regclass")) {

ebean-api/src/main/java/io/ebean/config/dbplatform/sqlserver/SqlServerBasePlatform.java

@@ -117,7 +117,7 @@ abstract class SqlServerBasePlatform extends DatabasePlatform {
   }
 
   @Override
-  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode) {
+  protected String withForUpdate(String sql, Query.ForUpdate forUpdateMode, Query.LockType lockType) {
     // for update are hints on from clause of base table
     return sql;
   }

ebean-core/src/main/java/io/ebeaninternal/server/expression/DefaultExpressionList.java

@@ -491,16 +491,31 @@ public class DefaultExpressionList<T> implements SpiExpressionList<T> {
     return query.forUpdate();
   }
 
+  @Override
+  public Query<T> forUpdate(Query.LockType lockType) {
+    return query.forUpdate(lockType);
+  }
+
   @Override
   public Query<T> forUpdateNoWait() {
     return query.forUpdateNoWait();
   }
 
+  @Override
+  public Query<T> forUpdateNoWait(Query.LockType lockType) {
+    return query.forUpdateNoWait(lockType);
+  }
+
   @Override
   public Query<T> forUpdateSkipLocked() {
     return query.forUpdateSkipLocked();
   }
 
+  @Override
+  public Query<T> forUpdateSkipLocked(Query.LockType lockType) {
+    return query.forUpdateSkipLocked(lockType);
+  }
+
   @Override
   public Query<T> select(String fetchProperties) {
     return query.select(fetchProperties);

ebean-core/src/main/java/io/ebeaninternal/server/query/DefaultFetchGroupQuery.java

@@ -559,6 +559,21 @@ class DefaultFetchGroupQuery<T> implements SpiFetchGroupQuery<T> {
     throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
   }
 
+  @Override
+  public Query<T> forUpdate(LockType lockType) {
+    throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
+  }
+
+  @Override
+  public Query<T> forUpdateNoWait(LockType lockType) {
+    throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
+  }
+
+  @Override
+  public Query<T> forUpdateSkipLocked(LockType lockType) {
+    throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
+  }
+
   @Override
   public boolean isForUpdate() {
     return false;
@@ -569,6 +584,11 @@ class DefaultFetchGroupQuery<T> implements SpiFetchGroupQuery<T> {
     throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
   }
 
+  @Override
+  public LockType getForUpdateLockType() {
+    throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
+  }
+
   @Override
   public Query<T> alias(String alias) {
     throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");

ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultOrmQuery.java

@@ -236,10 +236,8 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
    */
   private Boolean autoTune;
 
-  /**
-   * For update mode.
-   */
   private ForUpdate forUpdate;
+  private LockType lockType;
 
   private boolean singleAttribute;
 
@@ -967,21 +965,37 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
 
   @Override
   public DefaultOrmQuery<T> forUpdate() {
-    return setForUpdateWithMode(ForUpdate.BASE);
+    return setForUpdateWithMode(ForUpdate.BASE, LockType.Default);
+  }
+
+  @Override
+  public Query<T> forUpdate(LockType lockType) {
+    return setForUpdateWithMode(ForUpdate.BASE, lockType);
+  }
+
+  @Override
+  public Query<T> forUpdateNoWait(LockType lockType) {
+    return setForUpdateWithMode(ForUpdate.NOWAIT, lockType);
+  }
+
+  @Override
+  public Query<T> forUpdateSkipLocked(LockType lockType) {
+    return setForUpdateWithMode(ForUpdate.SKIPLOCKED, lockType);
   }
 
   @Override
   public DefaultOrmQuery<T> forUpdateNoWait() {
-    return setForUpdateWithMode(ForUpdate.NOWAIT);
+    return setForUpdateWithMode(ForUpdate.NOWAIT, LockType.Default);
   }
 
   @Override
   public DefaultOrmQuery<T> forUpdateSkipLocked() {
-    return setForUpdateWithMode(ForUpdate.SKIPLOCKED);
+    return setForUpdateWithMode(ForUpdate.SKIPLOCKED, LockType.Default);
   }
 
-  private DefaultOrmQuery<T> setForUpdateWithMode(ForUpdate mode) {
+  private DefaultOrmQuery<T> setForUpdateWithMode(ForUpdate mode, LockType lockType) {
     this.forUpdate = mode;
+    this.lockType = lockType;
     this.useBeanCache = CacheMode.OFF;
     return this;
   }
@@ -996,6 +1010,11 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
     return forUpdate;
   }
 
+  @Override
+  public LockType getForUpdateLockType() {
+    return lockType;
+  }
+
   @Override
   public ProfilingListener getProfilingListener() {
     return profilingListener;

ebean-core/src/test/java/io/ebean/config/dbplatform/PostgresPlatformTest.java

@@ -30,9 +30,24 @@ public class PostgresPlatformTest {
     platform.configure(config);
 
     assertThat(config.isLockWithKey()).isTrue();
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED)).isEqualTo("X for update skip locked");
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT)).isEqualTo("X for update nowait");
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE)).isEqualTo("X for update");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.Default)).isEqualTo("X for update skip locked");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.Default)).isEqualTo("X for update nowait");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.Default)).isEqualTo("X for update");
+
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.Update)).isEqualTo("X for update skip locked");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.NoKeyUpdate)).isEqualTo("X for no key update skip locked");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.Share)).isEqualTo("X for share skip locked");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.KeyShare)).isEqualTo("X for key share skip locked");
+
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.Update)).isEqualTo("X for update nowait");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.NoKeyUpdate)).isEqualTo("X for no key update nowait");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.Share)).isEqualTo("X for share nowait");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.KeyShare)).isEqualTo("X for key share nowait");
+
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.Update)).isEqualTo("X for update");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.NoKeyUpdate)).isEqualTo("X for no key update");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.Share)).isEqualTo("X for share");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.KeyShare)).isEqualTo("X for key share");
   }
 
   @Test
@@ -45,9 +60,9 @@ public class PostgresPlatformTest {
     platform.configure(config);
 
     assertThat(config.isLockWithKey()).isFalse();
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED)).isEqualTo("X for no key update skip locked");
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT)).isEqualTo("X for no key update nowait");
-    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE)).isEqualTo("X for no key update");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.SKIPLOCKED, Query.LockType.Default)).isEqualTo("X for no key update skip locked");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.NOWAIT, Query.LockType.Default)).isEqualTo("X for no key update nowait");
+    assertThat(platform.withForUpdate("X", Query.ForUpdate.BASE, Query.LockType.Default)).isEqualTo("X for no key update");
   }
 
 }

ebean-core/src/test/java/org/tests/basic/TestQueryForUpdatePostgresLock.java

@@ -5,7 +5,6 @@ import io.ebean.DB;
 import io.ebean.annotation.ForPlatform;
 import io.ebean.annotation.Platform;
 import io.ebean.annotation.Transactional;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -16,6 +15,7 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
+import static io.ebean.Query.LockType.NoKeyUpdate;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TestQueryForUpdatePostgresLock extends BaseTestCase {
@@ -27,7 +27,6 @@ public class TestQueryForUpdatePostgresLock extends BaseTestCase {
   private long timePreLock;
   private long timePostLock;
 
-  @Ignore // restore this test with explicit use of NO KEY lock strength
   @Test
   @ForPlatform(Platform.POSTGRES)
   public void testForUpdatePostgresLock() throws InterruptedException {
@@ -56,7 +55,7 @@ public class TestQueryForUpdatePostgresLock extends BaseTestCase {
   private void lockArticle(Integer id) {
     timePreLock = System.currentTimeMillis();
     log.info("lock start");
-    DB.find(Article.class).setId(id).forUpdate().findOne();
+    DB.find(Article.class).setId(id).forUpdate(NoKeyUpdate).findOne();
     sleep(1000);
     timePostLock = System.currentTimeMillis();
     log.info("lock done");

ebean-querybean/src/main/java/io/ebean/typequery/TQRootBean.java

@@ -686,24 +686,39 @@ public abstract class TQRootBean<T, R> {
   }
 
   /**
-   * executed the select with "for update" which should lock the record "on read"
+   * Execute using "for update" clause which results in the DB locking the record.
    */
   public R forUpdate() {
     query.forUpdate();
     return root;
   }
 
+  /**
+   * Execute using "for update" with given lock type (currently Postgres only).
+   */
+  public R forUpdate(Query.LockType lockType) {
+    query.forUpdate(lockType);
+    return root;
+  }
+
   /**
    * Execute using "for update" clause with "no wait" option.
    * <p>
    * This is typically a Postgres and Oracle only option at this stage.
-   * </p>
    */
   public R forUpdateNoWait() {
     query.forUpdateNoWait();
     return root;
   }
 
+  /**
+   * Execute using "for update nowait" with given lock type (currently Postgres only).
+   */
+  public R forUpdateNoWait(Query.LockType lockType) {
+    query.forUpdateNoWait(lockType);
+    return root;
+  }
+
   /**
    * Execute using "for update" clause with "skip locked" option.
    * <p>
@@ -715,6 +730,14 @@ public abstract class TQRootBean<T, R> {
     return root;
   }
 
+  /**
+   * Execute using "for update skip locked" with given lock type (currently Postgres only).
+   */
+  public R forUpdateSkipLocked(Query.LockType lockType) {
+    query.forUpdateSkipLocked(lockType);
+    return root;
+  }
+
   /**
    * Return this query as an UpdateQuery.
    *