Introduced pessimistic locks for derived queries.

Methods which use the derive query functionality now can be annotated with `@Lock` to used a given `LockMode`. Right now there are two different modes `PESSIMISTIC_READ` and `PESSIMISTIC_WRITE`. Based on the dialect the right select is generated. For example for H2 `Select ... FOR UPDATE`.

Closes spring-projects/spring-data-jdbc#1041
See #643,
Original pull request spring-projects/spring-data-jdbc/pull/1158

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/DefaultStatementMapper.java

@@ -45,6 +45,7 @@ import org.springframework.util.Assert;
  * @author Mark Paluch
  * @author Roman Chigvintsev
  * @author Mingyuan Wu
+ * @author Diego Krupitza
  */
 class DefaultStatementMapper implements StatementMapper {
 
@@ -132,6 +133,10 @@ class DefaultStatementMapper implements StatementMapper {
 			selectBuilder.offset(selectSpec.getOffset());
 		}
 
+		if (selectSpec.getLock() != null) {
+			selectBuilder.lock(selectSpec.getLock());
+		}
+
 		Select select = selectBuilder.build();
 		return new DefaultPreparedOperation<>(select, this.renderContext, bindings);
 	}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/StatementMapper.java

@@ -15,13 +15,7 @@
  */
 package org.springframework.data.r2dbc.core;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -33,6 +27,7 @@ import org.springframework.data.r2dbc.dialect.R2dbcDialect;
 import org.springframework.data.relational.core.query.Criteria;
 import org.springframework.data.relational.core.query.CriteriaDefinition;
 import org.springframework.data.relational.core.sql.Expression;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
 import org.springframework.data.relational.core.sql.Table;
 import org.springframework.data.relational.core.sql.render.RenderContext;
@@ -53,6 +48,7 @@ import org.springframework.util.Assert;
  * @author Mark Paluch
  * @author Roman Chigvintsev
  * @author Mingyuan Wu
