feat(spanner): support select for update

docs/src/main/asciidoc/spanner.adoc

@@ -1081,6 +1081,42 @@ Parameters can also be of type `Struct` or POJOs.
 If a POJO is given as a parameter, it will be converted to a `Struct` with the same type-conversion logic as used to create write mutations.
 Comparisons using Struct parameters are limited to https://cloud.google.com/spanner/docs/data-types#limited-comparisons-for-struct[what is available with Cloud Spanner].
 
+==== Locking rows with `FOR UPDATE`
+
+For workloads with high write contention, a query can acquire exclusive locks by setting
+`forUpdate = true` on the `@Query` annotation.
+The annotation can be used with a derived query without specifying a SQL string:
+
+[source,java]
+----
+public interface TradeRepository extends SpannerRepository<Trade, Key> {
+
+	@Query(forUpdate = true)
+	Optional<Trade> findBySymbol(String symbol);
+}
+----
+
+To retrieve an entity by its ID with exclusive locks, use `findByIdForUpdate`:
+
+[source,java]
+----
+@Transactional(transactionManager = "spannerTransactionManager")
+public void updateTrade(Key tradeId) {
+	Trade trade = tradeRepository.findByIdForUpdate(tradeId).orElseThrow();
+	trade.setAction("SELL");
+	tradeRepository.save(trade);
+}
+----
+
+Both forms must be executed in a read-write transaction.
+The generated queries also apply `FOR UPDATE` when loading eager or lazy interleaved children.
+For programmatic queries with `SpannerTemplate`, set `forUpdate` on `SpannerQueryOptions`.
+
+`FOR UPDATE` can reduce transaction aborts when concurrent transactions read and update the same
+data, but conflicting transactions wait for the locks and can therefore reduce throughput.
+See https://cloud.google.com/spanner/docs/use-select-for-update[Use SELECT FOR UPDATE] for details
+and restrictions.
+
 
 ==== Query methods by convention
 

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerQueryOptions.java

