[DATAJDBC-493] Add findOneWithLock for DataAccessStrategy

Author: mhyeon-lee

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/CascadingDataAccessStrategy.java

@@ -24,6 +24,7 @@
 import org.springframework.data.domain.Sort;
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
+import org.springframework.data.relational.core.sql.LockMode;
 
 /**
  * Delegates each methods to the {@link DataAccessStrategy}s passed to the constructor in turn until the first that does
@@ -33,6 +34,7 @@
  * @author Mark Paluch
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  * @since 1.1
  */
 public class CascadingDataAccessStrategy implements DataAccessStrategy {
@@ -133,6 +135,15 @@ public <T> T findById(Object id, Class<T> domainType) {
 		return collect(das -> das.findById(id, domainType));
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findById(java.lang.Object, org.springframework.data.relational.core.sql.LockMode, java.lang.Class)
+	 */
+	@Override
+	public <T> T findByIdWithLock(Object id, LockMode lockMode, Class<T> domainType) {
+		return collect(das -> das.findByIdWithLock(id, lockMode, domainType));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findAll(java.lang.Class)

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/DataAccessStrategy.java

@@ -23,6 +23,7 @@
 import org.springframework.data.jdbc.core.JdbcAggregateOperations;
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.lang.Nullable;
 
 /**
@@ -33,6 +34,7 @@
  * @author Jens Schauder
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  */
 public interface DataAccessStrategy extends RelationResolver {
 
@@ -148,6 +150,18 @@ public interface DataAccessStrategy extends RelationResolver {
 	@Nullable
 	<T> T findById(Object id, Class<T> domainType);
 
+	/**
+	 * Loads a single entity identified by type and id with lock.
+	 *
+	 * @param id the id of the entity to load. Must not be {@code null}.
+	 * @param lockMode the lock mode for select. Must not be {@code null}.
+	 * @param domainType the domain type of the entity. Must not be {@code null}.
+	 * @param <T> the type of the entity.
+	 * @return Might return {@code null}.
+	 */
+	@Nullable
+	<T> T findByIdWithLock(Object id, LockMode lockMode, Class<T> domainType);
+
 	/**
 	 * Loads all entities of the given type.
 	 *

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/DefaultDataAccessStrategy.java

@@ -41,6 +41,7 @@
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.sql.IdentifierProcessing;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
 import org.springframework.jdbc.core.RowMapper;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
@@ -62,6 +63,7 @@
  * @author Tom Hombergs
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  * @since 1.1
  */
 public class DefaultDataAccessStrategy implements DataAccessStrategy {
@@ -269,6 +271,24 @@ public <T> T findById(Object id, Class<T> domainType) {
 		}
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findById(java.lang.Object, org.springframework.data.relational.core.sql.LockMode, java.lang.Class)
+	 */
+	@Override
+	@SuppressWarnings("unchecked")
+	public <T> T findByIdWithLock(Object id, LockMode lockMode, Class<T> domainType) {
+
+		String findOneWithLockSql = sql(domainType).getFindOneWithLock(lockMode);
+		SqlIdentifierParameterSource parameter = createIdParameterSource(id, domainType);
+
+		try {
+			return operations.queryForObject(findOneWithLockSql, parameter, (RowMapper<T>) getEntityRowMapper(domainType));
+		} catch (EmptyResultDataAccessException e) {
+			return null;
+		}
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findAll(java.lang.Class)

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/DelegatingDataAccessStrategy.java

@@ -19,6 +19,7 @@
 import org.springframework.data.domain.Sort;
 import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.util.Assert;
 
 /**
@@ -28,6 +29,7 @@
  * @author Jens Schauder
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  * @since 1.1
  */
 public class DelegatingDataAccessStrategy implements DataAccessStrategy {
@@ -128,6 +130,18 @@ public <T> T findById(Object id, Class<T> domainType) {
 		return delegate.findById(id, domainType);
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findById(java.lang.Object, org.springframework.data.relational.core.sql.LockMode, java.lang.Class)
+	 */
+	@Override
+	public <T> T findByIdWithLock(Object id, LockMode lockMode, Class<T> domainType) {
+
+		Assert.notNull(delegate, "Delegate is null");
+
+		return delegate.findByIdWithLock(id, lockMode, domainType);
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findAll(java.lang.Class)

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/SqlGenerator.java

@@ -52,6 +52,7 @@
  * @author Tom Hombergs
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  */
 class SqlGenerator {
 
@@ -257,6 +258,16 @@ String getFindOne() {
 		return findOneSql.get();
 	}
 
+	/**
+	 * Create a {@code SELECT … FROM … WHERE :id = … (LOCK CLAUSE)} statement.
+	 *
+	 * @param lockMode Lock clause mode.
+	 * @return the statement as a {@link String}. Guaranteed to be not {@literal null}.
+	 */
+	String getFindOneWithLock(LockMode lockMode) {
+		return this.createFindOneWithLockSql(lockMode);
+	}
+
 	/**
 	 * Create a {@code INSERT INTO … (…) VALUES(…)} statement.
 	 *
@@ -358,6 +369,15 @@ private String createFindOneSql() {
 		return render(select);
 	}
 
+	private String createFindOneWithLockSql(LockMode lockMode) {
+
+		Select select = selectBuilder().where(getIdColumn().isEqualTo(getBindMarker(ID_SQL_PARAMETER))) //
+			.lock(lockMode)
+			.build();
+
+		return render(select);
+	}
+
 	private String createFindAllSql() {
 		return render(selectBuilder().build());
 	}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/mybatis/MyBatisDataAccessStrategy.java

@@ -41,6 +41,7 @@
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.sql.IdentifierProcessing;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.util.Assert;
@@ -60,6 +61,7 @@
  * @author Mark Paluch
  * @author Tyler Van Gorder
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  */
 public class MyBatisDataAccessStrategy implements DataAccessStrategy {
 
@@ -260,6 +262,17 @@ public <T> T findById(Object id, Class<T> domainType) {
 		return sqlSession().selectOne(statement, parameter);
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findById(java.lang.Object, org.springframework.data.relational.core.sql.LockMode, java.lang.Class)
+	 */
+	@Override
+	public <T> T findByIdWithLock(Object id, LockMode lockMode, Class<T> domainType) {
+		String statement = namespace(domainType) + ".findByIdWithLock";
+		MyBatisContext parameter = new MyBatisContext(id, null, domainType, Collections.emptyMap());
+		return sqlSession().selectOne(statement, parameter);
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.jdbc.core.DataAccessStrategy#findAll(java.lang.Class)

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/core/convert/SqlGeneratorUnitTests.java

@@ -46,6 +46,7 @@
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.sql.Aliased;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
 import org.springframework.data.relational.core.sql.Table;
 
@@ -59,6 +60,7 @@
  * @author Mark Paluch
  * @author Tom Hombergs
  * @author Milan Milanov
+ * @author Myeonghyeon Lee
  */
 public class SqlGeneratorUnitTests {
 
@@ -95,16 +97,36 @@ public void findOne() {
 
 		SoftAssertions softAssertions = new SoftAssertions();
 		softAssertions.assertThat(sql) //
-				.startsWith("SELECT") //
-				.contains("dummy_entity.id1 AS id1,") //
-				.contains("dummy_entity.x_name AS x_name,") //
-				.contains("dummy_entity.x_other AS x_other,") //
-				.contains("ref.x_l1id AS ref_x_l1id") //
-				.contains("ref.x_content AS ref_x_content").contains(" FROM dummy_entity") //
-				.contains("ON ref.dummy_entity = dummy_entity.id1") //
-				.contains("WHERE dummy_entity.id1 = :id") //
-				// 1-N relationships do not get loaded via join
-				.doesNotContain("Element AS elements");
+			.startsWith("SELECT") //
+			.contains("dummy_entity.id1 AS id1,") //
+			.contains("dummy_entity.x_name AS x_name,") //
+			.contains("dummy_entity.x_other AS x_other,") //
+			.contains("ref.x_l1id AS ref_x_l1id") //
+			.contains("ref.x_content AS ref_x_content").contains(" FROM dummy_entity") //
+			.contains("ON ref.dummy_entity = dummy_entity.id1") //
+			.contains("WHERE dummy_entity.id1 = :id") //
+			// 1-N relationships do not get loaded via join
+			.doesNotContain("Element AS elements");
+		softAssertions.assertAll();
+	}
+
+	@Test // DATAJDBC-493
+	public void findOneWithLock() {
+
+		String sql = sqlGenerator.getFindOneWithLock(LockMode.PESSIMISTIC_WRITE);
+
+		SoftAssertions softAssertions = new SoftAssertions();
+		softAssertions.assertThat(sql) //
+			.startsWith("SELECT") //
+			.contains("dummy_entity.id1 AS id1,") //
+			.contains("dummy_entity.x_name AS x_name,") //
+			.contains("dummy_entity.x_other AS x_other,") //
+			.contains("ref.x_l1id AS ref_x_l1id") //
+			.contains("ref.x_content AS ref_x_content").contains(" FROM dummy_entity") //
+			.contains("ON ref.dummy_entity = dummy_entity.id1") //
+			.contains("WHERE dummy_entity.id1 = :id") //
+			.contains("FOR UPDATE") //
+			.doesNotContain("Element AS elements");
 		softAssertions.assertAll();
 	}