diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java index 36b13f237b47..0ca21debb596 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java +++ b/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 getEnableOpenTelemetryTracing(); + public abstract @Nullable ValueProvider getExcludeTxnFromChangeStreams(); + + public abstract @Nullable ValueProvider getTransactionTag(); + abstract Builder toBuilder(); public static SpannerConfig create() { @@ -210,6 +214,11 @@ abstract Builder setWaitForSessionCreationDuration( abstract Builder setEnableOpenTelemetryTracing( ValueProvider enableOpenTelemetryTracing); + abstract Builder setExcludeTxnFromChangeStreams( + ValueProvider excludeTxnFromChangeStreams); + + abstract Builder setTransactionTag(ValueProvider transactionTag); + abstract Builder setClientCertPath(ValueProvider clientCertPath); abstract Builder setClientCertKeyPath(ValueProvider 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 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 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. * diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java index 4c709963fc68..cc573686357a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java +++ b/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 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 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); } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/com/google/cloud/spanner/OptionsImposter.java b/sdks/java/io/google-cloud-platform/src/test/java/com/google/cloud/spanner/OptionsImposter.java index 72fc3a0d3d20..9d8b87b44453 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/com/google/cloud/spanner/OptionsImposter.java +++ b/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(); + } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java index cf8b876d039b..ab227a9a239d 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java +++ b/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 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 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 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);