From 04837bd36a8ada175f7c745dd5cca08530777919 Mon Sep 17 00:00:00 2001 From: Weihan Jiang Date: Wed, 12 Aug 2026 21:57:40 -0700 Subject: [PATCH] Add per-destination GCS CMEK support to FileIO --- .../org/apache/beam/sdk/io/FileBasedSink.java | 48 ++++++++++++---- .../java/org/apache/beam/sdk/io/FileIO.java | 23 ++++++++ .../org/apache/beam/sdk/io/FileIOTest.java | 25 +++++++++ .../extensions/gcp/storage/GcsFileSystem.java | 15 ++++- .../gcp/storage/GcsMoveOptions.java | 42 ++++++++++++++ .../beam/sdk/extensions/gcp/util/GcsUtil.java | 9 +++ .../sdk/extensions/gcp/util/GcsUtilV1.java | 56 ++++++++++++++++--- .../gcp/storage/GcsFileSystemTest.java | 22 ++++++++ .../sdk/extensions/gcp/util/GcsUtilTest.java | 21 +++++++ 9 files changed, 242 insertions(+), 19 deletions(-) create mode 100644 sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsMoveOptions.java diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java index bba9b1f82f5b..cdf5a89ac30e 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java @@ -52,6 +52,7 @@ import org.apache.beam.sdk.io.fs.CreateOptions.StandardCreateOptions; import org.apache.beam.sdk.io.fs.MatchResult; import org.apache.beam.sdk.io.fs.MatchResult.Metadata; +import org.apache.beam.sdk.io.fs.MoveOptions; import org.apache.beam.sdk.io.fs.MoveOptions.StandardMoveOptions; import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions; import org.apache.beam.sdk.io.fs.ResourceId; @@ -326,6 +327,11 @@ final void setSideInputAccessorFromProcessContext(DoFn.ProcessContext cont /** Converts a destination into a {@link FilenamePolicy}. May not return null. */ public abstract FilenamePolicy getFilenamePolicy(DestinationT destination); + /** Returns filesystem-specific options to use when moving files to {@code destination}. */ + public List getMoveOptions(DestinationT destination) { + return ImmutableList.of(); + } + /** Populates the display data. */ @Override public void populateDisplayData(DisplayData.Builder builder) {} @@ -787,21 +793,41 @@ final void moveToOutputFiles( int numFiles = resultsToFinalFilenames.size(); LOG.debug("Copying {} files.", numFiles); - List srcFiles = new ArrayList<>(); - List dstFiles = new ArrayList<>(); + Map> moveOptionsByDestination = Maps.newLinkedHashMap(); + Map, List, ResourceId>>> resultsByMoveOptions = + Maps.newLinkedHashMap(); for (KV, ResourceId> entry : resultsToFinalFilenames) { - srcFiles.add(entry.getKey().getTempFilename()); - dstFiles.add(entry.getValue()); + DestinationT destination = entry.getKey().getDestination(); + List destinationMoveOptions = + moveOptionsByDestination.computeIfAbsent( + destination, + unused -> + ImmutableList.copyOf( + checkNotNull( + getSink().getDynamicDestinations().getMoveOptions(destination), + "DynamicDestinations.getMoveOptions() must not return null"))); + resultsByMoveOptions + .computeIfAbsent(destinationMoveOptions, unused -> new ArrayList<>()) + .add(entry); LOG.info( "Will copy temporary file {} to final location {}", entry.getKey(), entry.getValue()); } - // During a failure case, files may have been deleted in an earlier step. Thus - // we ignore missing files here. - FileSystems.rename( - srcFiles, - dstFiles, - StandardMoveOptions.IGNORE_MISSING_FILES, - StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); + for (Map.Entry, List, ResourceId>>> + moveOptionsEntry : resultsByMoveOptions.entrySet()) { + List srcFiles = new ArrayList<>(); + List dstFiles = new ArrayList<>(); + for (KV, ResourceId> entry : moveOptionsEntry.getValue()) { + srcFiles.add(entry.getKey().getTempFilename()); + dstFiles.add(entry.getValue()); + } + List moveOptions = new ArrayList<>(); + // During a failure case, files may have been deleted in an earlier step. Thus + // we ignore missing files here. + moveOptions.add(StandardMoveOptions.IGNORE_MISSING_FILES); + moveOptions.add(StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); + moveOptions.addAll(moveOptionsEntry.getKey()); + FileSystems.rename(srcFiles, dstFiles, moveOptions.toArray(new MoveOptions[0])); + } // The rename ensures that the source files are deleted. However we may still need to clean // up the directory or orphaned files. diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java index b7590a4c2d1a..a2075ec21232 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java @@ -41,6 +41,7 @@ import org.apache.beam.sdk.io.fs.EmptyMatchTreatment; import org.apache.beam.sdk.io.fs.MatchResult; import org.apache.beam.sdk.io.fs.MetadataCoderV2; +import org.apache.beam.sdk.io.fs.MoveOptions; import org.apache.beam.sdk.io.fs.ResourceId; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -1039,6 +1040,8 @@ public static FileNaming relativeFileNaming( abstract @Nullable Contextful> getFileNamingFn(); + abstract @Nullable SerializableFunction> getMoveOptionsFn(); + abstract @Nullable DestinationT getEmptyWindowDestination(); abstract @Nullable Coder getDestinationCoder(); @@ -1092,6 +1095,9 @@ abstract Builder setOutputDirectory( abstract Builder setFileNamingFn( Contextful> namingFn); + abstract Builder setMoveOptionsFn( + SerializableFunction> moveOptionsFn); + abstract Builder setEmptyWindowDestination( DestinationT emptyWindowDestination); @@ -1266,6 +1272,13 @@ public Write withNaming( return toBuilder().setFileNamingFn(namingFn).build(); } + /** Specifies filesystem-specific move options for final files in each destination. */ + public Write withMoveOptions( + SerializableFunction> moveOptionsFn) { + checkArgument(moveOptionsFn != null, "moveOptionsFn can not be null"); + return toBuilder().setMoveOptionsFn(moveOptionsFn).build(); + } + /** Specifies a directory into which all temporary files will be placed. */ public Write withTempDirectory(String tempDirectory) { checkArgument(tempDirectory != null, "tempDirectory can not be null"); @@ -1494,6 +1507,9 @@ public WriteFilesResult expand(PCollection input) { } resolvedSpec.setFileNamingFn(resolveFileNamingFn()); + if (getMoveOptionsFn() != null) { + resolvedSpec.setMoveOptionsFn(getMoveOptionsFn()); + } resolvedSpec.setEmptyWindowDestination(getEmptyWindowDestination()); if (getTempDirectory() == null) { checkArgument( @@ -1721,6 +1737,13 @@ public List> getSideInputs() { public @Nullable Coder getDestinationCoder() { return spec.getDestinationCoder(); } + + @Override + public List getMoveOptions(DestinationT destination) { + return spec.getMoveOptionsFn() == null + ? super.getMoveOptions(destination) + : spec.getMoveOptionsFn().apply(destination); + } } } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java index c5a227d46b82..fa1a5fc66e34 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java @@ -52,6 +52,7 @@ import org.apache.beam.sdk.io.fs.EmptyMatchTreatment; import org.apache.beam.sdk.io.fs.MatchResult; import org.apache.beam.sdk.io.fs.MatchResult.Metadata; +import org.apache.beam.sdk.io.fs.MoveOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; @@ -96,6 +97,11 @@ /** Tests for {@link FileIO}. */ @RunWith(JUnit4.class) public class FileIOTest implements Serializable { + private enum TestMoveOptions implements MoveOptions { + DESTINATION_A, + DESTINATION_B + } + @Rule public transient TestPipeline p = TestPipeline.create(); @Rule public transient TemporaryFolder tmpFolder = new TemporaryFolder(); @@ -104,6 +110,25 @@ public class FileIOTest implements Serializable { @Rule public transient Timeout globalTimeout = Timeout.seconds(1200); + @Test + public void testMoveOptionsCanDependOnDestination() { + FileIO.Write write = + FileIO.writeDynamic() + .withMoveOptions( + destination -> + Collections.singletonList( + destination.equals("a") + ? TestMoveOptions.DESTINATION_A + : TestMoveOptions.DESTINATION_B)); + + assertEquals( + Collections.singletonList(TestMoveOptions.DESTINATION_A), + write.getMoveOptionsFn().apply("a")); + assertEquals( + Collections.singletonList(TestMoveOptions.DESTINATION_B), + write.getMoveOptionsFn().apply("b")); + } + @Test @Category(NeedsRunner.class) public void testMatchAndMatchAll() throws IOException { diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystem.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystem.java index 1bee44eb38c0..6fbaa2f55990 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystem.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystem.java @@ -167,9 +167,22 @@ protected void rename( MoveOptions... moveOptions) throws IOException { Stopwatch stopwatch = Stopwatch.createStarted(); + @Nullable String destinationKmsKeyName = null; + for (MoveOptions moveOption : moveOptions) { + if (moveOption instanceof GcsMoveOptions) { + checkArgument( + destinationKmsKeyName == null, + "At most one GcsMoveOptions may be specified for a rename operation."); + destinationKmsKeyName = ((GcsMoveOptions) moveOption).destinationKmsKeyName(); + } + } options .getGcsUtil() - .rename(toFilenames(srcResourceIds), toFilenames(destResourceIds), moveOptions); + .rename( + toFilenames(srcResourceIds), + toFilenames(destResourceIds), + destinationKmsKeyName, + moveOptions); stopwatch.stop(); if (options.getGcsPerformanceMetrics()) { numRenames.inc(srcResourceIds.size()); diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsMoveOptions.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsMoveOptions.java new file mode 100644 index 000000000000..9f511b9bb768 --- /dev/null +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsMoveOptions.java @@ -0,0 +1,42 @@ +/* + * 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.extensions.gcp.storage; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.io.fs.MoveOptions; + +/** Google Cloud Storage-specific options for moving resources. */ +@AutoValue +public abstract class GcsMoveOptions implements MoveOptions { + + /** The Cloud KMS key to use to encrypt destination objects. */ + public abstract String destinationKmsKeyName(); + + /** Returns a {@link Builder}. */ + public static Builder builder() { + return new AutoValue_GcsMoveOptions.Builder(); + } + + /** A builder for {@link GcsMoveOptions}. */ + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setDestinationKmsKeyName(String destinationKmsKeyName); + + public abstract GcsMoveOptions build(); + } +} diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java index ed727d495cf8..7f76f756e423 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java @@ -371,6 +371,15 @@ public void rename( delegate.rename(srcFilenames, destFilenames, moveOptions); } + public void rename( + Iterable srcFilenames, + Iterable destFilenames, + @Nullable String destinationKmsKeyName, + MoveOptions... moveOptions) + throws IOException { + delegate.rename(srcFilenames, destFilenames, destinationKmsKeyName, moveOptions); + } + /** experimental api. */ public void renameV2( Iterable srcPaths, Iterable dstPaths, MoveOptions... moveOptions) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java index cfeb12dcae5c..f58c0608b8e9 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java @@ -712,7 +712,6 @@ public WritableByteChannel create(GcsPath path, CreateOptions options) throws IO if (options.getContentType() != null) { createBuilder = createBuilder.setContentType(options.getContentType()); } - HashMap baseLabels = new HashMap<>(); baseLabels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); baseLabels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); @@ -1076,7 +1075,12 @@ public void onFailure(GoogleJsonError e, HttpHeaders responseHeaders) }); } - public RewriteOp(GcsPath from, GcsPath to, boolean deleteSource, boolean ignoreMissingSource) + public RewriteOp( + GcsPath from, + GcsPath to, + boolean deleteSource, + boolean ignoreMissingSource, + @Nullable String destinationKmsKeyName) throws IOException { this.from = from; this.to = to; @@ -1086,6 +1090,9 @@ public RewriteOp(GcsPath from, GcsPath to, boolean deleteSource, boolean ignoreM storageClient .objects() .rewrite(from.getBucket(), from.getObject(), to.getBucket(), to.getObject(), null); + if (destinationKmsKeyName != null) { + rewriteRequest.setDestinationKmsKeyName(destinationKmsKeyName); + } if (maxBytesRewrittenPerCall != null) { rewriteRequest.setMaxBytesRewrittenPerCall(maxBytesRewrittenPerCall); } @@ -1173,12 +1180,22 @@ public void copy(Iterable srcFilenames, Iterable destFilenames) destFilenames, /*deleteSource=*/ false, /*ignoreMissingSource=*/ false, - /*ignoreExistingDest=*/ false); + /*ignoreExistingDest=*/ false, + /*destinationKmsKeyName=*/ null); } public void rename( Iterable srcFilenames, Iterable destFilenames, MoveOptions... moveOptions) throws IOException { + rename(srcFilenames, destFilenames, null, moveOptions); + } + + public void rename( + Iterable srcFilenames, + Iterable destFilenames, + @Nullable String destinationKmsKeyName, + MoveOptions... moveOptions) + throws IOException { // Rename is implemented as a rewrite followed by deleting the source. If the new object is in // the same location, the copy is a metadata-only operation. Set moveOptionSet = Sets.newHashSet(moveOptions); @@ -1187,7 +1204,12 @@ public void rename( final boolean ignoreExistingDest = moveOptionSet.contains(StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); rewriteHelper( - srcFilenames, destFilenames, /*deleteSource=*/ true, ignoreMissingSrc, ignoreExistingDest); + srcFilenames, + destFilenames, + /*deleteSource=*/ true, + ignoreMissingSrc, + ignoreExistingDest, + destinationKmsKeyName); } private void rewriteHelper( @@ -1195,11 +1217,17 @@ private void rewriteHelper( Iterable destFilenames, boolean deleteSource, boolean ignoreMissingSource, - boolean ignoreExistingDest) + boolean ignoreExistingDest, + @Nullable String destinationKmsKeyName) throws IOException { LinkedList rewrites = makeRewriteOps( - srcFilenames, destFilenames, deleteSource, ignoreMissingSource, ignoreExistingDest); + srcFilenames, + destFilenames, + deleteSource, + ignoreMissingSource, + ignoreExistingDest, + destinationKmsKeyName); org.apache.beam.sdk.util.BackOff backoff = BACKOFF_FACTORY.backoff(); while (true) { List batches = makeRewriteBatches(rewrites); // Removes completed rewrite ops. @@ -1247,6 +1275,18 @@ LinkedList makeRewriteOps( boolean ignoreMissingSource, boolean ignoreExistingDest) throws IOException { + return makeRewriteOps( + srcFilenames, destFilenames, deleteSource, ignoreMissingSource, ignoreExistingDest, null); + } + + LinkedList makeRewriteOps( + Iterable srcFilenames, + Iterable destFilenames, + boolean deleteSource, + boolean ignoreMissingSource, + boolean ignoreExistingDest, + @Nullable String destinationKmsKeyName) + throws IOException { List srcList = Lists.newArrayList(srcFilenames); List destList = Lists.newArrayList(destFilenames); checkArgument( @@ -1262,7 +1302,9 @@ LinkedList makeRewriteOps( throw new UnsupportedOperationException( "Skipping dest existence is only supported within a bucket."); } - rewrites.addLast(new RewriteOp(sourcePath, destPath, deleteSource, ignoreMissingSource)); + rewrites.addLast( + new RewriteOp( + sourcePath, destPath, deleteSource, ignoreMissingSource, destinationKmsKeyName)); } return rewrites; } diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystemTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystemTest.java index daa419abb576..4bfc632e7fcc 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystemTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystemTest.java @@ -20,6 +20,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; @@ -41,6 +42,7 @@ import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.fs.MatchResult; import org.apache.beam.sdk.io.fs.MatchResult.Status; +import org.apache.beam.sdk.io.fs.MoveOptions; import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.FluentIterable; @@ -123,6 +125,26 @@ public void testMatch() throws Exception { contains(toFilenames(matchResults.get(2)).toArray())); } + @Test + public void testRenameUsesDestinationKmsKey() throws IOException { + GcsResourceId source = GcsResourceId.fromGcsPath(GcsPath.fromUri("gs://testbucket/source")); + GcsResourceId destination = + GcsResourceId.fromGcsPath(GcsPath.fromUri("gs://testbucket/destination")); + String kmsKey = "projects/project/locations/location/keyRings/keyring/cryptoKeys/key"; + + gcsFileSystem.rename( + ImmutableList.of(source), + ImmutableList.of(destination), + GcsMoveOptions.builder().setDestinationKmsKeyName(kmsKey).build()); + + verify(mockGcsUtil) + .rename( + eq(ImmutableList.of(source.toString())), + eq(ImmutableList.of(destination.toString())), + eq(kmsKey), + any(MoveOptions[].class)); + } + @Test public void testGlobExpansion() throws IOException { Objects modelObjects = new Objects(); diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index 2f77f15dcffc..3087a9ef4898 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -1187,6 +1187,27 @@ public void testMakeRewriteOpsWithOptions() throws IOException { assertEquals(Long.valueOf(1337L), request.getMaxBytesRewrittenPerCall()); } + @Test + public void testMakeRewriteOpsWithDestinationKmsKey() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + LinkedList rewrites = + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 2), + makeStrings("d", 2), + true, + false, + false, + "projects/project/locations/location/keyRings/keyring/cryptoKeys/key"); + + assertEquals(2, rewrites.size()); + for (RewriteOp rewrite : rewrites) { + assertEquals( + "projects/project/locations/location/keyRings/keyring/cryptoKeys/key", + rewrite.rewriteRequest.getDestinationKmsKeyName()); + } + } + @Test public void testMakeRewriteBatches() throws IOException { GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();