Add per-destination GCS CMEK support to FileIO

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<MoveOptions> 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<ResourceId> srcFiles = new ArrayList<>();
-      List<ResourceId> dstFiles = new ArrayList<>();
+      Map<DestinationT, List<MoveOptions>> moveOptionsByDestination = Maps.newLinkedHashMap();
+      Map<List<MoveOptions>, List<KV<FileResult<DestinationT>, ResourceId>>> resultsByMoveOptions =
+          Maps.newLinkedHashMap();
       for (KV<FileResult<DestinationT>, ResourceId> entry : resultsToFinalFilenames) {
-        srcFiles.add(entry.getKey().getTempFilename());
-        dstFiles.add(entry.getValue());
+        DestinationT destination = entry.getKey().getDestination();
+        List<MoveOptions> 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<MoveOptions>, List<KV<FileResult<DestinationT>, ResourceId>>>
+          moveOptionsEntry : resultsByMoveOptions.entrySet()) {
+        List<ResourceId> srcFiles = new ArrayList<>();
+        List<ResourceId> dstFiles = new ArrayList<>();
+        for (KV<FileResult<DestinationT>, ResourceId> entry : moveOptionsEntry.getValue()) {
+          srcFiles.add(entry.getKey().getTempFilename());
+          dstFiles.add(entry.getValue());
+        }
+        List<MoveOptions> 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.

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<Fn<DestinationT, FileNaming>> getFileNamingFn();
 
+    abstract @Nullable SerializableFunction<DestinationT, List<MoveOptions>> getMoveOptionsFn();
+
     abstract @Nullable DestinationT getEmptyWindowDestination();
 
     abstract @Nullable Coder<DestinationT> getDestinationCoder();
@@ -1092,6 +1095,9 @@ abstract Builder<DestinationT, UserT> setOutputDirectory(
       abstract Builder<DestinationT, UserT> setFileNamingFn(
           Contextful<Fn<DestinationT, FileNaming>> namingFn);
 
+      abstract Builder<DestinationT, UserT> setMoveOptionsFn(
+          SerializableFunction<DestinationT, List<MoveOptions>> moveOptionsFn);
+
       abstract Builder<DestinationT, UserT> setEmptyWindowDestination(
           DestinationT emptyWindowDestination);
 
@@ -1266,6 +1272,13 @@ public Write<DestinationT, UserT> withNaming(
       return toBuilder().setFileNamingFn(namingFn).build();
     }
 
+    /** Specifies filesystem-specific move options for final files in each destination. */
+    public Write<DestinationT, UserT> withMoveOptions(
+        SerializableFunction<DestinationT, List<MoveOptions>> 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<DestinationT, UserT> withTempDirectory(String tempDirectory) {
       checkArgument(tempDirectory != null, "tempDirectory can not be null");
@@ -1494,6 +1507,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       }
 
       resolvedSpec.setFileNamingFn(resolveFileNamingFn());
+      if (getMoveOptionsFn() != null) {
+        resolvedSpec.setMoveOptionsFn(getMoveOptionsFn());
+      }
       resolvedSpec.setEmptyWindowDestination(getEmptyWindowDestination());
       if (getTempDirectory() == null) {
         checkArgument(
@@ -1721,6 +1737,13 @@ public List<PCollectionView<?>> getSideInputs() {
         public @Nullable Coder<DestinationT> getDestinationCoder() {
           return spec.getDestinationCoder();
         }
+
+        @Override
+        public List<MoveOptions> getMoveOptions(DestinationT destination) {
+          return spec.getMoveOptionsFn() == null
+              ? super.getMoveOptions(destination)
+              : spec.getMoveOptionsFn().apply(destination);
+        }
       }
     }
   }

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<String, String> write =
+        FileIO.<String, String>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 {

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());

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();
+  }
+}

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<String> srcFilenames,
+      Iterable<String> destFilenames,
+      @Nullable String destinationKmsKeyName,
+      MoveOptions... moveOptions)
+      throws IOException {
+    delegate.rename(srcFilenames, destFilenames, destinationKmsKeyName, moveOptions);
+  }
+
   /** experimental api. */
   public void renameV2(
       Iterable<GcsPath> srcPaths, Iterable<GcsPath> dstPaths, MoveOptions... moveOptions)

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<String, String> 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<String> srcFilenames, Iterable<String> destFilenames)
         destFilenames,
         /*deleteSource=*/ false,
         /*ignoreMissingSource=*/ false,
-        /*ignoreExistingDest=*/ false);
+        /*ignoreExistingDest=*/ false,
+        /*destinationKmsKeyName=*/ null);
   }
 
   public void rename(
       Iterable<String> srcFilenames, Iterable<String> destFilenames, MoveOptions... moveOptions)
       throws IOException {
+    rename(srcFilenames, destFilenames, null, moveOptions);
+  }
+
+  public void rename(
+      Iterable<String> srcFilenames,
+      Iterable<String> 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<MoveOptions> moveOptionSet = Sets.newHashSet(moveOptions);
@@ -1187,19 +1204,30 @@ 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(
       Iterable<String> srcFilenames,
       Iterable<String> destFilenames,
       boolean deleteSource,
       boolean ignoreMissingSource,
-      boolean ignoreExistingDest)
+      boolean ignoreExistingDest,
+      @Nullable String destinationKmsKeyName)
       throws IOException {
     LinkedList<RewriteOp> 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<BatchInterface> batches = makeRewriteBatches(rewrites); // Removes completed rewrite ops.
@@ -1247,6 +1275,18 @@ LinkedList<RewriteOp> makeRewriteOps(
       boolean ignoreMissingSource,
       boolean ignoreExistingDest)
       throws IOException {
+    return makeRewriteOps(
+        srcFilenames, destFilenames, deleteSource, ignoreMissingSource, ignoreExistingDest, null);
+  }
+
+  LinkedList<RewriteOp> makeRewriteOps(
+      Iterable<String> srcFilenames,
+      Iterable<String> destFilenames,
+      boolean deleteSource,
+      boolean ignoreMissingSource,
+      boolean ignoreExistingDest,
+      @Nullable String destinationKmsKeyName)
+      throws IOException {
     List<String> srcList = Lists.newArrayList(srcFilenames);
     List<String> destList = Lists.newArrayList(destFilenames);
     checkArgument(
@@ -1262,7 +1302,9 @@ LinkedList<RewriteOp> 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;
   }