@@ -30,6 +30,8 @@
  */
 public class SpannerQueryOptions extends AbstractSpannerRequestOptions<QueryOption> {
 
+  private boolean forUpdate;
+
   /**
    * Constructor to create an instance. Use the extension-style add/set functions to add options and
    * settings.
@@ -44,6 +46,22 @@ public SpannerQueryOptions addQueryOption(QueryOption queryOption) {
     return this;
   }
 
+  public boolean isForUpdate() {
+    return this.forUpdate;
+  }
+
+  /**
+   * Sets whether the query should acquire exclusive locks on the selected rows. Cloud Spanner only
+   * supports {@code FOR UPDATE} in read-write transactions.
+   *
+   * @param forUpdate whether {@code FOR UPDATE} is enabled
+   * @return this options instance
+   */
+  public SpannerQueryOptions setForUpdate(boolean forUpdate) {
+    this.forUpdate = forUpdate;
+    return this;
+  }
+
   @Override
   public SpannerQueryOptions setIncludeProperties(Set<String> includeProperties) {
     super.setIncludeProperties(includeProperties);

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java

@@ -291,8 +291,10 @@ public <T> List<T> queryAll(Class<T> entityClass, SpannerPageableQueryOptions op
     return query(
         entityClass,
         SpannerStatementQueryExecutor.buildStatementFromSqlWithArgs(
-            SpannerStatementQueryExecutor.applySortingPagingQueryOptions(
-                entityClass, options, sql, this.mappingContext, false),
+            SpannerStatementQueryExecutor.applyForUpdate(
+                SpannerStatementQueryExecutor.applySortingPagingQueryOptions(
+                    entityClass, options, sql, this.mappingContext, false),
+                options.isForUpdate()),
             null,
             null,
             null,
@@ -502,6 +504,10 @@ public <T> T performReadOnlyTransaction(
   }
 
   public ResultSet executeQuery(Statement statement, SpannerQueryOptions options) {
+    if (options != null && options.isForUpdate() && !isReadWriteTransactionActive()) {
+      throw new SpannerDataException(
+          "FOR UPDATE queries must be executed in a read-write transaction.");
+    }
     ResultSet resultSet = performQuery(statement, options);
     if (LOGGER.isDebugEnabled()) {
       String message;
@@ -640,28 +646,45 @@ private <T> List<T> queryAndResolveChildren(
         executeQuery(statement, options),
         entityClass,
         (options != null) ? options.getIncludeProperties() : null,
-        options != null && options.isAllowPartialRead());
+        options != null && options.isAllowPartialRead(),
+        options != null && options.isForUpdate());
   }
 
   private <T> List<T> mapToListAndResolveChildren(
       ResultSet resultSet,
       Class<T> entityClass,
       Set<String> includeProperties,
       boolean allowMissingColumns) {
+    return mapToListAndResolveChildren(
+        resultSet, entityClass, includeProperties, allowMissingColumns, false);
+  }
+
+  private <T> List<T> mapToListAndResolveChildren(
+      ResultSet resultSet,
+      Class<T> entityClass,
+      Set<String> includeProperties,
+      boolean allowMissingColumns,
+      boolean forUpdate) {
     return resolveChildEntities(
         this.spannerEntityProcessor.mapToList(
             resultSet, entityClass, includeProperties, allowMissingColumns),
-        includeProperties);
+        includeProperties,
+        forUpdate);
   }
 
   private <T> List<T> resolveChildEntities(List<T> entities, Set<String> includeProperties) {
+    return resolveChildEntities(entities, includeProperties, false);
+  }
+
+  private <T> List<T> resolveChildEntities(
+      List<T> entities, Set<String> includeProperties, boolean forUpdate) {
     for (Object entity : entities) {
-      resolveChildEntity(entity, includeProperties);
+      resolveChildEntity(entity, includeProperties, forUpdate);
     }
     return entities;
   }
 
-  private void resolveChildEntity(Object entity, Set<String> includeProperties) {
+  private void resolveChildEntity(Object entity, Set<String> includeProperties, boolean forUpdate) {
     SpannerPersistentEntity<?> spannerPersistentEntity =
         this.mappingContext.getPersistentEntityOrFail(entity.getClass());
 
@@ -675,7 +698,7 @@ private void resolveChildEntity(Object entity, Set<String> includeProperties) {
           // an interleaved property can only be List
           List propertyValue = (List) accessor.getProperty(spannerPersistentProperty);
           if (propertyValue != null) {
-            resolveChildEntities(propertyValue, null);
+            resolveChildEntities(propertyValue, null, forUpdate);
             return;
           }
           Class<?> childType = spannerPersistentProperty.getColumnInnerType();
@@ -688,8 +711,9 @@ private void resolveChildEntity(Object entity, Set<String> includeProperties) {
                           this.spannerSchemaUtils.getKey(entity),
                           spannerPersistentProperty,
                           this.spannerEntityProcessor.getWriteConverter(),
-                          this.mappingContext),
-                      null);
+                          this.mappingContext,
+                          forUpdate),
+                      new SpannerQueryOptions().setForUpdate(forUpdate));
 
           accessor.setProperty(
               spannerPersistentProperty,
@@ -718,6 +742,10 @@ private TransactionContext getTransactionContext() {
     return null;
   }
 
+  private boolean isReadWriteTransactionActive() {
+    return this instanceof ReadWriteTransactionSpannerTemplate || getTransactionContext() != null;
+  }
+
   private <A> A doWithOrWithoutTransactionContext(
       Function<TransactionContext, A> funcWithTransactionContext,
       Supplier<A> funcWithoutTransactionContext) {

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/repository/SpannerRepository.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spring.data.spanner.repository;
 
 import com.google.cloud.spring.data.spanner.core.SpannerOperations;
+import java.util.Optional;
 import java.util.function.Function;
 import org.springframework.data.repository.CrudRepository;
 import org.springframework.data.repository.PagingAndSortingRepository;
@@ -58,4 +59,13 @@ public interface SpannerRepository<T, I>
    * @return the final result of the transaction.
    */
   <A> A performReadOnlyTransaction(Function<SpannerRepository<T, I>, A> operations);
+
+  /**
+   * Retrieves an entity by its id and acquires exclusive locks on the selected row and its
+   * interleaved children. This method must be called in a read-write transaction.
+   *
+   * @param id the entity id
+   * @return the entity, or empty if none was found
+   */
+  Optional<T> findByIdForUpdate(I id);
 }

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/repository/query/PartTreeSpannerQuery.java

@@ -65,7 +65,8 @@ protected List executeRawResult(Object[] parameters) {
           paramAccessor,
           getQueryMethod().getQueryMethod().getParameters(),
           this.spannerTemplate,
-          this.spannerMappingContext);
+          this.spannerMappingContext,
+          false);
     }
     if (this.tree.isDelete()) {
       return this.spannerTemplate.performReadWriteTransaction(getDeleteFunction(parameters));
@@ -76,7 +77,8 @@ protected List executeRawResult(Object[] parameters) {
         paramAccessor,
         getQueryMethod().getQueryMethod().getParameters(),
         this.spannerTemplate,
-        this.spannerMappingContext);
+        this.spannerMappingContext,
+        this.queryMethod.isForUpdate());
   }
 
   private Function<SpannerTemplate, List> getDeleteFunction(Object[] parameters) {
@@ -90,7 +92,8 @@ private Function<SpannerTemplate, List> getDeleteFunction(Object[] parameters) {
               paramAccessor,
               getQueryMethod().getQueryMethod().getParameters(),
               transactionTemplate,
-              this.spannerMappingContext);
+              this.spannerMappingContext,
+              false);
       transactionTemplate.deleteAll(entitiesToDelete);
 
       List result = null;

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/repository/query/Query.java

@@ -53,4 +53,12 @@
    *     method is executed as a DML query.
    */
   boolean dmlStatement() default false;
+
+  /**
+   * Indicates if the query should acquire exclusive locks on the selected rows. This option is
+   * valid only for select queries executed in a read-write transaction.
+   *
+   * @return {@code true} if {@code FOR UPDATE} should be applied.
+   */
+  boolean forUpdate() default false;
 }

spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/repository/query/SpannerQueryMethod.java

@@ -85,4 +85,9 @@ Method getQueryMethod() {
   Query getQueryAnnotation() {
     return AnnotatedElementUtils.findMergedAnnotation(this.queryMethod, Query.class);
   }
+
+  boolean isForUpdate() {
+    Query query = getQueryAnnotation();
+    return query != null && query.forUpdate();
+  }
 }