feat: add support for tagging (#576)

Author: mayurkale22

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

@@ -563,6 +563,12 @@ RequestOptions buildRequestOptions(Options options) {
  if (options.hasPriority()) {
  builder.setPriority(options.priority());
  }
+ if (options.hasTag()) {
+ builder.setRequestTag(options.tag());
+ }
+ if (getTransactionTag() != null) {
+ builder.setTransactionTag(getTransactionTag());
+ }
  return builder.build();
  }
 
@@ -707,6 +713,15 @@ public void close() {
  @Nullable
  abstract TransactionSelector getTransactionSelector();
 
+ /**
+ * Returns the transaction tag for this {@link AbstractReadContext} or <code>null</code> if this
+ * {@link AbstractReadContext} does not have a transaction tag.
+ */
+ @Nullable
+ String getTransactionTag() {
+ return null;
+ }
+
  /** This method is called when a statement returned a new transaction as part of its results. */
  @Override
  public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {}

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

@@ -101,6 +101,14 @@ public static ReadQueryUpdateTransactionOption priority(RpcPriority priority) {
  return new PriorityOption(priority);
  }
 
+ /**
+ * Specifying this will cause the reads, queries, updates and writes operations statistics
+ * collection to be grouped by tag.
+ */
+ public static ReadQueryUpdateTransactionOption tag(String name) {
+ return new TagOption(name);
+ }
+
  /**
  * Specifying this will cause the list operations to fetch at most this many records in a page.
  */
@@ -194,6 +202,19 @@ void appendToOptions(Options options) {
  }
  }
 
+ static final class TagOption extends InternalOption implements ReadQueryUpdateTransactionOption {
+ private final String tag;
+
+ TagOption(String tag) {
+ this.tag = tag;
+ }
+
+ @Override
+ void appendToOptions(Options options) {
+ options.tag = tag;
+ }
+ }
+
  private boolean withCommitStats;
  private Long limit;
  private Integer prefetchChunks;
@@ -202,6 +223,7 @@ void appendToOptions(Options options) {
  private String pageToken;
  private String filter;
  private RpcPriority priority;
+ private String tag;
 
  // Construction is via factory methods below.
  private Options() {}
@@ -266,6 +288,14 @@ Priority priority() {
  return priority == null ? null : priority.proto;
  }
 
+ boolean hasTag() {
+ return tag != null;
+ }
+
+ String tag() {
+ return tag;
+ }
+
  @Override
  public String toString() {
  StringBuilder b = new StringBuilder();
@@ -290,6 +320,9 @@ public String toString() {
  if (priority != null) {
  b.append("priority: ").append(priority).append(' ');
  }
+ if (tag != null) {
+ b.append("tag: ").append(tag).append(' ');
+ }
  return b.toString();
  }
 
@@ -320,7 +353,8 @@ public boolean equals(Object o) {
  || hasPageSize() && that.hasPageSize() && Objects.equals(pageSize(), that.pageSize()))
  && Objects.equals(pageToken(), that.pageToken())
  && Objects.equals(filter(), that.filter())
- && Objects.equals(priority(), that.priority());
+ && Objects.equals(priority(), that.priority())
+ && Objects.equals(tag(), that.tag());
  }
 
  @Override
@@ -350,6 +384,9 @@ public int hashCode() {
  if (priority != null) {
  result = 31 * result + priority.hashCode();
  }
+ if (tag != null) {
+ result = 31 * result + tag.hashCode();
+ }
  return result;
  }
 

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

@@ -181,11 +181,16 @@ ExecuteSqlRequest newTransactionRequestFrom(final Statement statement, final Opt
 
  builder.setResumeToken(ByteString.EMPTY);
 
- if (options.hasPriority()) {
- builder.setRequestOptions(
- RequestOptions.newBuilder().setPriority(options.priority()).build());
+ if (options.hasPriority() || options.hasTag()) {
+ RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+ if (options.hasPriority()) {
+ requestOptionsBuilder.setPriority(options.priority());
+ }
+ if (options.hasTag()) {
+ requestOptionsBuilder.setRequestTag(options.tag());
+ }
+ builder.setRequestOptions(requestOptionsBuilder.build());
  }
-
  return builder.build();
  }
 

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

@@ -173,9 +173,15 @@ public CommitResponse writeAtLeastOnceWithOptions(
  .setSingleUseTransaction(
  TransactionOptions.newBuilder()
  .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()));
- if (commitRequestOptions.hasPriority()) {
- requestBuilder.setRequestOptions(
- RequestOptions.newBuilder().setPriority(commitRequestOptions.priority()).build());
+ if (commitRequestOptions.hasPriority() || commitRequestOptions.hasTag()) {
+ RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+ if (commitRequestOptions.hasPriority()) {
+ requestOptionsBuilder.setPriority(commitRequestOptions.priority());
+ }
+ if (commitRequestOptions.hasTag()) {
+ requestOptionsBuilder.setTransactionTag(commitRequestOptions.tag());
+ }
+ requestBuilder.setRequestOptions(requestOptionsBuilder.build());
  }
  Span span = tracer.spanBuilder(SpannerImpl.COMMIT).startSpan();
  try (Scope s = tracer.withSpan(span)) {

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

@@ -299,9 +299,15 @@ ApiFuture<CommitResponse> commitAsync() {
  CommitRequest.newBuilder()
  .setSession(session.getName())
  .setReturnCommitStats(options.withCommitStats());
- if (options.hasPriority()) {
- builder.setRequestOptions(
- RequestOptions.newBuilder().setPriority(options.priority()).build());
+ if (options.hasPriority() || getTransactionTag() != null) {
+ RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+ if (options.hasPriority()) {
+ requestOptionsBuilder.setPriority(options.priority());
+ }
+ if (getTransactionTag() != null) {
+ requestOptionsBuilder.setTransactionTag(getTransactionTag());
+ }
+ builder.setRequestOptions(requestOptionsBuilder.build());
  }
  synchronized (lock) {
  if (transactionIdFuture == null && transactionId == null && runningAsyncOperations == 0) {
@@ -349,9 +355,15 @@ public void run() {
  requestBuilder.setTransactionId(
  transactionId == null ? transactionIdFuture.get() : transactionId);
  }
- if (options.hasPriority()) {
- requestBuilder.setRequestOptions(
- RequestOptions.newBuilder().setPriority(options.priority()).build());
+ if (options.hasPriority() || getTransactionTag() != null) {
+ RequestOptions.Builder requestOptionsBuilder = RequestOptions.newBuilder();
+ if (options.hasPriority()) {
+ requestOptionsBuilder.setPriority(options.priority());
+ }
+ if (getTransactionTag() != null) {
+ requestOptionsBuilder.setTransactionTag(getTransactionTag());
+ }
+ requestBuilder.setRequestOptions(requestOptionsBuilder.build());
  }
  final CommitRequest commitRequest = requestBuilder.build();
  span.addAnnotation("Starting Commit");
@@ -531,6 +543,12 @@ public void onTransactionMetadata(Transaction transaction, boolean shouldInclude
  }
  }
 
+ @Nullable
+ String getTransactionTag() {
+ if (this.options.hasTag()) return this.options.tag();
+ return null;
+ }
+
  @Override
  public SpannerException onError(SpannerException e, boolean withBeginTransaction) {
  // If the statement that caused an error was the statement that included a BeginTransaction

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

@@ -64,6 +64,14 @@ TestReadContext build() {
  }
  }
 
+ class TestReadContextWithTagBuilder
+ extends AbstractReadContext.Builder<TestReadContextWithTagBuilder, TestReadContextWithTag> {
+ @Override
+ TestReadContextWithTag build() {
+ return new TestReadContextWithTag(this);
+ }
+ }
+
  private final class TestReadContext extends AbstractReadContext {
  TestReadContext(TestReadContextBuilder builder) {
  super(builder);
@@ -75,6 +83,21 @@ TransactionSelector getTransactionSelector() {
  }
  }
 
+ private final class TestReadContextWithTag extends AbstractReadContext {
+ TestReadContextWithTag(TestReadContextWithTagBuilder builder) {
+ super(builder);
+ }
+
+ @Override
+ TransactionSelector getTransactionSelector() {
+ return TransactionSelector.getDefaultInstance();
+ }
+
+ String getTransactionTag() {
+ return "app=spanner,env=test";
+ }
+ }
+
  private TestReadContext context;
 
  @Before
@@ -162,4 +185,46 @@ public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
  Options.fromQueryOptions(Options.priority(RpcPriority.LOW)));
  assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority());
  }
+
+ public void executeSqlRequestBuilderWithRequestOptions() {
+ ExecuteSqlRequest request =
+ context
+ .getExecuteSqlRequestBuilder(
+ Statement.newBuilder("SELECT FOO FROM BAR").build(),
+ QueryMode.NORMAL,
+ Options.fromUpdateOptions(Options.tag("app=spanner,env=test,action=query")),
+ false)
+ .build();
+ assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR");
+ assertThat(request.getRequestOptions().getRequestTag())
+ .isEqualTo("app=spanner,env=test,action=query");
+ assertThat(request.getRequestOptions().getTransactionTag()).isEmpty();
+ }
+
+ @Test
+ public void executeSqlRequestBuilderWithRequestOptionsWithTxnTag() {
+ SessionImpl session = mock(SessionImpl.class);
+ when(session.getName()).thenReturn("session-1");
+ TestReadContextWithTagBuilder builder = new TestReadContextWithTagBuilder();
+ TestReadContextWithTag contextWithTag =
+ builder
+ .setSession(session)
+ .setRpc(mock(SpannerRpc.class))
+ .setDefaultQueryOptions(defaultQueryOptions)
+ .setExecutorProvider(mock(ExecutorProvider.class))
+ .build();
+
+ ExecuteSqlRequest request =
+ contextWithTag
+ .getExecuteSqlRequestBuilder(
+ Statement.newBuilder("SELECT FOO FROM BAR").build(),
+ QueryMode.NORMAL,
+ Options.fromUpdateOptions(Options.tag("app=spanner,env=test,action=query")),
+ false)
+ .build();
+ assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR");
+ assertThat(request.getRequestOptions().getRequestTag())
+ .isEqualTo("app=spanner,env=test,action=query");
+ assertThat(request.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test");
+ }
 }