feat: add LockHint feature

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -58,6 +58,7 @@
 import java.util.Map;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.logging.Logger;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
 
@@ -67,6 +68,7 @@
  */
 abstract class AbstractReadContext
     implements ReadContext, AbstractResultSet.Listener, SessionTransaction {
+  private static final Logger logger = Logger.getLogger(AbstractReadContext.class.getName());
 
   abstract static class Builder<B extends Builder<?, T>, T extends AbstractReadContext> {
     private SessionImpl session;
@@ -951,6 +953,15 @@ ResultSet readInternalWithOptions(
     } else if (defaultDirectedReadOptions != null) {
       builder.setDirectedReadOptions(defaultDirectedReadOptions);
     }
+    if (readOptions.hasLockHint()) {
+      if (isReadOnly()) {
+        logger.warning(
+            "Lock hint is only supported for ReadWrite transactions. "
+                + "Overriding lock hint to default unspecified.");
+      } else {
+        builder.setLockHint(readOptions.lockHint());
+      }
+    }
     final int prefetchChunks =
         readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
     ResumableStreamIterator stream =

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java

@@ -18,6 +18,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.spanner.v1.DirectedReadOptions;
+import com.google.spanner.v1.ReadRequest.LockHint;
 import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import java.io.Serializable;
@@ -75,6 +76,25 @@ public static RpcOrderBy fromProto(OrderBy proto) {
     }
   }
 
+  public enum RpcLockHint {
+    UNSPECIFIED(LockHint.LOCK_HINT_UNSPECIFIED),
+    SHARED(LockHint.LOCK_HINT_SHARED),
+    EXCLUSIVE(LockHint.LOCK_HINT_EXCLUSIVE);
+
+    private final LockHint proto;
+
+    RpcLockHint(LockHint proto) {
+      this.proto = Preconditions.checkNotNull(proto);
+    }
+
+    public static RpcLockHint fromProto(LockHint proto) {
+      for (RpcLockHint e : RpcLockHint.values()) {
+        if (e.proto.equals(proto)) return e;
+      }
+      return RpcLockHint.UNSPECIFIED;
+    }
+  }
+
   /** Marker interface to mark options applicable to both Read and Query operations */
   public interface ReadAndQueryOption extends ReadOption, QueryOption {}
 
@@ -160,6 +180,10 @@ public static ReadOption orderBy(RpcOrderBy orderBy) {
     return new OrderByOption(orderBy);
   }
 
+  public static ReadOption lockHint(RpcLockHint orderBy) {
+    return new LockHintOption(orderBy);
+  }
+
   /**
    * Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
    * PartialResultSet} chunks for read and query. The data size of each chunk depends on the server
@@ -469,6 +493,7 @@ void appendToOptions(Options options) {
   private DirectedReadOptions directedReadOptions;
   private DecodeMode decodeMode;
   private RpcOrderBy orderBy;
+  private RpcLockHint lockHint;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -605,6 +630,14 @@ OrderBy orderBy() {
     return orderBy == null ? null : orderBy.proto;
   }
 
+  boolean hasLockHint() {
+    return lockHint != null;
+  }
+
+  LockHint lockHint() {
+    return lockHint == null ? null : lockHint.proto;
+  }
+
   @Override
   public String toString() {
     StringBuilder b = new StringBuilder();
@@ -661,6 +694,9 @@ public String toString() {
     if (orderBy != null) {
       b.append("orderBy: ").append(orderBy).append(' ');
     }
+    if (lockHint != null) {
+      b.append("lockHint: ").append(lockHint).append(' ');
+    }
     return b.toString();
   }
 
@@ -700,7 +736,8 @@ public boolean equals(Object o) {
         && Objects.equals(withExcludeTxnFromChangeStreams(), that.withExcludeTxnFromChangeStreams())
         && Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
         && Objects.equals(directedReadOptions(), that.directedReadOptions())
-        && Objects.equals(orderBy(), that.orderBy());
+        && Objects.equals(orderBy(), that.orderBy())
+        && Objects.equals(lockHint(), that.lockHint());
   }
 
   @Override
@@ -760,6 +797,9 @@ public int hashCode() {
     if (orderBy != null) {
       result = 31 * result + orderBy.hashCode();
     }
+    if (lockHint != null) {
+      result = 31 * result + lockHint.hashCode();
+    }
     return result;
   }
 
@@ -853,6 +893,19 @@ void appendToOptions(Options options) {
     }
   }
 
+  static class LockHintOption extends InternalOption implements ReadOption {
+    private final RpcLockHint lockHint;
+
+    LockHintOption(RpcLockHint lockHint) {
+      this.lockHint = lockHint;
+    }
+
+    @Override
+    void appendToOptions(Options options) {
+      options.lockHint = lockHint;
+    }
+  }
+
   static final class DataBoostQueryOption extends InternalOption implements ReadAndQueryOption {
 
     private final Boolean dataBoostEnabled;

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java

@@ -34,6 +34,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.ReadRequest.LockHint;
 import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions;
 import com.google.spanner.v1.RequestOptions.Priority;
@@ -241,6 +242,21 @@ public void testGetReadRequestBuilderWithOrderBy() {
     assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
   }
 
+  @Test
+  public void testGetReadRequestBuilderWithLockHint() {
+    ReadRequest request =
+        ReadRequest.newBuilder()
+            .setSession(
+                SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]").toString())
+            .setTransaction(TransactionSelector.newBuilder().build())
+            .setTable("table110115790")
+            .setIndex("index100346066")
+            .addAllColumns(new ArrayList<String>())
+            .setLockHintValue(2)
+            .build();
+    assertEquals(LockHint.LOCK_HINT_EXCLUSIVE, request.getLockHint());
+  }
+
   @Test
   public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
     ExecuteBatchDmlRequest.Builder request =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -52,6 +52,7 @@
 import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Options.RpcLockHint;
 import com.google.cloud.spanner.Options.RpcOrderBy;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Options.TransactionOption;
@@ -90,6 +91,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.ReadRequest.LockHint;
 import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import com.google.spanner.v1.ResultSetMetadata;
@@ -1754,6 +1756,53 @@ public void testExecuteReadWithOrderByOption() {
     assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
   }
 
+  @Test
+  public void testUnsupportedTransactionWithLockHintOption() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    try (ResultSet resultSet =
+        client
+            .singleUse()
+            .read(
+                READ_TABLE_NAME,
+                KeySet.singleKey(Key.of(1L)),
+                READ_COLUMN_NAMES,
+                Options.lockHint(RpcLockHint.EXCLUSIVE))) {
+      consumeResults(resultSet);
+    }
+
+    List<ReadRequest> requests = mockSpanner.getRequestsOfType(ReadRequest.class);
+    assertThat(requests).hasSize(1);
+    ReadRequest request = requests.get(0);
+    // lock hint is only supported in  ReadWriteTransaction
+    assertEquals(LockHint.LOCK_HINT_UNSPECIFIED, request.getLockHint());
+  }
+
+  @Test
+  public void testReadWriteTransactionWithLockHint() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+
+    TransactionRunner runner = client.readWriteTransaction();
+    runner.run(
+        transaction -> {
+          try (ResultSet resultSet =
+              transaction.read(
+                  READ_TABLE_NAME,
+                  KeySet.singleKey(Key.of(1L)),
+                  READ_COLUMN_NAMES,
+                  Options.lockHint(RpcLockHint.EXCLUSIVE))) {
+            consumeResults(resultSet);
+          }
+          return null;
+        });
+
+    List<ReadRequest> requests = mockSpanner.getRequestsOfType(ReadRequest.class);
+    assertThat(requests).hasSize(1);
+    ReadRequest request = requests.get(0);
+    assertEquals(LockHint.LOCK_HINT_EXCLUSIVE, request.getLockHint());
+  }
+
   @Test
   public void testExecuteReadWithDirectedReadOptions() {
     DatabaseClient client =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java

@@ -25,11 +25,13 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.Options.RpcLockHint;
 import com.google.cloud.spanner.Options.RpcOrderBy;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
 import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
+import com.google.spanner.v1.ReadRequest.LockHint;
 import com.google.spanner.v1.ReadRequest.OrderBy;
 import com.google.spanner.v1.RequestOptions.Priority;
 import org.junit.Test;
@@ -83,7 +85,8 @@ public void allOptionsPresent() {
             Options.prefetchChunks(1),
             Options.dataBoostEnabled(true),
             Options.directedRead(DIRECTED_READ_OPTIONS),
-            Options.orderBy(RpcOrderBy.NO_ORDER));
+            Options.orderBy(RpcOrderBy.NO_ORDER),
+            Options.lockHint(Options.RpcLockHint.SHARED));
     assertThat(options.hasLimit()).isTrue();
     assertThat(options.limit()).isEqualTo(10);
     assertThat(options.hasPrefetchChunks()).isTrue();
@@ -92,6 +95,7 @@ public void allOptionsPresent() {
     assertTrue(options.dataBoostEnabled());
     assertTrue(options.hasDirectedReadOptions());
     assertTrue(options.hasOrderBy());
+    assertTrue(options.hasLockHint());
     assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
   }
 
@@ -107,6 +111,7 @@ public void allOptionsAbsent() {
     assertThat(options.hasDataBoostEnabled()).isFalse();
     assertThat(options.hasDirectedReadOptions()).isFalse();
     assertThat(options.hasOrderBy()).isFalse();
+    assertThat(options.hasLockHint()).isFalse();
     assertNull(options.withExcludeTxnFromChangeStreams());
     assertThat(options.toString()).isEqualTo("");
     assertThat(options.equals(options)).isTrue();
@@ -189,7 +194,8 @@ public void readOptionsTest() {
             Options.tag(tag),
             Options.dataBoostEnabled(true),
             Options.directedRead(DIRECTED_READ_OPTIONS),
-            Options.orderBy(RpcOrderBy.NO_ORDER));
+            Options.orderBy(RpcOrderBy.NO_ORDER),
+            Options.lockHint(RpcLockHint.SHARED));
 
     assertThat(options.toString())
         .isEqualTo(
@@ -207,11 +213,15 @@ public void readOptionsTest() {
                 + " "
                 + "orderBy: "
                 + RpcOrderBy.NO_ORDER
+                + " "
+                + "lockHint: "
+                + RpcLockHint.SHARED
                 + " ");
     assertThat(options.tag()).isEqualTo(tag);
     assertEquals(dataBoost, options.dataBoostEnabled());
     assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
     assertEquals(OrderBy.ORDER_BY_NO_ORDER, options.orderBy());
+    assertEquals(LockHint.LOCK_HINT_SHARED, options.lockHint());
   }
 
   @Test
@@ -373,6 +383,14 @@ public void testReadOptionsOrderBy() {
     assertEquals("orderBy: " + orderBy + " ", options.toString());
   }
 
+  @Test
+  public void testReadOptionsLockHint() {
+    RpcLockHint lockHint = RpcLockHint.SHARED;
+    Options options = Options.fromReadOptions(Options.lockHint(lockHint));
+    assertTrue(options.hasLockHint());
+    assertEquals("lockHint: " + lockHint + " ", options.toString());
+  }
+
   @Test
   public void testReadOptionsWithOrderByEquality() {
     Options optionsWithNoOrderBy1 = Options.fromReadOptions(Options.orderBy(RpcOrderBy.NO_ORDER));
@@ -383,6 +401,19 @@ public void testReadOptionsWithOrderByEquality() {
     assertFalse(optionsWithNoOrderBy1.equals(optionsWithPkOrder));
   }
 
+  @Test
+  public void testReadOptionsWithLockHintEquality() {
+    Options optionsWithSharedLockHint1 =
+        Options.fromReadOptions(Options.lockHint(RpcLockHint.SHARED));
+    Options optionsWithSharedLockHint2 =
+        Options.fromReadOptions(Options.lockHint(RpcLockHint.SHARED));
+    assertEquals(optionsWithSharedLockHint1, optionsWithSharedLockHint2);
+
+    Options optionsWithExclusiveLock =
+        Options.fromReadOptions(Options.lockHint(RpcLockHint.EXCLUSIVE));
+    assertNotEquals(optionsWithSharedLockHint1, optionsWithExclusiveLock);
+  }
+
   @Test
   public void testQueryOptionsPriority() {
     RpcPriority priority = RpcPriority.MEDIUM;