[SpannerIO] Support excludeTxnFromChangeStreams and transactionTag in SpannerIO.Write

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java

@@ -117,6 +117,10 @@ public String getHostValue() {
 
   public abstract @Nullable ValueProvider<Boolean> getEnableOpenTelemetryTracing();
 
+  public abstract @Nullable ValueProvider<Boolean> getExcludeTxnFromChangeStreams();
+
+  public abstract @Nullable ValueProvider<String> getTransactionTag();
+
   abstract Builder toBuilder();
 
   public static SpannerConfig create() {
@@ -210,6 +214,11 @@ abstract Builder setWaitForSessionCreationDuration(
     abstract Builder setEnableOpenTelemetryTracing(
         ValueProvider<Boolean> enableOpenTelemetryTracing);
 
+    abstract Builder setExcludeTxnFromChangeStreams(
+        ValueProvider<Boolean> excludeTxnFromChangeStreams);
+
+    abstract Builder setTransactionTag(ValueProvider<String> transactionTag);
+
     abstract Builder setClientCertPath(ValueProvider<String> clientCertPath);
 
     abstract Builder setClientCertKeyPath(ValueProvider<String> clientCertKeyPath);
@@ -479,6 +488,34 @@ public SpannerConfig withEnableOpenTelemetryTracing(boolean enableOpenTelemetryT
         ValueProvider.StaticValueProvider.of(enableOpenTelemetryTracing));
   }
 
+  /**
+   * Specifies whether transactions should be excluded from being recorded in change streams with
+   * the DDL option {@code allow_txn_exclusion=true}.
+   */
+  public SpannerConfig withExcludeTxnFromChangeStreams(
+      ValueProvider<Boolean> excludeTxnFromChangeStreams) {
+    return toBuilder().setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams).build();
+  }
+
+  /**
+   * Specifies whether transactions should be excluded from being recorded in change streams with
+   * the DDL option {@code allow_txn_exclusion=true}.
+   */
+  public SpannerConfig withExcludeTxnFromChangeStreams(boolean excludeTxnFromChangeStreams) {
+    return withExcludeTxnFromChangeStreams(
+        ValueProvider.StaticValueProvider.of(excludeTxnFromChangeStreams));
+  }
+
+  /** Specifies the transaction tag to use for write transactions. */
+  public SpannerConfig withTransactionTag(ValueProvider<String> transactionTag) {
+    return toBuilder().setTransactionTag(transactionTag).build();
+  }
+
+  /** Specifies the transaction tag to use for write transactions. */
+  public SpannerConfig withTransactionTag(String transactionTag) {
+    return withTransactionTag(ValueProvider.StaticValueProvider.of(transactionTag));
+  }
+
   /**
    * Specifies certificate paths to use for mTLS channel.
    *

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java

@@ -1616,6 +1616,41 @@ public Write withMaxCommitDelay(long millis) {
       return withSpannerConfig(config.withMaxCommitDelay(millis));
     }
 
+    /**
+     * Specifies whether transactions should be excluded from being recorded in change streams with
+     * the DDL option {@code allow_txn_exclusion=true}. This does not exclude the transaction from
+     * being recorded in the change streams with the DDL option {@code allow_txn_exclusion} being
+     * false or unset.
+     */
+    public Write withExcludeTxnFromChangeStreams(boolean excludeTxnFromChangeStreams) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams));
+    }
+
+    /**
+     * Specifies whether transactions should be excluded from being recorded in change streams with
+     * the DDL option {@code allow_txn_exclusion=true}. This does not exclude the transaction from
+     * being recorded in the change streams with the DDL option {@code allow_txn_exclusion} being
+     * false or unset.
+     */
+    public Write withExcludeTxnFromChangeStreams(
+        ValueProvider<Boolean> excludeTxnFromChangeStreams) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams));
+    }
+
+    /** Specifies the transaction tag to use for write transactions. */
+    public Write withTransactionTag(String transactionTag) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withTransactionTag(transactionTag));
+    }
+
+    /** Specifies the transaction tag to use for write transactions. */
+    public Write withTransactionTag(ValueProvider<String> transactionTag) {
+      SpannerConfig config = getSpannerConfig();
+      return withSpannerConfig(config.withTransactionTag(transactionTag));
+    }
+
     /**
      * Specifies the maximum cumulative backoff time when retrying after DEADLINE_EXCEEDED errors.
      * Default is 15 mins.
@@ -3063,7 +3098,15 @@ private Options.TransactionOption[] getTransactionOptions() {
                       java.time.Duration.ofMillis(
                           spannerConfig.getMaxCommitDelay().get().getMillis()))
                   : null,
-              commitStats ? Options.commitStats() : null)
+              commitStats ? Options.commitStats() : null,
+              spannerConfig.getExcludeTxnFromChangeStreams() != null
+                      && Boolean.TRUE.equals(spannerConfig.getExcludeTxnFromChangeStreams().get())
+                  ? Options.excludeTxnFromChangeStreams()
+                  : null,
+              spannerConfig.getTransactionTag() != null
+                      && !Strings.isNullOrEmpty(spannerConfig.getTransactionTag().get())
+                  ? Options.tag(spannerConfig.getTransactionTag().get())
+                  : null)
           .filter(Objects::nonNull)
           .toArray(Options.TransactionOption[]::new);
     }

sdks/java/io/google-cloud-platform/src/test/java/com/google/cloud/spanner/OptionsImposter.java

@@ -36,4 +36,12 @@ public static boolean withCommitStats(Options options) {
   public static CommitResponse createCommitResponse(com.google.spanner.v1.CommitResponse proto) {
     return new CommitResponse(proto);
   }
+
+  public static Boolean withExcludeTxnFromChangeStreams(Options options) {
+    return options.withExcludeTxnFromChangeStreams();
+  }
+
+  public static String tag(Options options) {
+    return options.tag();
+  }
 }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java

@@ -24,6 +24,7 @@
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.hasSize;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -325,6 +326,81 @@ public void runBatchQueryTestWithMaxCommitDelay() {
                 }));
   }
 
+  @Test
+  public void runBatchQueryTestWithExcludeTxnFromChangeStreams() {
+    SpannerIO.Write write =
+        SpannerIO.write()
+            .withSpannerConfig(SPANNER_CONFIG)
+            .withServiceFactory(serviceFactory)
+            .withExcludeTxnFromChangeStreams(true);
+    assertTrue(write.getSpannerConfig().getExcludeTxnFromChangeStreams().get());
+
+    Mutation mutation = buildUpsertMutation(2L);
+    PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
+    mutations.apply(write);
+    pipeline.run();
+
+    verify(serviceFactory.mockDatabaseClient(), times(1))
+        .writeAtLeastOnceWithOptions(
+            mutationsInNoOrder(buildMutationBatch(mutation)),
+            any(ReadQueryUpdateTransactionOption.class),
+            argThat(
+                opts -> {
+                  Options options = OptionsImposter.fromTransactionOptions(opts);
+                  return Boolean.TRUE.equals(
+                      OptionsImposter.withExcludeTxnFromChangeStreams(options));
+                }));
+  }
+
+  @Test
+  public void runBatchQueryTestWithExcludeTxnFromChangeStreamsFalse() {
+    SpannerIO.Write write =
+        SpannerIO.write()
+            .withSpannerConfig(SPANNER_CONFIG)
+            .withServiceFactory(serviceFactory)
+            .withExcludeTxnFromChangeStreams(false);
+    assertFalse(write.getSpannerConfig().getExcludeTxnFromChangeStreams().get());
+
+    Mutation mutation = buildUpsertMutation(2L);
+    PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
+    mutations.apply(write);
+    pipeline.run();
+
+    verify(serviceFactory.mockDatabaseClient(), times(1))
+        .writeAtLeastOnceWithOptions(
+            mutationsInNoOrder(buildMutationBatch(mutation)),
+            argThat(
+                (Options.TransactionOption opt) -> {
+                  Options options = OptionsImposter.fromTransactionOptions(opt);
+                  return OptionsImposter.withExcludeTxnFromChangeStreams(options) == null;
+                }));
+  }
+
+  @Test
+  public void runBatchQueryTestWithTransactionTag() {
+    SpannerIO.Write write =
+        SpannerIO.write()
+            .withSpannerConfig(SPANNER_CONFIG)
+            .withServiceFactory(serviceFactory)
+            .withTransactionTag("txBy=testJobId");
+    assertEquals("txBy=testJobId", write.getSpannerConfig().getTransactionTag().get());
+
+    Mutation mutation = buildUpsertMutation(2L);
+    PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
+    mutations.apply(write);
+    pipeline.run();
+
+    verify(serviceFactory.mockDatabaseClient(), times(1))
+        .writeAtLeastOnceWithOptions(
+            mutationsInNoOrder(buildMutationBatch(mutation)),
+            any(ReadQueryUpdateTransactionOption.class),
+            argThat(
+                opts -> {
+                  Options options = OptionsImposter.fromTransactionOptions(opts);
+                  return "txBy=testJobId".equals(OptionsImposter.tag(options));
+                }));
+  }
+
   @Test
   public void singleMutationPipeline() throws Exception {
     Mutation mutation = buildUpsertMutation(2L);