+ * @author Diego Krupitza
  */
 public interface StatementMapper {
 
@@ -227,9 +223,10 @@ public interface StatementMapper {
 		private final long offset;
 		private final int limit;
 		private final boolean distinct;
+		private final LockMode lockMode;
 
 		protected SelectSpec(Table table, List<String> projectedFields, List<Expression> selectList,
-				@Nullable CriteriaDefinition criteria, Sort sort, int limit, long offset, boolean distinct) {
+				@Nullable CriteriaDefinition criteria, Sort sort, int limit, long offset, boolean distinct, LockMode lockMode) {
 			this.table = table;
 			this.projectedFields = projectedFields;
 			this.selectList = selectList;
@@ -238,6 +235,7 @@ public interface StatementMapper {
 			this.offset = offset;
 			this.limit = limit;
 			this.distinct = distinct;
+			this.lockMode = lockMode;
 		}
 
 		/**
@@ -262,7 +260,7 @@ public interface StatementMapper {
 			List<String> projectedFields = Collections.emptyList();
 			List<Expression> selectList = Collections.emptyList();
 			return new SelectSpec(Table.create(table), projectedFields, selectList, Criteria.empty(), Sort.unsorted(), -1, -1,
-					false);
+					false, null);
 		}
 
 		public SelectSpec doWithTable(BiFunction<Table, SelectSpec, SelectSpec> function) {
@@ -304,7 +302,7 @@ public interface StatementMapper {
 			selectList.addAll(Arrays.asList(expressions));
 
 			return new SelectSpec(this.table, projectedFields, selectList, this.criteria, this.sort, this.limit, this.offset,
-					this.distinct);
+					this.distinct, this.lockMode);
 		}
 
 		/**
@@ -320,7 +318,7 @@ public interface StatementMapper {
 			selectList.addAll(projectedFields);
 
 			return new SelectSpec(this.table, this.projectedFields, selectList, this.criteria, this.sort, this.limit,
-					this.offset, this.distinct);
+					this.offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -331,7 +329,7 @@ public interface StatementMapper {
 		 */
 		public SelectSpec withCriteria(CriteriaDefinition criteria) {
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, criteria, this.sort, this.limit,
-					this.offset, this.distinct);
+					this.offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -344,11 +342,11 @@ public interface StatementMapper {
 
 			if (sort.isSorted()) {
 				return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, sort, this.limit,
-						this.offset, this.distinct);
+						this.offset, this.distinct, this.lockMode);
 			}
 
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, this.limit,
-					this.offset, this.distinct);
+					this.offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -364,11 +362,11 @@ public interface StatementMapper {
 				Sort sort = page.getSort();
 
 				return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria,
-						sort.isSorted() ? sort : this.sort, page.getPageSize(), page.getOffset(), this.distinct);
+						sort.isSorted() ? sort : this.sort, page.getPageSize(), page.getOffset(), this.distinct, this.lockMode);
 			}
 
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, this.limit,
-					this.offset, this.distinct);
+					this.offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -379,7 +377,7 @@ public interface StatementMapper {
 		 */
 		public SelectSpec offset(long offset) {
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, this.limit,
-					offset, this.distinct);
+					offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -390,7 +388,7 @@ public interface StatementMapper {
 		 */
 		public SelectSpec limit(int limit) {
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, limit,
-					this.offset, this.distinct);
+					this.offset, this.distinct, this.lockMode);
 		}
 
 		/**
@@ -400,7 +398,28 @@ public interface StatementMapper {
 		 */
 		public SelectSpec distinct() {
 			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, limit,
-					this.offset, true);
+					this.offset, true, this.lockMode);
+		}
+
+		/**
+		 * Associate a lock mode with the select and create a new {@link SelectSpec}.
+		 *
+		 * @param lockMode the {@link LockMode} we want to use. This might be null
+		 * @return the {@link SelectSpec}.
+		 */
+		public SelectSpec lock(LockMode lockMode) {
+			return new SelectSpec(this.table, this.projectedFields, this.selectList, this.criteria, this.sort, limit,
+					this.offset, this.distinct, lockMode);
+		}
+
+		/**
+		 * The used lockmode
+		 * 
+		 * @return might be null if no lockmode defined.
+		 */
+		@Nullable
+		public LockMode getLock() {
+			return this.lockMode;
 		}
 
 		public Table getTable() {
@@ -440,6 +459,7 @@ public interface StatementMapper {
 		public boolean isDistinct() {
 			return this.distinct;
 		}
+
 	}
 
 	/**

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/dialect/H2Dialect.java

@@ -1,5 +1,7 @@
 package org.springframework.data.r2dbc.dialect;
 
+import org.springframework.data.relational.core.dialect.AnsiDialect;
+import org.springframework.data.relational.core.dialect.LockClause;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
 
 /**
@@ -7,6 +9,7 @@ import org.springframework.data.relational.core.sql.SqlIdentifier;
  *
  * @author Mark Paluch
  * @author Jens Schauder
+ * @author Diego Krupitza
  */
 public class H2Dialect extends PostgresDialect {
 
@@ -19,4 +22,17 @@ public class H2Dialect extends PostgresDialect {
 	public String renderForGeneratedValues(SqlIdentifier identifier) {
 		return identifier.getReference(getIdentifierProcessing());
 	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.relational.core.dialect.Dialect#lock()
+	 */
+	@Override
+	public LockClause lock() {
+		// H2 Dialect does not support the same lock keywords as PostgreSQL, but it supports the ANSI SQL standard.
+		// see https://www.h2database.com/html/commands.html
+		// and https://www.h2database.com/html/features.html#compatibility
+		return AnsiDialect.INSTANCE.lock();
+	}
+
 }

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/repository/query/PartTreeR2dbcQuery.java

@@ -39,6 +39,7 @@ import org.springframework.r2dbc.core.PreparedOperation;
  *
  * @author Roman Chigvintsev
  * @author Mark Paluch
+ * @author Diego Krupitza
  * @since 1.1
  */
 public class PartTreeR2dbcQuery extends AbstractR2dbcQuery {
@@ -119,7 +120,7 @@ public class PartTreeR2dbcQuery extends AbstractR2dbcQuery {
 
 			RelationalEntityMetadata<?> entityMetadata = getQueryMethod().getEntityInformation();
 			R2dbcQueryCreator queryCreator = new R2dbcQueryCreator(tree, dataAccessStrategy, entityMetadata, accessor,
-					projectedProperties);
+					projectedProperties, this.getQueryMethod().getLock());
 			return queryCreator.createQuery(getDynamicSort(accessor));
 		});
 	}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/repository/query/R2dbcQueryCreator.java

@@ -18,21 +18,18 @@ package org.springframework.data.r2dbc.repository.query;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.stream.Collectors;
 
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
 import org.springframework.data.r2dbc.core.StatementMapper;
+import org.springframework.data.relational.repository.Lock;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.query.Criteria;
-import org.springframework.data.relational.core.sql.Column;
-import org.springframework.data.relational.core.sql.Expression;
-import org.springframework.data.relational.core.sql.Expressions;
-import org.springframework.data.relational.core.sql.Functions;
-import org.springframework.data.relational.core.sql.SqlIdentifier;
-import org.springframework.data.relational.core.sql.Table;
+import org.springframework.data.relational.core.sql.*;
 import org.springframework.data.relational.repository.query.RelationalEntityMetadata;
 import org.springframework.data.relational.repository.query.RelationalParameterAccessor;
 import org.springframework.data.relational.repository.query.RelationalQueryCreator;
@@ -48,6 +45,7 @@ import org.springframework.r2dbc.core.PreparedOperation;
  * @author Mark Paluch
  * @author Mingyuan Wu
  * @author Myeonghyeon Lee
+ * @author Diego Krupitza
  * @since 1.1
  */
 class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
@@ -58,6 +56,7 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 	private final RelationalEntityMetadata<?> entityMetadata;
 	private final List<String> projectedProperties;
 	private final Class<?> entityToRead;
+	private final Optional<Lock> lock;
 
 	/**
 	 * Creates new instance of this class with the given {@link PartTree}, {@link ReactiveDataAccessStrategy},
@@ -71,7 +70,7 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 	 */
 	public R2dbcQueryCreator(PartTree tree, ReactiveDataAccessStrategy dataAccessStrategy,
 			RelationalEntityMetadata<?> entityMetadata, RelationalParameterAccessor accessor,
-			List<String> projectedProperties) {
+			List<String> projectedProperties, Optional<Lock> lock) {
 		super(tree, accessor);
 
 		this.tree = tree;
@@ -81,6 +80,7 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 		this.entityMetadata = entityMetadata;
 		this.projectedProperties = projectedProperties;
 		this.entityToRead = entityMetadata.getTableEntity().getType();
+		this.lock = lock;
 	}
 
 	/**
@@ -138,6 +138,10 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 			selectSpec = selectSpec.distinct();
 		}
 
+		if (this.lock.isPresent()) {
+			selectSpec = selectSpec.lock(this.lock.get().value());
+		}
+
 		return statementMapper.getMappedObject(selectSpec);
 	}
 
@@ -154,15 +158,15 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 			for (String projectedProperty : projectedProperties) {
 
 				RelationalPersistentProperty property = entity.getPersistentProperty(projectedProperty);
-				Column column = table.column(property != null ? property.getColumnName() : SqlIdentifier.unquoted(projectedProperty));
+				Column column = table
+						.column(property != null ? property.getColumnName() : SqlIdentifier.unquoted(projectedProperty));
 				expressions.add(column);
 			}
 
 		} else if (tree.isExistsProjection()) {
 
-			expressions = dataAccessStrategy.getIdentifierColumns(entityToRead).stream()
-				.map(table::column)
-				.collect(Collectors.toList());
+			expressions = dataAccessStrategy.getIdentifierColumns(entityToRead).stream().map(table::column)
+					.collect(Collectors.toList());
 		} else if (tree.isCountProjection()) {
 
 			Expression countExpression = entityMetadata.getTableEntity().hasIdProperty()
@@ -171,9 +175,8 @@ class R2dbcQueryCreator extends RelationalQueryCreator<PreparedOperation<?>> {
 
 			expressions = Collections.singletonList(Functions.count(countExpression));
 		} else {
-			expressions = dataAccessStrategy.getAllColumns(entityToRead).stream()
-				.map(table::column)
-				.collect(Collectors.toList());
+			expressions = dataAccessStrategy.getAllColumns(entityToRead).stream().map(table::column)
+					.collect(Collectors.toList());
 		}
 
 		return expressions.toArray(new Expression[0]);

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/repository/query/R2dbcQueryMethod.java

@@ -28,6 +28,7 @@ import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.projection.ProjectionFactory;
+import org.springframework.data.relational.repository.Lock;
 import org.springframework.data.r2dbc.repository.Modifying;
 import org.springframework.data.r2dbc.repository.Query;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
@@ -53,6 +54,7 @@ import org.springframework.util.ClassUtils;
  *
  * @author Mark Paluch
  * @author Stephen Cohen
+ * @author Diego Krupitza
  */
 public class R2dbcQueryMethod extends QueryMethod {
 
@@ -66,6 +68,7 @@ public class R2dbcQueryMethod extends QueryMethod {
 	private final Optional<Query> query;
 	private final boolean modifying;
 	private final Lazy<Boolean> isCollectionQuery;
+	private final Optional<Lock> lock;
 
 	private @Nullable RelationalEntityMetadata<?> metadata;
 
@@ -113,11 +116,11 @@ public class R2dbcQueryMethod extends QueryMethod {
 			}
 		}
 
-		this.query = Optional.ofNullable(
-				AnnotatedElementUtils.findMergedAnnotation(method, Query.class));
+		this.query = Optional.ofNullable(AnnotatedElementUtils.findMergedAnnotation(method, Query.class));
 		this.modifying = AnnotatedElementUtils.hasAnnotation(method, Modifying.class);
 		this.isCollectionQuery = Lazy.of(() -> (!(isPageQuery() || isSliceQuery())
 				&& ReactiveWrappers.isMultiValueType(metadata.getReturnType(method).getType())) || super.isCollectionQuery());
+		this.lock = Optional.ofNullable(AnnotatedElementUtils.findMergedAnnotation(method, Lock.class));
 	}
 
 	/* (non-Javadoc)
@@ -144,6 +147,22 @@ public class R2dbcQueryMethod extends QueryMethod {
 		return modifying;
 	}
 
+	/**
+	 * @return is a {@link Lock} annotation present or not.
+	 */
+	public boolean hasLockMode() {
+		return this.lock.isPresent();
+	}
+
+	/**
+	 * Looks up the {@link Lock} annotation from the query method.
+	 *
+	 * @return the {@link Optional} wrapped {@link Lock} annotation.
+	 */
+	Optional<Lock> getLock() {
+		return this.lock;
+	}
+
 	/*
 	 * All reactive query methods are streaming queries.
 	 * (non-Javadoc)
@@ -166,8 +185,7 @@ public class R2dbcQueryMethod extends QueryMethod {
 			Class<?> returnedObjectType = getReturnedObjectType();
 			Class<?> domainClass = getDomainClass();
 
-			if (ClassUtils.isPrimitiveOrWrapper(returnedObjectType)
-					|| ReflectionUtils.isVoid(returnedObjectType)) {
+			if (ClassUtils.isPrimitiveOrWrapper(returnedObjectType) || ReflectionUtils.isVoid(returnedObjectType)) {
 
 				this.metadata = new SimpleRelationalEntityMetadata<>((Class<Object>) domainClass,
 						mappingContext.getRequiredPersistentEntity(domainClass));
@@ -214,8 +232,8 @@ public class R2dbcQueryMethod extends QueryMethod {
 	}
 
 	/**
-	 * Returns the required query string declared in a {@link Query} annotation
-	 * or throws {@link IllegalStateException} if neither the annotation found nor the attribute was specified.
+	 * Returns the required query string declared in a {@link Query} annotation or throws {@link IllegalStateException} if
+	 * neither the annotation found nor the attribute was specified.
 	 *
 	 * @return the query string.
 	 * @throws IllegalStateException in case query method has no annotated query.
@@ -235,8 +253,7 @@ public class R2dbcQueryMethod extends QueryMethod {
 	}
 
 	/**
-	 * @return {@literal true} if the {@link Method} is annotated with
-	 *         {@link Query}.
+	 * @return {@literal true} if the {@link Method} is annotated with {@link Query}.
 	 */
 	public boolean hasAnnotatedQuery() {
 		return getQueryAnnotation().isPresent();

spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/StatementMapperUnitTests.java

@@ -26,6 +26,7 @@ import org.springframework.data.r2dbc.core.StatementMapper.UpdateSpec;
 import org.springframework.data.r2dbc.dialect.PostgresDialect;
 import org.springframework.data.relational.core.query.Criteria;
 import org.springframework.data.relational.core.query.Update;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.r2dbc.core.PreparedOperation;
 import org.springframework.r2dbc.core.binding.BindTarget;
 
@@ -33,6 +34,7 @@ import org.springframework.r2dbc.core.binding.BindTarget;
  * Unit tests for {@link DefaultStatementMapper}.
  *
  * @author Mark Paluch
+ * @author Diego Krupitza
  */
 class StatementMapperUnitTests {
 
@@ -80,4 +82,26 @@ class StatementMapperUnitTests {
 		assertThat(preparedOperation.toQuery())
 				.isEqualTo("SELECT table.* FROM table ORDER BY table.id DESC LIMIT 2 OFFSET 2");
 	}
+
+	@Test // gh-1041
+	void shouldMapSelectWithSharedLock() {
+
+		StatementMapper.SelectSpec selectSpec = StatementMapper.SelectSpec.create("table").withProjection("*")
+				.lock(LockMode.PESSIMISTIC_READ);
+
+		PreparedOperation<?> preparedOperation = mapper.getMappedObject(selectSpec);
+
+		assertThat(preparedOperation.toQuery()).isEqualTo("SELECT table.* FROM table FOR SHARE OF table");
+	}
+
+	@Test // gh-1041
+	void shouldMapSelectWithWriteLock() {
+
+		StatementMapper.SelectSpec selectSpec = StatementMapper.SelectSpec.create("table").withProjection("*")
+				.lock(LockMode.PESSIMISTIC_WRITE);
+
+		PreparedOperation<?> preparedOperation = mapper.getMappedObject(selectSpec);
+
+		assertThat(preparedOperation.toQuery()).isEqualTo("SELECT table.* FROM table FOR UPDATE OF table");
+	}
 }

spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/AbstractR2dbcRepositoryIntegrationTests.java

@@ -23,6 +23,8 @@ import lombok.Getter;
 import lombok.NoArgsConstructor;
 import lombok.Setter;
 import lombok.Value;
+import org.springframework.data.relational.core.sql.LockMode;
+import org.springframework.data.relational.repository.Lock;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
@@ -58,6 +60,7 @@ import org.springframework.transaction.reactive.TransactionalOperator;
  *
  * @author Mark Paluch
  * @author Manousos Mathioudakis
+ * @author Diego Krupitza
  */
 public abstract class AbstractR2dbcRepositoryIntegrationTests extends R2dbcIntegrationTestSupport {
 
@@ -380,6 +383,33 @@ public abstract class AbstractR2dbcRepositoryIntegrationTests extends R2dbcInteg
 				.verifyComplete();
 	}
 
+	@Test
+	void getAllByNameWithWriteLock() {
+
+		shouldInsertNewItems();
+
+		repository.getAllByName("SCHAUFELRADBAGGER") //
+				.as(StepVerifier::create) //
+				.consumeNextWith(actual -> {
+					assertThat(actual.getName()).isEqualTo("SCHAUFELRADBAGGER");
+				}) //
+				.verifyComplete();
+	}
+
+	@Test
+	void findByNameAndFlagWithReadLock() {
+
+		shouldInsertNewItems();
+
+		repository.findByNameAndFlag("SCHAUFELRADBAGGER", true) //
+				.as(StepVerifier::create) //
+				.consumeNextWith(actual -> {
+					assertThat(actual.getName()).isEqualTo("SCHAUFELRADBAGGER");
+					assertThat(actual.isFlag()).isTrue();
+				}) //
+				.verifyComplete();
+	}
+
 	private static Object getCount(Map<String, Object> map) {
 		return map.getOrDefault("count", map.get("COUNT"));
 	}
@@ -393,6 +423,12 @@ public abstract class AbstractR2dbcRepositoryIntegrationTests extends R2dbcInteg
 	@NoRepositoryBean
 	interface LegoSetRepository extends ReactiveCrudRepository<LegoSet, Integer> {
 
+		@Lock(LockMode.PESSIMISTIC_WRITE)
+		Flux<LegoSet> getAllByName(String name);
+
+		@Lock(LockMode.PESSIMISTIC_READ)
+		Flux<LegoSet> findByNameAndFlag(String name, Boolean flag);
+
 		Flux<LegoSet> findByNameContains(String name);
 
 		Flux<LegoSet> findFirst10By();

spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/query/PartTreeR2dbcQueryUnitTests.java

@@ -36,7 +36,6 @@ import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
-
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
@@ -48,8 +47,10 @@ import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
 import org.springframework.data.r2dbc.dialect.DialectResolver;
 import org.springframework.data.r2dbc.dialect.R2dbcDialect;
 import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
+import org.springframework.data.relational.repository.Lock;
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.Table;
+import org.springframework.data.relational.core.sql.LockMode;
 import org.springframework.data.relational.repository.query.RelationalParametersParameterAccessor;
 import org.springframework.data.repository.Repository;
 import org.springframework.data.repository.core.support.DefaultRepositoryMetadata;
@@ -64,6 +65,7 @@ import org.springframework.r2dbc.core.binding.BindTarget;
  * @author Mark Paluch
  * @author Mingyuan Wu
  * @author Myeonghyeon Lee
+ * @author Diego Krupitza
  */
 @ExtendWith(MockitoExtension.class)
 @MockitoSettings(strictness = Strictness.LENIENT)
@@ -660,8 +662,38 @@ class PartTreeR2dbcQueryUnitTests {
 			dataAccessStrategy);
 		PreparedOperation<?> query = createQuery(queryMethod, r2dbcQuery, "John");
 
-		assertThat(query.get())
-			.isEqualTo("SELECT COUNT(users.id) FROM " + TABLE + " WHERE " + TABLE + ".first_name = $1");
+		assertThat(query.get()) //
+				.isEqualTo("SELECT COUNT(users.id) FROM " + TABLE + " WHERE " + TABLE + ".first_name = $1");
+	}
+
+	@Test // gh-1041
+	void createQueryWithPessimisticWriteLock() throws Exception {
+
+		R2dbcQueryMethod queryMethod = getQueryMethod("findAllByFirstNameAndLastName", String.class, String.class);
+		PartTreeR2dbcQuery r2dbcQuery = new PartTreeR2dbcQuery(queryMethod, operations, r2dbcConverter, dataAccessStrategy);
+
+		String firstname = "Diego";
+		String lastname = "Krupitza";
+
+		PreparedOperation<?> query = createQuery(queryMethod, r2dbcQuery, firstname, lastname);
+
+		assertThat(query.get()).isEqualTo(
+				"SELECT users.id, users.first_name, users.last_name, users.date_of_birth, users.age, users.active FROM users WHERE users.first_name = $1 AND (users.last_name = $2) FOR UPDATE OF users");
+	}
+
+	@Test // gh-1041
+	void createQueryWithPessimisticReadLock() throws Exception {
+
+		R2dbcQueryMethod queryMethod = getQueryMethod("findAllByFirstNameAndAge", String.class, Integer.class);
+		PartTreeR2dbcQuery r2dbcQuery = new PartTreeR2dbcQuery(queryMethod, operations, r2dbcConverter, dataAccessStrategy);
+
+		String firstname = "Diego";
+		Integer age = 22;
+
+		PreparedOperation<?> query = createQuery(queryMethod, r2dbcQuery, firstname, age);
+
+		assertThat(query.get()).isEqualTo(
+				"SELECT users.id, users.first_name, users.last_name, users.date_of_birth, users.age, users.active FROM users WHERE users.first_name = $1 AND (users.age = $2) FOR SHARE OF users");
 	}
 
 	private PreparedOperation<?> createQuery(R2dbcQueryMethod queryMethod, PartTreeR2dbcQuery r2dbcQuery,
@@ -687,6 +719,12 @@ class PartTreeR2dbcQueryUnitTests {
 	@SuppressWarnings("ALL")
 	interface UserRepository extends Repository<User, Long> {
 
+		@Lock(LockMode.PESSIMISTIC_WRITE)
+		Flux<User> findAllByFirstNameAndLastName(String firstName, String lastName);
+
+		@Lock(LockMode.PESSIMISTIC_READ)
+		Flux<User> findAllByFirstNameAndAge(String firstName, Integer age);
+
 		Flux<User> findAllByFirstName(String firstName);
 
 		Flux<User> findAllByLastNameAndFirstName(String lastName, String firstName);

spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/query/R2dbcQueryMethodUnitTests.java

@@ -18,6 +18,8 @@ package org.springframework.data.r2dbc.repository.query;
 import static org.assertj.core.api.Assertions.*;
 
 import kotlin.Unit;
+import org.springframework.data.relational.repository.Lock;
+import org.springframework.data.relational.core.sql.LockMode;
 import reactor.core.publisher.Mono;
 
 import java.lang.annotation.Retention;
@@ -47,6 +49,7 @@ import org.springframework.data.repository.core.support.DefaultRepositoryMetadat
  *
  * @author Mark Paluch
  * @author Stephen Cohen
+ * @author Diego Krupitza
  */
 class R2dbcQueryMethodUnitTests {
 
@@ -141,6 +144,32 @@ class R2dbcQueryMethodUnitTests {
 		assertThat(method.getEntityInformation().getJavaType()).isAssignableFrom(Contact.class);
 	}
 
+	@Test // GH-1041
+	void returnsQueryMethodWithCorrectLockTypeWriteLock() throws Exception {
+
+		R2dbcQueryMethod queryMethodWithWriteLock = queryMethod(PersonRepository.class, "queryMethodWithWriteLock");
+
+		assertThat(queryMethodWithWriteLock.getLock()).isPresent();
+		assertThat(queryMethodWithWriteLock.getLock().get().value()).isEqualTo(LockMode.PESSIMISTIC_WRITE);
+	}
+
+	@Test // GH-1041
+	void returnsQueryMethodWithCorrectLockTypeReadLock() throws Exception {
+
+		R2dbcQueryMethod queryMethodWithReadLock = queryMethod(PersonRepository.class, "queryMethodWithReadLock");
+
+		assertThat(queryMethodWithReadLock.getLock()).isPresent();
+		assertThat(queryMethodWithReadLock.getLock().get().value()).isEqualTo(LockMode.PESSIMISTIC_READ);
+	}
+
+	@Test // GH-1041
+	void returnsQueryMethodWithCorrectLockTypeNoLock() throws Exception {
+
+		R2dbcQueryMethod queryMethodWithWriteLock = queryMethod(SampleRepository.class, "methodReturningAnInterface");
+
+		assertThat(queryMethodWithWriteLock.getLock()).isEmpty();
+	}
+
 	private R2dbcQueryMethod queryMethod(Class<?> repository, String name, Class<?>... parameters) throws Exception {
 
 		Method method = repository.getMethod(name, parameters);
@@ -150,6 +179,12 @@ class R2dbcQueryMethodUnitTests {
 
 	interface PersonRepository extends Repository<Contact, Long> {
 
+		@Lock(LockMode.PESSIMISTIC_WRITE)
+		Mono<Contact> queryMethodWithWriteLock();
+
+		@Lock(LockMode.PESSIMISTIC_READ)
+		Mono<Contact> queryMethodWithReadLock();
+
 		Mono<Contact> findMonoByLastname(String lastname, Pageable pageRequest);
 
 		Mono<Page<Contact>> findMonoPageByLastname(String lastname, Pageable pageRequest);