This is an automated email from the ASF dual-hosted git repository.
stankiewicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 112739c3fea [SpannerIO] Add withWriteResults() to SpannerIO.Write to
emit commit metadata (#40480)
112739c3fea is described below
commit 112739c3fea2469377e3a9a34b271824f47b28d1
Author: Radosław Stankiewicz <[email protected]>
AuthorDate: Mon Oct 12 20:09:41 2026 +0200
[SpannerIO] Add withWriteResults() to SpannerIO.Write to emit commit
metadata (#40480)
---
.../apache/beam/sdk/io/gcp/spanner/SpannerIO.java | 156 ++++++++++++++--
.../sdk/io/gcp/spanner/SpannerWriteResult.java | 26 ++-
.../beam/sdk/io/gcp/spanner/WriteResults.java | 122 +++++++++++++
.../com/google/cloud/spanner/OptionsImposter.java | 8 +
.../sdk/io/gcp/spanner/SpannerIOWriteTest.java | 199 +++++++++++++++++++++
.../beam/sdk/io/gcp/spanner/SpannerWriteIT.java | 38 ++++
6 files changed, 532 insertions(+), 17 deletions(-)
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 1d892fb11e4..d2ea020c471 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
@@ -30,6 +30,7 @@ import static
org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsCons
import static
org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.LOW_LATENCY_REAL_TIME_CHECKPOINT_INTERVAL;
import static
org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.MAX_INCLUSIVE_END_AT;
import static
org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamsConstants.THROUGHPUT_WINDOW_SECONDS;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
@@ -41,6 +42,7 @@ import com.google.cloud.NoCredentials;
import com.google.cloud.ServiceFactory;
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.AbortedException;
+import com.google.cloud.spanner.CommitResponse;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Dialect;
@@ -529,6 +531,7 @@ public class SpannerIO {
.setMaxNumMutations(DEFAULT_MAX_NUM_MUTATIONS)
.setMaxNumRows(DEFAULT_MAX_NUM_ROWS)
.setFailureMode(FailureMode.FAIL_FAST)
+ .setCommitStats(false)
.build();
}
@@ -1427,6 +1430,8 @@ public class SpannerIO {
abstract @Nullable PCollectionView<Dialect> getDialectView();
+ abstract boolean getCommitStats();
+
abstract Builder toBuilder();
@AutoValue.Builder
@@ -1448,6 +1453,8 @@ public class SpannerIO {
abstract Builder setDialectView(PCollectionView<Dialect> dialect);
+ abstract Builder setCommitStats(boolean commitStats);
+
abstract Write build();
}
@@ -1693,6 +1700,14 @@ public class SpannerIO {
return withSpannerConfig(config.withRpcPriority(RpcPriority.HIGH));
}
+ /**
+ * Returns a {@link WriteWithResults} transform that enables {@link
Options#commitStats()} and
+ * emits a {@link WriteResults} for each batch of mutations written to
Cloud Spanner.
+ */
+ public WriteWithResults withWriteResults() {
+ return new WriteWithResults(toBuilder().setCommitStats(true).build());
+ }
+
@Override
public SpannerWriteResult expand(PCollection<Mutation> input) {
getSpannerConfig().validate();
@@ -1731,6 +1746,65 @@ public class SpannerIO {
}
}
+ /**
+ * A {@link PTransform} that writes {@link Mutation} objects to Google Cloud
Spanner with {@link
+ * Options#commitStats()} enabled and emits a {@link WriteResults} for each
committed batch.
+ */
+ public static class WriteWithResults
+ extends PTransform<PCollection<Mutation>, PCollection<WriteResults>> {
+
+ private final Write spec;
+
+ WriteWithResults(Write spec) {
+ this.spec = spec;
+ }
+
+ /** Same transform but can be applied to {@link PCollection} of {@link
MutationGroup}. */
+ public WriteGroupedWithResults grouped() {
+ return new WriteGroupedWithResults(spec);
+ }
+
+ @Override
+ public PCollection<WriteResults> expand(PCollection<Mutation> input) {
+ spec.getSpannerConfig().validate();
+
+ return input
+ .apply("To mutation group", ParDo.of(new ToMutationGroupFn()))
+ .apply("Write mutations to Cloud Spanner", new
WriteGroupedWithResults(spec));
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ spec.populateDisplayDataWithParamaters(builder);
+ }
+ }
+
+ /** Same as {@link WriteWithResults} but supports grouped mutations. */
+ public static class WriteGroupedWithResults
+ extends PTransform<PCollection<MutationGroup>,
PCollection<WriteResults>> {
+
+ private final Write spec;
+
+ public WriteGroupedWithResults(Write spec) {
+ this.spec = spec;
+ }
+
+ @Override
+ public PCollection<WriteResults> expand(PCollection<MutationGroup> input) {
+ SpannerWriteResult result = input.apply(new WriteGrouped(spec));
+ return checkStateNotNull(
+ result.getWriteResults(),
+ "WriteResults PCollection must be present when commitStats is
enabled.");
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ spec.populateDisplayDataWithParamaters(builder);
+ }
+ }
+
static class WriteRows extends PTransform<PCollection<Row>, PDone> {
private final Write write;
@@ -1769,6 +1843,8 @@ public class SpannerIO {
new TupleTag<Iterable<MutationGroup>>("unbatchableMutations") {};
private static final TupleTag<Void> MAIN_OUT_TAG = new
TupleTag<Void>("mainOut") {};
+ private static final TupleTag<WriteResults> WRITE_RESULTS_TAG =
+ new TupleTag<WriteResults>("writeResults") {};
private static final TupleTag<MutationGroup> FAILED_MUTATIONS_TAG =
new TupleTag<MutationGroup>("failedMutations") {};
private static final SerializableCoder<MutationGroup> CODER =
@@ -1778,6 +1854,15 @@ public class SpannerIO {
this.spec = spec;
}
+ /**
+ * Returns a {@link WriteGroupedWithResults} transform that enables {@link
+ * Options#commitStats()} and emits a {@link WriteResults} for each batch
of mutations written
+ * to Cloud Spanner.
+ */
+ public WriteGroupedWithResults withWriteResults() {
+ return new
WriteGroupedWithResults(spec.toBuilder().setCommitStats(true).build());
+ }
+
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
@@ -1876,12 +1961,19 @@ public class SpannerIO {
"Write batches to Spanner",
ParDo.of(
new WriteToSpannerFn(
- spec.getSpannerConfig(), spec.getFailureMode(),
FAILED_MUTATIONS_TAG))
- .withOutputTags(MAIN_OUT_TAG,
TupleTagList.of(FAILED_MUTATIONS_TAG)));
+ spec.getSpannerConfig(),
+ spec.getFailureMode(),
+ FAILED_MUTATIONS_TAG,
+ spec.getCommitStats(),
+ WRITE_RESULTS_TAG))
+ .withOutputTags(
+ MAIN_OUT_TAG,
TupleTagList.of(FAILED_MUTATIONS_TAG).and(WRITE_RESULTS_TAG)));
return new SpannerWriteResult(
input.getPipeline(),
result.get(MAIN_OUT_TAG),
+ spec.getCommitStats() ? result.get(WRITE_RESULTS_TAG) : null,
+ spec.getCommitStats() ? WRITE_RESULTS_TAG : null,
result.get(FAILED_MUTATIONS_TAG),
FAILED_MUTATIONS_TAG);
}
@@ -2501,6 +2593,8 @@ public class SpannerIO {
private long sortableNumCells = 0;
// total number of rows mutated in mutationsToSort
private long sortableNumRows = 0;
+ // minimum timestamp of elements currently in mutationsToSort
+ private @Nullable Instant minBundleTimestamp = null;
GatherSortCreateBatchesFn(
long maxBatchSizeBytes,
@@ -2529,11 +2623,13 @@ public class SpannerIO {
sortableSizeBytes = 0;
sortableNumCells = 0;
sortableNumRows = 0;
+ minBundleTimestamp = null;
}
@FinishBundle
public synchronized void finishBundle(FinishBundleContext c) throws
Exception {
- sortAndOutputBatches(new OutputReceiverForFinishBundle(c));
+ Instant outputTimestamp = minBundleTimestamp != null ?
minBundleTimestamp : Instant.now();
+ sortAndOutputBatches(new OutputReceiverForFinishBundle(c,
outputTimestamp));
}
private synchronized void
sortAndOutputBatches(OutputReceiver<Iterable<MutationGroup>> out)
@@ -2617,6 +2713,12 @@ public class SpannerIO {
sortAndOutputBatches(out);
}
+ Instant elementTimestamp = c.timestamp();
+ if (elementTimestamp != null
+ && (minBundleTimestamp == null ||
elementTimestamp.isBefore(minBundleTimestamp))) {
+ minBundleTimestamp = elementTimestamp;
+ }
+
mutationsToSort.add(
new MutationGroupContainer(
mg, groupSize, groupCells, groupRows,
encoder.encodeTableNameAndKey(mg.primary())));
@@ -2663,7 +2765,7 @@ public class SpannerIO {
private final OutputBuilderSupplier outputBuilderSupplier;
private final DoFn<MutationGroup,
Iterable<MutationGroup>>.FinishBundleContext context;
- OutputReceiverForFinishBundle(FinishBundleContext context) {
+ OutputReceiverForFinishBundle(FinishBundleContext context, Instant
timestamp) {
this.context = context;
this.outputBuilderSupplier =
new OutputBuilderSupplier() {
@@ -2671,7 +2773,7 @@ public class SpannerIO {
public <OutputT> WindowedValues.Builder<OutputT> builder(OutputT
value) {
return WindowedValues.<OutputT>builder()
.setValue(value)
- .setTimestamp(Instant.now())
+ .setTimestamp(timestamp)
.setPaneInfo(PaneInfo.NO_FIRING)
.setWindow(GlobalWindow.INSTANCE);
}
@@ -2799,6 +2901,8 @@ public class SpannerIO {
Metrics.counter(WriteGrouped.class, "spanner_write_retries");
private final TupleTag<MutationGroup> failedTag;
+ private final boolean commitStats;
+ private final @Nullable TupleTag<WriteResults> writeResultsTag;
// Fluent Backoff is not serializable so create at runtime in setup().
private transient FluentBackoff bundleWriteBackoff;
@@ -2806,9 +2910,20 @@ public class SpannerIO {
WriteToSpannerFn(
SpannerConfig spannerConfig, FailureMode failureMode,
TupleTag<MutationGroup> failedTag) {
+ this(spannerConfig, failureMode, failedTag, false, null);
+ }
+
+ WriteToSpannerFn(
+ SpannerConfig spannerConfig,
+ FailureMode failureMode,
+ TupleTag<MutationGroup> failedTag,
+ boolean commitStats,
+ @Nullable TupleTag<WriteResults> writeResultsTag) {
this.spannerConfig = spannerConfig;
this.failureMode = failureMode;
this.failedTag = failedTag;
+ this.commitStats = commitStats;
+ this.writeResultsTag = writeResultsTag;
}
@Setup
@@ -2851,9 +2966,12 @@ public class SpannerIO {
mutationGroupBatchesReceived.inc();
mutationGroupsReceived.inc(mutations.size());
Iterable<Mutation> batch = Iterables.concat(mutations);
- writeMutations(batch);
+ CommitResponse response = writeMutations(batch);
mutationGroupBatchesWriteSuccess.inc();
mutationGroupsWriteSuccess.inc(mutations.size());
+ if (commitStats && writeResultsTag != null) {
+ c.output(writeResultsTag, WriteResults.fromCommitResponse(response));
+ }
return;
} catch (SpannerException e) {
mutationGroupBatchesWriteFail.inc();
@@ -2872,8 +2990,11 @@ public class SpannerIO {
for (MutationGroup mg : mutations) {
try {
spannerWriteRetries.inc();
- writeMutations(mg);
+ CommitResponse response = writeMutations(mg);
mutationGroupsWriteSuccess.inc();
+ if (commitStats && writeResultsTag != null) {
+ c.output(writeResultsTag,
WriteResults.fromCommitResponse(response));
+ }
} catch (SpannerException e) {
mutationGroupsWriteFail.inc();
LOG.warn("Failed to write the mutation group: {}", mg, e);
@@ -2886,13 +3007,15 @@ public class SpannerIO {
Spanner aborts all inflight transactions during a schema change. Client
is expected
to retry silently. These must not be counted against retry backoff.
*/
- private void spannerWriteWithRetryIfSchemaChange(List<Mutation> batch)
throws SpannerException {
+ private CommitResponse spannerWriteWithRetryIfSchemaChange(List<Mutation>
batch)
+ throws SpannerException {
Set<String> tableNames =
batch.stream().map(Mutation::getTable).collect(Collectors.toSet());
for (int retry = 1; ; retry++) {
try {
- spannerAccessor
- .getDatabaseClient()
- .writeAtLeastOnceWithOptions(batch, getTransactionOptions());
+ CommitResponse response =
+ spannerAccessor
+ .getDatabaseClient()
+ .writeAtLeastOnceWithOptions(batch, getTransactionOptions());
// Get names of all tables in batch of mutations.
reportServiceCallMetricsForBatch(tableNames, "ok");
for (String tableName : tableNames) {
@@ -2906,7 +3029,7 @@ public class SpannerIO {
spannerConfig.getDatabaseId().get(),
tableName));
}
- return;
+ return response;
} catch (AbortedException e) {
reportServiceCallMetricsForBatch(
tableNames, e.getErrorCode().getGrpcStatusCode().toString());
@@ -2937,7 +3060,8 @@ public class SpannerIO {
? Options.maxCommitDelay(
java.time.Duration.ofMillis(
spannerConfig.getMaxCommitDelay().get().getMillis()))
- : null)
+ : null,
+ commitStats ? Options.commitStats() : null)
.filter(Objects::nonNull)
.toArray(Options.TransactionOption[]::new);
}
@@ -2964,7 +3088,7 @@ public class SpannerIO {
}
/** Write the Mutations to Spanner, handling DEADLINE_EXCEEDED with
backoff/retries. */
- private void writeMutations(Iterable<Mutation> mutationIterable)
+ private CommitResponse writeMutations(Iterable<Mutation> mutationIterable)
throws SpannerException, IOException {
BackOff backoff = bundleWriteBackoff.backoff();
List<Mutation> mutations = ImmutableList.copyOf(mutationIterable);
@@ -2973,9 +3097,9 @@ public class SpannerIO {
Stopwatch timer = Stopwatch.createStarted();
// loop is broken on success, timeout backoff/retry attempts exceeded,
or other failure.
try {
- spannerWriteWithRetryIfSchemaChange(mutations);
+ CommitResponse response =
spannerWriteWithRetryIfSchemaChange(mutations);
spannerWriteSuccess.inc();
- return;
+ return response;
} catch (SpannerException exception) {
if (exception.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED) {
spannerWriteTimeouts.inc();
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteResult.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteResult.java
index 87f67345036..31f612b00e8 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteResult.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteResult.java
@@ -26,6 +26,7 @@ import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.checkerframework.checker.nullness.qual.Nullable;
/**
* The results of a {@link SpannerIO#write()} transform.
@@ -42,6 +43,8 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Immuta
public class SpannerWriteResult implements POutput {
private final Pipeline pipeline;
private final PCollection<Void> output;
+ private final @Nullable PCollection<WriteResults> writeResults;
+ private final @Nullable TupleTag<WriteResults> writeResultsTag;
private final PCollection<MutationGroup> failedMutations;
private final TupleTag<MutationGroup> failedMutationsTag;
@@ -50,8 +53,20 @@ public class SpannerWriteResult implements POutput {
PCollection<Void> output,
PCollection<MutationGroup> failedMutations,
TupleTag<MutationGroup> failedMutationsTag) {
+ this(pipeline, output, null, null, failedMutations, failedMutationsTag);
+ }
+
+ public SpannerWriteResult(
+ Pipeline pipeline,
+ PCollection<Void> output,
+ @Nullable PCollection<WriteResults> writeResults,
+ @Nullable TupleTag<WriteResults> writeResultsTag,
+ PCollection<MutationGroup> failedMutations,
+ TupleTag<MutationGroup> failedMutationsTag) {
this.pipeline = pipeline;
this.output = output;
+ this.writeResults = writeResults;
+ this.writeResultsTag = writeResultsTag;
this.failedMutations = failedMutations;
this.failedMutationsTag = failedMutationsTag;
}
@@ -63,7 +78,12 @@ public class SpannerWriteResult implements POutput {
@Override
public Map<TupleTag<?>, PValue> expand() {
- return ImmutableMap.of(failedMutationsTag, failedMutations);
+ ImmutableMap.Builder<TupleTag<?>, PValue> builder = ImmutableMap.builder();
+ builder.put(failedMutationsTag, failedMutations);
+ if (writeResults != null && writeResultsTag != null) {
+ builder.put(writeResultsTag, writeResults);
+ }
+ return builder.build();
}
public PCollection<MutationGroup> getFailedMutations() {
@@ -74,6 +94,10 @@ public class SpannerWriteResult implements POutput {
return output;
}
+ public @Nullable PCollection<WriteResults> getWriteResults() {
+ return writeResults;
+ }
+
@Override
public void finishSpecifyingOutput(
String transformName, PInput input, PTransform<?, ?> transform) {}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/WriteResults.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/WriteResults.java
new file mode 100644
index 00000000000..a85cd3e7ffa
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/WriteResults.java
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.spanner;
+
+import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.CommitResponse;
+import java.io.Serializable;
+import java.util.Objects;
+import org.apache.avro.reflect.AvroEncode;
+import org.apache.beam.sdk.coders.DefaultCoder;
+import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
+import
org.apache.beam.sdk.io.gcp.spanner.changestreams.encoder.TimestampEncoding;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * The results of committing a batch of mutations to Cloud Spanner via {@link
+ * SpannerIO.Write#withWriteResults()}.
+ */
+@SuppressWarnings("initialization.fields.uninitialized") // Avro requires the
default constructor
+@DefaultCoder(AvroCoder.class)
+public class WriteResults implements Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ private long mutationCount;
+
+ @AvroEncode(using = TimestampEncoding.class)
+ private @Nullable Timestamp commitTimestamp;
+
+ @AvroEncode(using = TimestampEncoding.class)
+ private @Nullable Timestamp snapshotTimestamp;
+
+ /** Default constructor for Avro serialization only. */
+ private WriteResults() {}
+
+ private WriteResults(
+ long mutationCount,
+ @Nullable Timestamp commitTimestamp,
+ @Nullable Timestamp snapshotTimestamp) {
+ this.mutationCount = mutationCount;
+ this.commitTimestamp = commitTimestamp;
+ this.snapshotTimestamp = snapshotTimestamp;
+ }
+
+ /** The number of mutations applied by Cloud Spanner in this commit, or 0 if
unavailable. */
+ public long getMutationCount() {
+ return mutationCount;
+ }
+
+ /** The timestamp at which the write was committed in Cloud Spanner. */
+ public @Nullable Timestamp getCommitTimestamp() {
+ return commitTimestamp;
+ }
+
+ /** The snapshot timestamp of the transaction, if returned by Cloud Spanner.
*/
+ public @Nullable Timestamp getSnapshotTimestamp() {
+ return snapshotTimestamp;
+ }
+
+ public static WriteResults create(
+ long mutationCount,
+ @Nullable Timestamp commitTimestamp,
+ @Nullable Timestamp snapshotTimestamp) {
+ return new WriteResults(mutationCount, commitTimestamp, snapshotTimestamp);
+ }
+
+ public static WriteResults fromCommitResponse(@Nullable CommitResponse
commitResponse) {
+ if (commitResponse == null) {
+ return create(0L, null, null);
+ }
+ long mutationCount =
+ commitResponse.hasCommitStats() ?
commitResponse.getCommitStats().getMutationCount() : 0L;
+ return create(
+ mutationCount, commitResponse.getCommitTimestamp(),
commitResponse.getSnapshotTimestamp());
+ }
+
+ @Override
+ public boolean equals(@Nullable Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof WriteResults)) {
+ return false;
+ }
+ WriteResults that = (WriteResults) o;
+ return mutationCount == that.mutationCount
+ && Objects.equals(commitTimestamp, that.commitTimestamp)
+ && Objects.equals(snapshotTimestamp, that.snapshotTimestamp);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(mutationCount, commitTimestamp, snapshotTimestamp);
+ }
+
+ @Override
+ public String toString() {
+ return "WriteResults{"
+ + "mutationCount="
+ + mutationCount
+ + ", commitTimestamp="
+ + commitTimestamp
+ + ", snapshotTimestamp="
+ + snapshotTimestamp
+ + '}';
+ }
+}
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 9d2954950e9..72fc3a0d3d2 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
@@ -28,4 +28,12 @@ public class OptionsImposter {
public static Duration maxCommitDelay(Options options) {
return options.maxCommitDelay();
}
+
+ public static boolean withCommitStats(Options options) {
+ return options.withCommitStats();
+ }
+
+ public static CommitResponse
createCommitResponse(com.google.spanner.v1.CommitResponse proto) {
+ return new CommitResponse(proto);
+ }
}
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 4adfe9fb0ea..cf8b876d039 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
@@ -67,6 +67,7 @@ import
org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.FailureMode;
@@ -76,6 +77,7 @@ import
org.apache.beam.sdk.io.gcp.spanner.SpannerIO.WriteToSpannerFn;
import org.apache.beam.sdk.metrics.Lineage;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.testing.CoderProperties;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestStream;
@@ -84,15 +86,19 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.FinishBundleContext;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.ProcessContext;
+import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.util.Sleeper;
+import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.TimestampedValue;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.joda.time.Duration;
+import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -1626,6 +1632,199 @@ public class SpannerIOWriteTest implements Serializable
{
verify(serviceFactory.mockSpanner(), times(2)).close();
}
+ @Test
+ public void testWriteResultsCoder() throws Exception {
+ Coder<WriteResults> coder =
pipeline.getCoderRegistry().getCoder(WriteResults.class);
+ Timestamp commitTs = Timestamp.ofTimeSecondsAndNanos(1700000000L,
123456789);
+ Timestamp snapshotTs = Timestamp.ofTimeSecondsAndNanos(1699999999L,
987654321);
+
+ WriteResults full = WriteResults.create(5L, commitTs, snapshotTs);
+ WriteResults withoutSnapshot = WriteResults.create(3L, commitTs, null);
+ WriteResults allNullTimestamps = WriteResults.create(0L, null, null);
+
+ CoderProperties.coderDecodeEncodeEqual(coder, full);
+ CoderProperties.coderDecodeEncodeEqual(coder, withoutSnapshot);
+ CoderProperties.coderDecodeEncodeEqual(coder, allNullTimestamps);
+ }
+
+ @Test
+ public void testGatherSortAndBatchFn_usesMinBundleTimestamp() throws
Exception {
+ GatherSortCreateBatchesFn testFn =
+ new GatherSortCreateBatchesFn(
+ 10000000, // batch bytes
+ 100, // batch cells
+ 5, // batch rows
+ 100, // groupingFactor
+ null);
+
+ DoFn<MutationGroup, Iterable<MutationGroup>>.ProcessContext
mockProcessContext =
+ Mockito.mock(ProcessContext.class);
+ DoFn<MutationGroup, Iterable<MutationGroup>>.FinishBundleContext
mockFinishBundleContext =
+ Mockito.mock(FinishBundleContext.class);
+ when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
+
+ Instant t1 = Instant.ofEpochMilli(3000L);
+ Instant t2 = Instant.ofEpochMilli(1000L);
+ Instant t3 = Instant.ofEpochMilli(2000L);
+
+
when(mockProcessContext.element()).thenReturn(buildMutationGroup(buildUpsertMutation(1L)));
+ when(mockProcessContext.timestamp()).thenReturn(t1);
+ testFn.processElement(mockProcessContext, null);
+
+
when(mockProcessContext.element()).thenReturn(buildMutationGroup(buildUpsertMutation(2L)));
+ when(mockProcessContext.timestamp()).thenReturn(t2);
+ testFn.processElement(mockProcessContext, null);
+
+
when(mockProcessContext.element()).thenReturn(buildMutationGroup(buildUpsertMutation(3L)));
+ when(mockProcessContext.timestamp()).thenReturn(t3);
+ testFn.processElement(mockProcessContext, null);
+
+ testFn.finishBundle(mockFinishBundleContext);
+
+ verify(mockFinishBundleContext, times(1)).output(any(), eq(t2), any());
+ }
+
+ @Test
+ public void testWriteWithResultsEmitsCommitMetadataAndEnablesCommitStats()
throws Exception {
+ Timestamp expectedCommitTs = Timestamp.ofTimeSecondsAndNanos(1700000000L,
500);
+ Timestamp expectedSnapshotTs =
Timestamp.ofTimeSecondsAndNanos(1699999995L, 250);
+ CommitResponse commitResponse =
+ OptionsImposter.createCommitResponse(
+ com.google.spanner.v1.CommitResponse.newBuilder()
+ .setCommitTimestamp(expectedCommitTs.toProto())
+ .setSnapshotTimestamp(expectedSnapshotTs.toProto())
+ .setCommitStats(
+
com.google.spanner.v1.CommitResponse.CommitStats.newBuilder()
+ .setMutationCount(2L)
+ .build())
+ .build());
+
+ when(serviceFactory
+ .mockDatabaseClient()
+ .writeAtLeastOnceWithOptions(
+ mutationBatchesCaptor.capture(),
+ any(ReadQueryUpdateTransactionOption.class),
+ eq(Options.commitStats())))
+ .thenReturn(commitResponse);
+
+ Mutation m1 = buildUpsertMutation(1L);
+
+ PCollection<WriteResults> results =
+ pipeline
+ .apply(Create.of(m1))
+ .apply(
+ SpannerIO.write()
+ .withProjectId(PROJECT_NAME)
+ .withInstanceId(INSTANCE_NAME)
+ .withDatabaseId(DATABASE_NAME)
+ .withServiceFactory(serviceFactory)
+ .withWriteResults());
+
+ PAssert.that(results)
+ .containsInAnyOrder(WriteResults.create(2L, expectedCommitTs,
expectedSnapshotTs));
+
+ pipeline.run();
+
+ verify(serviceFactory.mockDatabaseClient(), times(1))
+ .writeAtLeastOnceWithOptions(
+ mutationsInNoOrder(buildMutationBatch(m1)),
+ any(ReadQueryUpdateTransactionOption.class),
+ eq(Options.commitStats()));
+ }
+
+ @Test
+ public void testWriteGroupedWithResultsEmitsCommitMetadata() throws
Exception {
+ Timestamp expectedCommitTs = Timestamp.ofTimeSecondsAndNanos(1700000100L,
0);
+ CommitResponse commitResponse =
+ OptionsImposter.createCommitResponse(
+ com.google.spanner.v1.CommitResponse.newBuilder()
+ .setCommitTimestamp(expectedCommitTs.toProto())
+ .setCommitStats(
+
com.google.spanner.v1.CommitResponse.CommitStats.newBuilder()
+ .setMutationCount(2L)
+ .build())
+ .build());
+
+ when(serviceFactory
+ .mockDatabaseClient()
+ .writeAtLeastOnceWithOptions(
+ mutationBatchesCaptor.capture(),
+ any(ReadQueryUpdateTransactionOption.class),
+ eq(Options.commitStats())))
+ .thenReturn(commitResponse);
+
+ Mutation m1 = buildUpsertMutation(1L);
+ Mutation m2 = buildUpsertMutation(2L);
+ MutationGroup g = buildMutationGroup(m1, m2);
+
+ PCollection<WriteResults> results =
+ pipeline
+ .apply(Create.of(Arrays.asList(g)))
+ .apply(
+ SpannerIO.write()
+ .withProjectId(PROJECT_NAME)
+ .withInstanceId(INSTANCE_NAME)
+ .withDatabaseId(DATABASE_NAME)
+ .withServiceFactory(serviceFactory)
+ .grouped()
+ .withWriteResults());
+
+ PAssert.that(results).containsInAnyOrder(WriteResults.create(2L,
expectedCommitTs, null));
+
+ pipeline.run();
+ }
+
+ @Test
+ public void
testWriteWithResultsPreservesElementTimestampWhenBatchingDisabled() throws
Exception {
+ Timestamp expectedCommitTs = Timestamp.ofTimeSecondsAndNanos(1700000200L,
0);
+ CommitResponse commitResponse =
+ OptionsImposter.createCommitResponse(
+ com.google.spanner.v1.CommitResponse.newBuilder()
+ .setCommitTimestamp(expectedCommitTs.toProto())
+ .setCommitStats(
+
com.google.spanner.v1.CommitResponse.CommitStats.newBuilder()
+ .setMutationCount(1L)
+ .build())
+ .build());
+
+ when(serviceFactory
+ .mockDatabaseClient()
+ .writeAtLeastOnceWithOptions(
+ mutationBatchesCaptor.capture(),
+ any(ReadQueryUpdateTransactionOption.class),
+ eq(Options.commitStats())))
+ .thenReturn(commitResponse);
+
+ Mutation m1 = buildUpsertMutation(1L);
+ Instant recordTimestamp = Instant.ofEpochMilli(123456789L);
+
+ PCollection<KV<Instant, WriteResults>> timestampedResults =
+ pipeline
+ .apply(Create.timestamped(TimestampedValue.of(m1,
recordTimestamp)))
+ .apply(
+ SpannerIO.write()
+ .withProjectId(PROJECT_NAME)
+ .withInstanceId(INSTANCE_NAME)
+ .withDatabaseId(DATABASE_NAME)
+ .withBatchSizeBytes(0)
+ .withServiceFactory(serviceFactory)
+ .withWriteResults())
+ .apply(
+ ParDo.of(
+ new DoFn<WriteResults, KV<Instant, WriteResults>>() {
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ c.output(KV.of(c.timestamp(), c.element()));
+ }
+ }));
+
+ PAssert.that(timestampedResults)
+ .containsInAnyOrder(
+ KV.of(recordTimestamp, WriteResults.create(1L, expectedCommitTs,
null)));
+
+ pipeline.run();
+ }
+
static MutationGroup buildMutationGroup(Mutation m, Mutation... other) {
return MutationGroup.create(m, other);
}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
index 318237e1118..cd4a01e7321 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
@@ -21,6 +21,8 @@ import static
org.apache.beam.sdk.io.gcp.spanner.SpannerTestHelper.isOmni;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
import static org.junit.Assume.assumeFalse;
import com.google.api.gax.longrunning.OperationFuture;
@@ -45,6 +47,7 @@ import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.transforms.Create;
@@ -53,6 +56,7 @@ import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.Wait;
+import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
@@ -230,6 +234,40 @@ public class SpannerWriteIT {
assertThat(countNumberOfRecords(pgDatabaseName), equalTo((long)
numRecords));
}
+ @Test
+ public void testWriteWithResults() throws Exception {
+ int numRecords = 100;
+ PCollection<WriteResults> writeResults =
+ p.apply("Init", GenerateSequence.from(0).to(numRecords))
+ .apply("Generate mu", ParDo.of(new
GenerateMutations(options.getTable())))
+ .apply(
+ "Write db with results",
+ SpannerTestHelper.setUpSpannerIO(
+ SpannerIO.write()
+ .withProjectId(project)
+ .withInstanceId(options.getInstanceId())
+ .withDatabaseId(databaseName))
+ .withWriteResults());
+
+ PAssert.that(writeResults)
+ .satisfies(
+ results -> {
+ long totalMutations = 0;
+ for (WriteResults r : results) {
+ assertNotNull(r.getCommitTimestamp());
+ assertTrue(r.getMutationCount() > 0);
+ totalMutations += r.getMutationCount();
+ }
+ assertThat(totalMutations, equalTo(2L * numRecords));
+ return null;
+ });
+
+ PipelineResult result = p.run();
+ result.waitUntilFinish();
+ assertThat(result.getState(), is(PipelineResult.State.DONE));
+ assertThat(countNumberOfRecords(databaseName), equalTo((long) numRecords));
+ }
+
@Test
public void testWriteViaSchemaTransform() throws Exception {
assumeFalse(
