You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2025/10/01 13:57:46 UTC
(beam) branch users/damccorm/enforce-java-gbek created (now 90fa3fd2887)
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a change to branch users/damccorm/enforce-java-gbek
in repository https://gitbox.apache.org/repos/asf/beam.git
at 90fa3fd2887 [WIP] Add pipeline option to force GBEK (Java)
This branch includes the following new commits:
new 90fa3fd2887 [WIP] Add pipeline option to force GBEK (Java)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(beam) 01/01: [WIP] Add pipeline option to force GBEK (Java)
Posted by da...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch users/damccorm/enforce-java-gbek
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 90fa3fd2887acfdf221d6c3ba9fa3a2a0c47db35
Author: Danny Mccormick <da...@google.com>
AuthorDate: Wed Oct 1 09:57:27 2025 -0400
[WIP] Add pipeline option to force GBEK (Java)
---
.../dataflow/internal/DataflowGroupByKey.java | 49 ++++++++++++++++
.../apache/beam/sdk/options/PipelineOptions.java | 35 +++++++++++
.../beam/sdk/transforms/GroupByEncryptedKey.java | 37 ++++++++++--
.../org/apache/beam/sdk/transforms/GroupByKey.java | 34 +++++++++++
.../java/org/apache/beam/sdk/util/GcpSecret.java | 9 +++
.../main/java/org/apache/beam/sdk/util/Secret.java | 49 ++++++++++++++++
.../util/construction/GroupByKeyTranslation.java | 12 ++++
.../java/org/apache/beam/sdk/util/SecretTest.java | 67 ++++++++++++++++++++++
8 files changed, 288 insertions(+), 4 deletions(-)
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/internal/DataflowGroupByKey.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/internal/DataflowGroupByKey.java
index 89135641689..82affe15f0c 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/internal/DataflowGroupByKey.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/internal/DataflowGroupByKey.java
@@ -25,10 +25,13 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.Coder.NonDeterministicException;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.GroupByEncryptedKey;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.util.Secret;
import org.apache.beam.sdk.util.construction.PTransformTranslation;
import org.apache.beam.sdk.util.construction.SdkComponents;
import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
@@ -46,9 +49,13 @@ public class DataflowGroupByKey<K, V>
// Plumbed from Redistribute transform.
private final boolean allowDuplicates;
+ private boolean insideGBEK;
+ private boolean surroundsGBEK;
private DataflowGroupByKey(boolean allowDuplicates) {
this.allowDuplicates = allowDuplicates;
+ this.insideGBEK = false;
+ this.surroundsGBEK = false;
}
/**
@@ -79,6 +86,22 @@ public class DataflowGroupByKey<K, V>
return allowDuplicates;
}
+ /**
+ * For Beam internal use only. Tells runner that this is an inner GBK inside of a
+ * GroupByEncryptedKey
+ */
+ public void setInsideGBEK() {
+ this.insideGBEK = true;
+ }
+
+ /**
+ * For Beam internal use only. Tells runner that this is a GBK wrapped around of a
+ * GroupByEncryptedKey
+ */
+ public boolean surroundsGBEK() {
+ return this.surroundsGBEK;
+ }
+
/////////////////////////////////////////////////////////////////////////////
public static void applicableTo(PCollection<?> input) {
@@ -117,6 +140,20 @@ public class DataflowGroupByKey<K, V>
"the keyCoder of a DataflowGroupByKey must be deterministic", e);
}
+ PipelineOptions options = input.getPipeline().getOptions();
+ String gbekOveride = options.getGBEK();
+ if (!this.insideGBEK && gbekOveride != null && !gbekOveride.trim().isEmpty()) {
+ this.surroundsGBEK = true;
+ Secret hmacSecret = Secret.parseSecretOption(gbekOveride);
+ DataflowGroupByKey<byte[], KV<byte[], byte[]>> gbk = DataflowGroupByKey.create();
+ if (this.allowDuplicates) {
+ gbk = DataflowGroupByKey.createWithAllowDuplicates();
+ }
+ gbk.setInsideGBEK();
+ GroupByEncryptedKey<K, V> gbek = GroupByEncryptedKey.createWithCustomGbk(hmacSecret, gbk);
+ return input.apply(gbek);
+ }
+
// This primitive operation groups by the combination of key and window,
// merging windows as needed, using the windows assigned to the
// key/value input elements and the window merge operation of the
@@ -171,10 +208,22 @@ public class DataflowGroupByKey<K, V>
return PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
}
+ @Override
+ public String getUrn(DataflowGroupByKey<?, ?> transform) {
+ if (transform.surroundsGBEK()) {
+ return "beam:transform:group_by_key_wrapper:v1";
+ }
+ return PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
+ }
+
@Override
@SuppressWarnings("nullness")
public RunnerApi.FunctionSpec translate(
AppliedPTransform<?, ?, DataflowGroupByKey<?, ?>> transform, SdkComponents components) {
+ if (transform.getTransform().surroundsGBEK()) {
+ // Can use null for spec for empty composite.
+ return null;
+ }
return RunnerApi.FunctionSpec.newBuilder().setUrn(getUrn(transform.getTransform())).build();
}
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptions.java
index 2eba8c6ef68..8d34345306c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptions.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptions.java
@@ -37,6 +37,7 @@ import org.apache.beam.sdk.transforms.display.HasDisplayData;
import org.apache.beam.sdk.util.ReleaseInfo;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.DateTimeUtils;
import org.joda.time.DateTimeZone;
import org.joda.time.format.DateTimeFormat;
@@ -413,6 +414,40 @@ public interface PipelineOptions extends HasDisplayData {
void setUserAgent(String userAgent);
+ /**
+ * A string defining whether GroupByKey transforms should be replaced by GroupByEncryptedKey
+ *
+ * <p>Beam will infer the secret type and value based on the secret itself. This guarantees that
+ * any data at rest during the performing a GBK, so this can be used to guarantee that data is not
+ * unencrypted. Runners with this behavior include the Dataflow, Flink, and Spark runners. The
+ * option should be structured like:
+ *
+ * <pre><code>
+ * --encrypt=type:<secret_type>;<secret_param>:<value>
+ * </code></pre>
+ *
+ * for example:
+ *
+ * <pre><code>
+ * --encrypt=type:GcpSecret;version_name:my_secret/versions/latest"
+ * </code></pre>
+ *
+ * All variables should use snake case to allow consistency across languages.
+ */
+ @Description(
+ "When set, will replace all GroupByKey transforms in the pipeline the option. Beam will"
+ + " infer the secret type and value based on the secret itself. This guarantees that"
+ + " any data at rest during the performing a GBK, so this can be used to guarantee"
+ + " that data is not unencrypted. Runners with this behavior include the Dataflow,"
+ + " Flink, and Spark runners. The option should be structured like:"
+ + " --encrypt=type:<secret_type>;<secret_param>:<value>, for example "
+ + " --encrypt=type:GcpSecret;version_name:my_secret/versions/latest. All variables "
+ + " should use snake case to allow consistency across languages.")
+ @Nullable
+ String getGBEK();
+
+ void setGBEK(String gbek);
+
/**
* Returns a user agent string constructed from {@link ReleaseInfo#getName()} and {@link
* ReleaseInfo#getVersion()}, in the format {@code [name]/[version]}.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByEncryptedKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByEncryptedKey.java
index 0e5a873799a..56806633c3c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByEncryptedKey.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByEncryptedKey.java
@@ -52,9 +52,19 @@ public class GroupByEncryptedKey<K, V>
extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> {
private final Secret hmacKey;
-
- private GroupByEncryptedKey(Secret hmacKey) {
+ private final PTransform<
+ PCollection<KV<byte[], KV<byte[], byte[]>>>,
+ PCollection<KV<byte[], Iterable<KV<byte[], byte[]>>>>>
+ gbk;
+
+ private GroupByEncryptedKey(
+ Secret hmacKey,
+ PTransform<
+ PCollection<KV<byte[], KV<byte[], byte[]>>>,
+ PCollection<KV<byte[], Iterable<KV<byte[], byte[]>>>>>
+ gbk) {
this.hmacKey = hmacKey;
+ this.gbk = gbk;
}
/**
@@ -66,7 +76,26 @@ public class GroupByEncryptedKey<K, V>
* @return A {@link GroupByEncryptedKey} transform.
*/
public static <K, V> GroupByEncryptedKey<K, V> create(Secret hmacKey) {
- return new GroupByEncryptedKey<>(hmacKey);
+ GroupByKey<byte[], KV<byte[], byte[]>> gbk = GroupByKey.create();
+ return new GroupByEncryptedKey<>(hmacKey, GroupByKey.create());
+ }
+
+ /**
+ * Creates a {@link GroupByEncryptedKey} transform with a custom GBK in the middle.
+ *
+ * @param hmacKey The {@link Secret} key to use for encryption.
+ * @param gbk The custom GBK transform to use in the middle of the GBEK.
+ * @param <K> The type of the keys in the input PCollection.
+ * @param <V> The type of the values in the input PCollection.
+ * @return A {@link GroupByEncryptedKey} transform.
+ */
+ public static <K, V> GroupByEncryptedKey<K, V> createWithCustomGbk(
+ Secret hmacKey,
+ PTransform<
+ PCollection<KV<byte[], KV<byte[], byte[]>>>,
+ PCollection<KV<byte[], Iterable<KV<byte[], byte[]>>>>>
+ gbk) {
+ return new GroupByEncryptedKey<>(hmacKey, gbk);
}
@Override
@@ -92,7 +121,7 @@ public class GroupByEncryptedKey<K, V>
.apply(
"EncryptMessage",
ParDo.of(new EncryptMessage<>(this.hmacKey, keyCoder, valueCoder)))
- .apply(GroupByKey.create());
+ .apply(this.gbk);
return grouped
.apply("DecryptMessage", ParDo.of(new DecryptMessage<>(this.hmacKey, keyCoder, valueCoder)))
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
index d0b320a8765..da3fcd1de6a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
@@ -32,6 +32,7 @@ import org.apache.beam.sdk.transforms.windowing.Never.NeverTrigger;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.util.Secret;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
@@ -115,9 +116,13 @@ public class GroupByKey<K, V>
extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> {
private final boolean fewKeys;
+ private boolean insideGBEK;
+ private boolean surroundsGBEK;
private GroupByKey(boolean fewKeys) {
this.fewKeys = fewKeys;
+ this.insideGBEK = false;
+ surroundsGBEK = false;
}
/**
@@ -148,6 +153,21 @@ public class GroupByKey<K, V>
return fewKeys;
}
+ /**
+ * For Beam internal use only. Tells runner that this is an inner GBK inside a GroupByEncryptedKey
+ */
+ public void setInsideGBEK() {
+ this.insideGBEK = true;
+ }
+
+ /**
+ * For Beam internal use only. Tells runner that this is a GBK wrapped around of a
+ * GroupByEncryptedKey
+ */
+ public boolean surroundsGBEK() {
+ return this.surroundsGBEK;
+ }
+
/////////////////////////////////////////////////////////////////////////////
public static void applicableTo(PCollection<?> input) {
@@ -244,6 +264,20 @@ public class GroupByKey<K, V>
throw new IllegalStateException("the keyCoder of a GroupByKey must be deterministic", e);
}
+ PipelineOptions options = input.getPipeline().getOptions();
+ String gbekOveride = options.getGBEK();
+ if (!this.insideGBEK && gbekOveride != null && !gbekOveride.trim().isEmpty()) {
+ this.surroundsGBEK = true;
+ Secret hmacSecret = Secret.parseSecretOption(gbekOveride);
+ GroupByKey<byte[], KV<byte[], byte[]>> gbk = GroupByKey.create();
+ if (this.fewKeys) {
+ gbk = GroupByKey.createWithFewKeys();
+ }
+ gbk.setInsideGBEK();
+ GroupByEncryptedKey<K, V> gbek = GroupByEncryptedKey.createWithCustomGbk(hmacSecret, gbk);
+ return input.apply(gbek);
+ }
+
// This primitive operation groups by the combination of key and window,
// merging windows as needed, using the windows assigned to the
// key/value input elements and the window merge operation of the
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/GcpSecret.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/GcpSecret.java
index 80bc3a54535..8effae7f61c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/GcpSecret.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/GcpSecret.java
@@ -55,4 +55,13 @@ public class GcpSecret implements Secret {
throw new RuntimeException("Failed to retrieve secret bytes", e);
}
}
+
+ /**
+ * Returns the version name of the secret.
+ *
+ * @return The version name as a String.
+ */
+ public String getVersionName() {
+ return versionName;
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Secret.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Secret.java
index fe476ef6cb1..a75e01c9543 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Secret.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Secret.java
@@ -18,6 +18,11 @@
package org.apache.beam.sdk.util;
import java.io.Serializable;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
/**
* A secret management interface used for handling sensitive data.
@@ -33,4 +38,48 @@ public interface Secret extends Serializable {
* @return The secret as a byte array.
*/
byte[] getSecretBytes();
+
+ static Secret parseSecretOption(String secretOption) {
+ Map<String, String> paramMap = new HashMap<>();
+ for (String param : secretOption.split(";", -1)) {
+ String[] parts = param.split(":", 2);
+ if (parts.length == 2) {
+ paramMap.put(parts[0], parts[1]);
+ }
+ }
+
+ if (!paramMap.containsKey("type")) {
+ throw new RuntimeException("Secret string must contain a valid type parameter");
+ }
+
+ String secretType = paramMap.get("type");
+ paramMap.remove("type");
+
+ if (secretType == null) {
+ throw new RuntimeException("Secret string must contain a valid value for type parameter");
+ }
+
+ switch (secretType.toLowerCase()) {
+ case "gcpsecret":
+ Set<String> gcpSecretParams = new HashSet<>(Arrays.asList("version_name"));
+ for (String paramName : paramMap.keySet()) {
+ if (!gcpSecretParams.contains(paramName)) {
+ throw new RuntimeException(
+ String.format(
+ "Invalid secret parameter %s, GcpSecret only supports the following parameters: %s",
+ paramName, gcpSecretParams));
+ }
+ }
+ String versionName = paramMap.get("version_name");
+ if (versionName == null) {
+ throw new RuntimeException(
+ "version_name must contain a valid value for versionName parameter");
+ }
+ return new GcpSecret(versionName);
+ default:
+ throw new RuntimeException(
+ String.format(
+ "Invalid secret type %s, currently only GcpSecret is supported", secretType));
+ }
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/GroupByKeyTranslation.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/GroupByKeyTranslation.java
index d08a48d0e5e..1471c790a68 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/GroupByKeyTranslation.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/GroupByKeyTranslation.java
@@ -43,9 +43,21 @@ public class GroupByKeyTranslation {
return PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
}
+ @Override
+ public String getUrn(GroupByKey<?, ?> transform) {
+ if (transform.surroundsGBEK()) {
+ return "beam:transform:group_by_key_wrapper:v1";
+ }
+ return PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
+ }
+
@Override
public FunctionSpec translate(
AppliedPTransform<?, ?, GroupByKey<?, ?>> transform, SdkComponents components) {
+ if (transform.getTransform().surroundsGBEK()) {
+ // Can use null for spec for empty composite.
+ return null;
+ }
return FunctionSpec.newBuilder().setUrn(getUrn(transform.getTransform())).build();
}
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/SecretTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/SecretTest.java
new file mode 100644
index 00000000000..dd4b125d73f
--- /dev/null
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/SecretTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.util;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link org.apache.beam.sdk.util.Secret}. */
+@RunWith(JUnit4.class)
+public class SecretTest {
+
+ @Test
+ public void testParseSecretOptionWithValidGcpSecret() {
+ String secretOption = "type:gcpsecret;version_name:my_secret/versions/latest";
+ Secret secret = Secret.parseSecretOption(secretOption);
+ assertTrue(secret instanceof GcpSecret);
+ assertEquals("my_secret/versions/latest", ((GcpSecret) secret).getVersionName());
+ }
+
+ @Test
+ public void testParseSecretOptionWithMissingType() {
+ String secretOption = "version_name:my_secret/versions/latest";
+ Exception exception =
+ assertThrows(RuntimeException.class, () -> Secret.parseSecretOption(secretOption));
+ assertEquals("Secret string must contain a valid type parameter", exception.getMessage());
+ }
+
+ @Test
+ public void testParseSecretOptionWithUnsupportedType() {
+ String secretOption = "type:unsupported;version_name:my_secret/versions/latest";
+ Exception exception =
+ assertThrows(RuntimeException.class, () -> Secret.parseSecretOption(secretOption));
+ assertEquals(
+ "Invalid secret type unsupported, currently only GcpSecret is supported",
+ exception.getMessage());
+ }
+
+ @Test
+ public void testParseSecretOptionWithInvalidGcpSecretParameter() {
+ String secretOption = "type:gcpsecret;invalid_param:some_value";
+ Exception exception =
+ assertThrows(RuntimeException.class, () -> Secret.parseSecretOption(secretOption));
+ assertEquals(
+ "Invalid secret parameter invalid_param, GcpSecret only supports the following parameters: [version_name]",
+ exception.getMessage());
+ }
+}