#2362 - ENH: Add lock(bean) method as convenience for DB pessimistic locking

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

@@ -1272,6 +1272,21 @@ public final class DB {
     return beanId(bean);
   }
 
+  /**
+   * Load and lock the bean using {@code select for update}.
+   * <p>
+   * This should be executed inside a transaction.
+   * <p>
+   * The bean needs to have an ID property set and can be a reference bean (only has ID)
+   * or partially or fully populated bean. This will load all the properties of the bean
+   * from the database using {@code select for update}.
+   *
+   * @param bean The entity bean that we wish to obtain a database lock on.
+   */
+  public static void lock(Object bean) {
+    getDefault().lock(bean);
+  }
+
   /**
    * Deprecated migrate to cacheManager().
    */

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

@@ -1714,6 +1714,20 @@ public interface Database {
    */
   <T> Set<String> validateQuery(Query<T> query);
 
+  /**
+   * Load and lock the bean using {@code select for update}.
+   * <p>
+   * This should be executed inside a transaction and results in the bean being loaded or
+   * refreshed from the database and a database row lock held via {@code select for update}.
+   * <p>
+   * The bean needs to have an ID property set and can be a reference bean (only has ID)
+   * or partially or fully populated bean. This will load all the properties of the bean
+   * from the database using {@code select for update} obtaining a database row lock (using WAIT).
+   *
+   * @param bean The entity bean that we wish to obtain a database lock on.
+   */
+  void lock(Object bean);
+
   /**
    * Truncate all the given tables.
    */

ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java

@@ -129,16 +129,7 @@ import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.time.Clock;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.Spliterator;
+import java.util.*;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
@@ -930,6 +921,20 @@ public final class DefaultServer implements SpiServer, SpiEbeanServer {
     executeInTrans((txn) -> persister.merge(desc, checkEntityBean(bean), options, txn), transaction);
   }
 
+  @Override
+  public void lock(Object bean) {
+    BeanDescriptor<?> desc = descriptor(bean.getClass());
+    if (desc == null) {
+      throw new PersistenceException(bean.getClass() + " is NOT an Entity Bean registered with this server?");
+    }
+    Object id = desc.id(bean);
+    Objects.requireNonNull(id, "Bean missing an @Id value which is required to lock");
+    new DefaultOrmQuery<>(desc, this, expressionFactory)
+      .setId(id)
+      .withLock(Query.LockType.DEFAULT, Query.LockWait.NOWAIT)
+      .findOne();
+  }
+
   @Override
   public <T> Query<T> find(Class<T> beanType) {
     return createQuery(beanType);

ebean-test/src/test/java/io/ebeaninternal/api/TDSpiServer.java

@@ -137,6 +137,10 @@ public class TDSpiServer implements SpiServer {
     return null;
   }
 
+  @Override
+  public void lock(Object bean) {
+  }
+
   @Override
   public Object nextId(Class<?> beanType) {
     return null;

ebean-test/src/test/java/org/tests/basic/TestQueryForUpdate.java

@@ -7,12 +7,13 @@ import io.ebean.test.LoggedSql;
 import org.junit.jupiter.api.Test;
 import org.tests.model.basic.Customer;
 import org.tests.model.basic.EBasic;
+import org.tests.model.basic.Order;
 import org.tests.model.basic.ResetBasicData;
 
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class TestQueryForUpdate extends BaseTestCase {
 
@@ -40,6 +41,69 @@ public class TestQueryForUpdate extends BaseTestCase {
     }
   }
 
+  @Test
+  @ForPlatform({ Platform.H2, Platform.ORACLE, Platform.POSTGRES, Platform.SQLSERVER, Platform.MYSQL, Platform.MARIADB})
+  public void testForUpdate_when_alreadyInPCAsReference() {
+    ResetBasicData.reset();
+    Order o0 = DB.find(Order.class).orderBy("id").setMaxRows(1).findOne();
+    Integer customerId = o0.getCustomer().getId();
+
+    try (Transaction transaction = DB.beginTransaction()) {
+
+      LoggedSql.start();
+      Order order = DB.find(Order.class, o0.getId());
+      assert order != null;
+
+      Customer customer = order.getCustomer();
+      assertTrue(DB.beanState(customer).isReference());
+      assertEquals(customerId, customer.getId());
+
+      Customer customer1 = DB.find(Customer.class).where().idEq(customer.getId()).forUpdate().findOne();
+
+      assertThat(customer).isSameAs(customer1);
+      assertThat(customer.getName()).isNotNull();
+
+      List<String> sql = LoggedSql.stop();
+      assertThat(sql).hasSize(2);
+      assertThat(sql.get(0)).contains("from o_order");
+      assertThat(sql.get(1)).contains("from o_customer t0 where t0.id = ?   for update");
+
+      transaction.commit();
+    }
+  }
+
+  @Test
+  @ForPlatform({ Platform.H2, Platform.ORACLE, Platform.POSTGRES, Platform.SQLSERVER, Platform.MYSQL, Platform.MARIADB})
+  public void testForUpdate_when_alreadyInPCAsReference_usingLock() {
+    ResetBasicData.reset();
+    Order o0 = DB.find(Order.class).orderBy("id").setMaxRows(1).findOne();
+    Integer customerId = o0.getCustomer().getId();
+
+    try (Transaction transaction = DB.beginTransaction()) {
+
+      LoggedSql.start();
+      Order order = DB.find(Order.class, o0.getId());
+      assert order != null;
+
+      Customer customer = order.getCustomer();
+      assertEquals(customerId, customer.getId());
+      assertTrue(DB.beanState(customer).isReference());
+
+      DB.lock(customer); // load the customer bean using select for update
+      // bean is now loaded and database row lock held until commit
+      assertFalse(DB.beanState(customer).isReference());
+
+      assertThat(customer.getName()).isNotNull();
+
+      List<String> sql = LoggedSql.stop();
+      assertThat(sql).hasSize(2);
+      assertThat(sql.get(0)).contains("from o_order");
+      assertThat(sql.get(1)).contains("from o_customer t0 where t0.id = ?   for update");
+
+      transaction.commit();
+    }
+  }
+
   @Test
   @ForPlatform({ Platform.H2, Platform.ORACLE, Platform.POSTGRES, Platform.SQLSERVER, Platform.MYSQL, Platform.MARIADB})
   public void testForUpdate_when_alreadyInPC() {