diff --git a/build.gradle.kts b/build.gradle.kts
index f72e12af176e..04fd7ee3a218 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -353,6 +353,7 @@ tasks.register("javaioPreCommit") {
dependsOn(":sdks:java:io:jms:build")
dependsOn(":sdks:java:io:kafka:build")
dependsOn(":sdks:java:io:kafka:upgrade:build")
+ dependsOn(":sdks:java:io:kafka:file-aware-factories:build")
dependsOn(":sdks:java:io:kudu:build")
dependsOn(":sdks:java:io:mongodb:build")
dependsOn(":sdks:java:io:mqtt:build")
diff --git a/sdks/java/extensions/kafka-factories/build.gradle b/sdks/java/extensions/kafka-factories/build.gradle
new file mode 100644
index 000000000000..30c5d3fd6642
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/build.gradle
@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+plugins { id 'org.apache.beam.module' }
+applyJavaNature(
+ automaticModuleName: 'org.apache.beam.sdk.extensions.kafka.factories',
+ publish: 'False'
+)
+
+description = "Apache Beam :: SDKs :: Java :: Extensions :: Kafka :: Factories"
+ext.summary = "Library to instantiate kafka clients with files from GCS or SecretManager."
+
+dependencies {
+ // ------------------------- CORE DEPENDENCIES -------------------------
+ implementation project(path: ":sdks:java:core", configuration: "shadow")
+ provided library.java.kafka_clients
+ implementation 'com.google.cloud:google-cloud-secretmanager:2.72.0'
+ implementation library.java.slf4j_api
+ implementation library.java.vendored_guava_32_1_2_jre
+ implementation project(path: ":sdks:java:extensions:google-cloud-platform-core")
+ permitUnusedDeclared project(path: ":sdks:java:extensions:google-cloud-platform-core")
+ // ------------------------- TEST DEPENDENCIES -------------------------
+ testImplementation 'org.apache.kafka:kafka-clients:3.9.0'
+ testImplementation library.java.junit
+ testImplementation library.java.mockito_core
+ testRuntimeOnly library.java.mockito_inline
+ testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
+}
diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java
new file mode 100644
index 000000000000..a0f15b42382d
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java
@@ -0,0 +1,258 @@
+/*
+ * 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.kafka.factories;
+
+import com.google.cloud.secretmanager.v1.AccessSecretVersionResponse;
+import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
+import com.google.cloud.secretmanager.v1.SecretVersionName;
+import java.io.File;
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * An abstract {@link SerializableFunction} that serves as a base class for factories that need to
+ * process a configuration map to handle external resources like files and secrets.
+ *
+ *
This class is designed to be extended by concrete factory implementations (e.g., for creating
+ * Kafka consumers). It automates the process of detecting special URI strings within the
+ * configuration values and transforming them before passing the processed configuration to the
+ * subclass.
+ *
+ *
Supported Patterns:
+ *
+ *
+ * External File Paths: It recognizes paths prefixed with schemes like {@code gs://} or
+ * {@code s3://} that are supported by the Beam {@link FileSystems} API. It downloads these
+ * files to a local temporary directory (under {@code /tmp//...}) and replaces
+ * the original path in the configuration with the new local file path.
+ * Secret Manager Values: It recognizes strings prefixed with {@code secretValue:}. It
+ * interprets the rest of the string as a Google Secret Manager secret version name (e.g.,
+ * "projects/p/secrets/s/versions/v"), fetches the secret payload, and replaces the original
+ * {@code secretValue:...} identifier with the plain-text secret.
+ *
+ *
+ * Usage:
+ *
+ * A subclass must implement the {@link #createObject(Map)} method, which receives the fully
+ * processed configuration map with all paths localized and secrets resolved. Subclasses can also
+ * override {@link #downloadAndProcessExtraFiles()} to handle specific preliminary file downloads
+ * (e.g., a krb5.conf file) before the main configuration processing begins.
+ *
+ * @param The type of object this factory creates.
+ */
+public abstract class FileAwareFactoryFn
+ implements SerializableFunction, T> {
+
+ public static final String SECRET_VALUE_PREFIX = "secretValue:";
+ public static final String DIRECTORY_PREFIX = "/tmp";
+ private static final Pattern PATH_PATTERN =
+ Pattern.compile("([a-zA-Z0-9]+://[^\"]+)|(secretValue:[^\"]+)|(secretFile:[^\"]+)");
+
+ private static final Map secretCache = new ConcurrentHashMap<>();
+
+ private final String factoryType;
+ private static final Logger LOG = LoggerFactory.getLogger(FileAwareFactoryFn.class);
+
+ public FileAwareFactoryFn(String factoryType) {
+ Preconditions.checkNotNull(factoryType);
+ this.factoryType = factoryType;
+ }
+
+ protected abstract T createObject(Map config);
+
+ @Override
+ public T apply(Map config) {
+ if (config == null) {
+ return createObject(config);
+ }
+
+ Map processedConfig = new HashMap<>(config);
+
+ String key = "";
+ Object value = null;
+ try {
+ downloadAndProcessExtraFiles();
+
+ for (Map.Entry e : config.entrySet()) {
+ try {
+ key = e.getKey();
+ value = e.getValue();
+ if (value instanceof String) {
+ String originalValue = (String) value;
+ Matcher matcher = PATH_PATTERN.matcher(originalValue);
+ StringBuffer sb = new StringBuffer();
+
+ while (matcher.find()) {
+ String externalPath = matcher.group(1);
+ String secretValue = matcher.group(2);
+ String secretFile = matcher.group(3);
+
+ if (externalPath != null) {
+ try {
+ String tmpPath = replacePathWithLocal(externalPath);
+ String localPath = downloadExternalFile(externalPath, tmpPath);
+ matcher.appendReplacement(sb, Matcher.quoteReplacement(localPath));
+ LOG.info("Downloaded {} to {}", externalPath, localPath);
+ } catch (IOException io) {
+ throw new IOException("Failed to download file : " + externalPath, io);
+ }
+ } else if (secretValue != null) {
+ try {
+ String secretId = secretValue.substring(SECRET_VALUE_PREFIX.length());
+ String processedSecret =
+ processSecret(originalValue, secretId, getSecretWithCache(secretId));
+
+ matcher.appendReplacement(sb, Matcher.quoteReplacement(processedSecret));
+ } catch (IllegalArgumentException ia) {
+ throw new IllegalArgumentException("Failed to get secret.", ia);
+ }
+ } else if (secretFile != null) {
+ throw new UnsupportedOperationException("Not yet implemented.");
+ }
+ }
+ matcher.appendTail(sb);
+ String processedValue = sb.toString();
+ processedConfig.put(key, processedValue);
+ }
+ } catch (IOException ex) {
+ throw new RuntimeException("Failed trying to process value for key " + key + ".", ex);
+ }
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed trying to process extra files.", e);
+ }
+
+ return createObject(processedConfig);
+ }
+
+ /**
+ * A function to download files from their specified external storage path and copy them to the
+ * provided local filepath. The local filepath is provided by the replacePathWithLocal.
+ *
+ * @param externalFilePath
+ * @param outputFileString
+ * @return
+ * @throws IOException
+ */
+ protected static synchronized String downloadExternalFile(
+ String externalFilePath, String outputFileString) throws IOException {
+ // create the file only if it doesn't exist
+ if (new File(outputFileString).exists()) {
+ return outputFileString;
+ }
+ Path outputFilePath = Paths.get(outputFileString);
+ Path parentDir = outputFilePath.getParent();
+ if (parentDir != null) {
+ Files.createDirectories(parentDir);
+ }
+ LOG.info("Staging external file [{}] to [{}]", externalFilePath, outputFileString);
+ Set options = new HashSet<>(2);
+ options.add(StandardOpenOption.CREATE);
+ options.add(StandardOpenOption.WRITE);
+
+ // Copy the external file into a local file and will throw an I/O exception in case file not
+ // found.
+ try (ReadableByteChannel readerChannel =
+ FileSystems.open(FileSystems.matchSingleFileSpec(externalFilePath).resourceId())) {
+ try (FileChannel writeChannel = FileChannel.open(outputFilePath, options)) {
+ writeChannel.transferFrom(readerChannel, 0, Long.MAX_VALUE);
+ }
+ }
+ return outputFileString;
+ }
+
+ protected byte[] getSecretWithCache(String secretId) {
+ return secretCache.computeIfAbsent(secretId, this::getSecret);
+ }
+
+ /**
+ * A helper method to create a new string with the external paths replaced with their local path
+ * and subdirectory based on the factory type in the /tmp directory. For example, the kerberos
+ * factory type will replace the file paths with /tmp/kerberos/file.path
+ *
+ * @param externalPath
+ * @return a string with all instances of external paths converted to the local paths where the
+ * files sit.
+ */
+ private String replacePathWithLocal(String externalPath) throws IOException {
+ String externalBucketPrefixIdentifier = "://";
+ int externalBucketPrefixIndex = externalPath.lastIndexOf(externalBucketPrefixIdentifier);
+ if (externalBucketPrefixIndex == -1) {
+ // if we don't find a known bucket prefix then we will error early.
+ throw new RuntimeException(
+ "The provided external bucket could not be matched to a known source.");
+ }
+
+ int prefixLength = externalBucketPrefixIndex + externalBucketPrefixIdentifier.length();
+ return DIRECTORY_PREFIX + "/" + factoryType + "/" + externalPath.substring(prefixLength);
+ }
+
+ /**
+ * @throws IOException A hook for subclasses to download and process specific files before the
+ * main configuration is handled. For example, the kerberos factory can use this to download a
+ * krb5.conf and set a system property.
+ */
+ protected void downloadAndProcessExtraFiles() throws IOException {
+ // Default implementation should do nothing.
+ }
+
+ protected String getBaseDirectory() {
+ return DIRECTORY_PREFIX;
+ }
+
+ protected byte[] getSecret(String secretVersion) {
+ SecretVersionName secretVersionName;
+ if (SecretVersionName.isParsableFrom(secretVersion)) {
+ secretVersionName = SecretVersionName.parse(secretVersion);
+ } else {
+ throw new IllegalArgumentException(
+ "Provided Secret must be in the form"
+ + " projects/{project}/secrets/{secret}/versions/{secret_version}");
+ }
+ try (SecretManagerServiceClient client = SecretManagerServiceClient.create()) {
+ AccessSecretVersionResponse response = client.accessSecretVersion(secretVersionName);
+ return response.getPayload().getData().toByteArray();
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ protected String processSecret(String originalValue, String secretId, byte[] secretValue) {
+ // By Default, this will return the secret value directly. This function can be overridden by
+ // derived classes.
+ return new String(secretValue, StandardCharsets.UTF_8);
+ }
+}
diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java
new file mode 100644
index 000000000000..409904b667f9
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java
@@ -0,0 +1,150 @@
+/*
+ * 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.kafka.factories;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.PosixFilePermission;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.security.auth.login.Configuration;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class KerberosConsumerFactoryFn extends FileAwareFactoryFn> {
+ private static final String LOCAL_FACTORY_TYPE = "kerberos";
+ private String krb5ConfigPath = "";
+ private static volatile String localKrb5ConfPath = "";
+
+ private static final Object lock = new Object();
+
+ // Standard Kafka property for SASL JAAS configuration
+ private static final String JAAS_CONFIG_PROPERTY = "sasl.jaas.config";
+ private static final String KEYTAB_SECRET_PREFIX = "keyTab=\"secretValue:";
+ private static final Pattern KEYTAB_SECRET_PATTERN =
+ Pattern.compile("(keyTab=\"secretValue:[^\"]+)");
+
+ private static final Logger LOG = LoggerFactory.getLogger(KerberosConsumerFactoryFn.class);
+
+ public KerberosConsumerFactoryFn(String krb5ConfigPath) {
+ super("kerberos");
+ this.krb5ConfigPath = krb5ConfigPath;
+ }
+
+ @Override
+ protected Consumer createObject(Map config) {
+ // This will be called after the config map processing has occurred. Therefore, we know that the
+ // property will have had it's value replaced with a local directory.
+ // We don't need to worry about the external bucket prefix in this case.
+ try {
+ String jaasConfig = (String) config.get(JAAS_CONFIG_PROPERTY);
+ String localKeytabPath = "";
+ if (jaasConfig != null && !jaasConfig.isEmpty()) {
+ localKeytabPath =
+ jaasConfig.substring(
+ jaasConfig.indexOf("keyTab=\"") + 8, jaasConfig.lastIndexOf("\" principal"));
+ }
+
+ // Set the permissions on the file to be as strict as possible for security reasons. The
+ // keytab contains sensitive information and should be as locked down as possible.
+ Path path = Paths.get(localKeytabPath);
+ Set perms = new HashSet<>();
+ perms.add(PosixFilePermission.OWNER_READ);
+ Files.setPosixFilePermissions(path, perms);
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Could not access keytab file. Make sure that the sasl.jaas.config config property "
+ + "is set correctly.",
+ e);
+ }
+ return new KafkaConsumer<>(config);
+ }
+
+ @Override
+ protected void downloadAndProcessExtraFiles() throws IOException {
+ synchronized (lock) {
+ // we only want a new krb5 file if there is not already one present.
+ if (localKrb5ConfPath.isEmpty()) {
+ if (this.krb5ConfigPath != null && !this.krb5ConfigPath.isEmpty()) {
+ String localPath =
+ super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + "krb5.conf";
+ localKrb5ConfPath = downloadExternalFile(this.krb5ConfigPath, localPath);
+
+ System.setProperty("java.security.krb5.conf", localKrb5ConfPath);
+ Configuration.getConfiguration().refresh();
+ }
+ }
+ }
+ }
+
+ @Override
+ protected String processSecret(String originalValue, String secretId, byte[] secretValue)
+ throws RuntimeException {
+ Matcher matcher = KEYTAB_SECRET_PATTERN.matcher(originalValue);
+ String localFileString = "";
+ while (matcher.find()) {
+ String currentSecretId = matcher.group(1);
+ if (currentSecretId == null || currentSecretId.isEmpty()) {
+ throw new RuntimeException(
+ "Error matching values. Secret was discovered but its value is null");
+ }
+ currentSecretId = currentSecretId.substring(KEYTAB_SECRET_PREFIX.length());
+ if (!currentSecretId.equals(secretId)) {
+ // A sasl.jaas.config can contain multiple keytabs in one string. Therefore, we must assume
+ // that there can
+ // also be multiple keytab secrets in the same string. If the currently matched secret does
+ // not equal
+ // the secret that we are processing (passed in via secretId) then we do not want to create
+ // a keytab file and overwrite it.
+ continue;
+ }
+ String filename = "kafka-client-" + UUID.randomUUID().toString() + ".keytab";
+
+ localFileString = super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + filename;
+ Path localFilePath = Paths.get(localFileString);
+ Path parentDir = localFilePath.getParent();
+ try {
+ if (parentDir != null) {
+ Files.createDirectories(parentDir);
+ }
+ Files.write(localFilePath, secretValue);
+ if (!new File(localFileString).canRead()) {
+ LOG.warn("The file is not readable");
+ }
+ LOG.info("Successfully wrote file to path: {}", localFilePath);
+ } catch (IOException e) {
+ throw new RuntimeException("Unable to create the keytab file for the provided secret.");
+ }
+ }
+ // if no localFile was created, then we can assume that the secret is meant to be kept as a
+ // value.
+ return localFileString.isEmpty()
+ ? new String(secretValue, StandardCharsets.UTF_8)
+ : localFileString;
+ }
+}
diff --git a/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java
new file mode 100644
index 000000000000..da12c8203a64
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+/** ConsumerFactoryFns for file paths that exist in GCS or Google SecretManager. */
+package org.apache.beam.sdk.extensions.kafka.factories;
diff --git a/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java
new file mode 100644
index 000000000000..0ad096e856dc
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java
@@ -0,0 +1,268 @@
+/*
+ * 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.kafka.factories;
+
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.IOException;
+import java.nio.channels.Channels;
+import java.nio.channels.FileChannel;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.MatchResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentMatchers;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class FileAwareFactoryFnTest {
+
+ @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private TestFactoryFn factory;
+ private String baseDir;
+ private static final String TEST_FACTORY_TYPE = "test-factory";
+
+ // A concrete implementation for testing the abstract FileAwareFactoryFn
+ static class TestFactoryFn extends FileAwareFactoryFn {
+ public TestFactoryFn() {
+ super(TEST_FACTORY_TYPE);
+ }
+
+ @Override
+ protected Object createObject(Map config) {
+ // Return the processed config for easy assertion
+ return config;
+ }
+ }
+
+ @Before
+ public void setup() throws IOException {
+ baseDir = "/tmp/" + TEST_FACTORY_TYPE;
+ factory = Mockito.spy(new TestFactoryFn());
+ Mockito.doReturn(baseDir).when(factory).getBaseDirectory();
+ }
+
+ @Test
+ public void testHappyPathReplacesExternalPath() {
+ // Arrange
+ String gcsPath = "gs://test-bucket/config-file.json";
+ String expectedLocalPath =
+ FileAwareFactoryFn.DIRECTORY_PREFIX
+ + "/"
+ + TEST_FACTORY_TYPE
+ + "/test-bucket/config-file.json";
+ Map config = new HashMap<>();
+ config.put("config.file.path", gcsPath);
+
+ // Act & Assert
+ // Use try-with-resources to manage the scope of the static mock on FileSystems
+ try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) {
+ // 1. Mock the underlying static FileSystems calls to avoid real network I/O
+ MatchResult.Metadata metadata = Mockito.mock(MatchResult.Metadata.class);
+ ResourceId resourceId = Mockito.mock(ResourceId.class);
+ Mockito.when(metadata.resourceId()).thenReturn(resourceId);
+ mockedFileSystems.when(() -> FileSystems.matchSingleFileSpec(gcsPath)).thenReturn(metadata);
+
+ // 2. Mock 'open' to return a channel with no data, simulating a successful download
+ ReadableByteChannel channel = Channels.newChannel(new ByteArrayInputStream(new byte[0]));
+ mockedFileSystems.when(() -> FileSystems.open(resourceId)).thenReturn(channel);
+
+ // Act
+ Map processedConfig = (Map) factory.apply(config);
+
+ // Assert
+ Assert.assertEquals(expectedLocalPath, processedConfig.get("config.file.path"));
+ Assert.assertTrue(
+ "Local file should have been created", new File(expectedLocalPath).exists());
+ }
+ }
+
+ @Test
+ public void testApplyFailurePathThrowsRuntimeExceptionOnDownloadFailure() {
+ // Arrange
+ String gcsPath = "gs://test-bucket/failing-file.txt";
+ Map config = new HashMap<>();
+ config.put("critical.file", gcsPath);
+
+ // Mock the static FileSystems.matchSingleFileSpec to throw an exception
+ try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) {
+ mockedFileSystems
+ .when(() -> FileSystems.matchSingleFileSpec(gcsPath))
+ .thenThrow(new IOException("GCS file not found"));
+ // Act & Assert
+ RuntimeException exception =
+ Assert.assertThrows(RuntimeException.class, () -> factory.apply(config));
+ Assert.assertTrue(exception.getMessage().contains("Failed trying to process value"));
+ Assert.assertTrue(exception.getCause() instanceof IOException);
+ Assert.assertTrue(exception.getCause().getMessage().contains("Failed to download file"));
+ }
+ }
+
+ @Test
+ public void testApplyHappyPathIgnoresNonExternalValues() {
+ // Arrange
+ Map config = new HashMap<>();
+ config.put("some.string", "/local/path/file.txt");
+ config.put("some.number", 42);
+ config.put("some.boolean", false);
+
+ // Act
+ Map processedConfig = (Map) factory.apply(config);
+
+ // Assert
+ Assert.assertEquals(config, processedConfig);
+ }
+
+ @Test
+ public void testApplyEdgeCaseMultipleExternalPathsInSingleValue() {
+ // Arrange
+ String gcsPath1 = "gs://bucket/keytab.keytab";
+ String gcsPath2 = "gs://bucket/trust.jks";
+ String originalValue =
+ "jaas_config keyTab=\"" + gcsPath1 + "\" trustStore=\"" + gcsPath2 + "\"";
+
+ String expectedLocalPath1 =
+ FileAwareFactoryFn.DIRECTORY_PREFIX + "/" + TEST_FACTORY_TYPE + "/bucket/keytab.keytab";
+ String expectedLocalPath2 =
+ FileAwareFactoryFn.DIRECTORY_PREFIX + "/" + TEST_FACTORY_TYPE + "/bucket/trust.jks";
+ String expectedProcessedValue =
+ "jaas_config keyTab=\""
+ + expectedLocalPath1
+ + "\" trustStore=\""
+ + expectedLocalPath2
+ + "\"";
+
+ Map config = new HashMap<>();
+ config.put("jaas.config", originalValue);
+
+ try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class)) {
+ // Mock GCS calls for both paths
+ mockSuccessfulDownload(mockedFileSystems, gcsPath1);
+ mockSuccessfulDownload(mockedFileSystems, gcsPath2);
+
+ // Act
+ Map processedConfig = (Map) factory.apply(config);
+
+ // Assert
+ Assert.assertEquals(expectedProcessedValue, processedConfig.get("jaas.config"));
+ }
+ }
+
+ @Test
+ public void testApplyEdgeCaseLocalFileWriteFails() throws IOException {
+ // Arrange
+ String gcsPath = "gs://test-bucket/some-file.txt";
+ Map config = new HashMap<>();
+ config.put("a.file", gcsPath);
+
+ // Mock GCS part to succeed
+ try (MockedStatic mockedFileSystems = Mockito.mockStatic(FileSystems.class);
+ MockedStatic mockedFileChannel = Mockito.mockStatic(FileChannel.class)) {
+ mockSuccessfulDownload(mockedFileSystems, gcsPath);
+
+ // Mock the local file writing part to fail
+ mockedFileChannel
+ .when(
+ () ->
+ FileChannel.open(
+ ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class)))
+ .thenThrow(new IOException("Permission denied"));
+
+ // Act & Assert
+ RuntimeException exception =
+ Assert.assertThrows(RuntimeException.class, () -> factory.apply(config));
+ Assert.assertTrue(exception.getMessage().contains("Failed trying to process value"));
+ Assert.assertTrue(exception.getCause() instanceof IOException);
+ // Check that the root cause is our "Permission denied" mock
+ Assert.assertTrue(exception.getCause().getCause().getMessage().contains("Permission denied"));
+ }
+ }
+
+ @Test
+ public void testApplyHappyPathResolvesSecretValue() {
+ // Arrange
+ String secretVersion = "secretValue:projects/p/secrets/s/versions/v";
+ String secretVersionParsed = "projects/p/secrets/s/versions/v";
+ String secretValue = "my-secret-password";
+ String originalValue = "password=" + secretVersion;
+ String expectedProcessedValue = "password=" + secretValue;
+
+ Map config = new HashMap<>();
+ config.put("db.password", originalValue);
+
+ TestFactoryFn factoryWithMockedSecret =
+ new TestFactoryFn() {
+ @Override
+ public byte[] getSecret(String secretIdentifier) {
+ // Assert that the correct identifier is passed
+ Assert.assertEquals(secretVersionParsed, secretIdentifier);
+ // Return a predictable, hardcoded value for the test
+ return secretValue.getBytes(StandardCharsets.UTF_8);
+ }
+ };
+
+ // Act
+ @SuppressWarnings("unchecked")
+ Map processedConfig =
+ (Map) factoryWithMockedSecret.apply(config);
+
+ // Assert
+ Assert.assertEquals(expectedProcessedValue, processedConfig.get("db.password"));
+ }
+
+ @Test
+ public void testApplyFailurePathThrowsExceptionForInvalidSecretFormat() {
+ // Arrange
+ String invalidSecret = "secretValue:not-a-valid-secret-path";
+ Map config = new HashMap<>();
+ config.put("db.password", "password=" + invalidSecret);
+
+ // Act & Assert
+ RuntimeException ex = Assert.assertThrows(RuntimeException.class, () -> factory.apply(config));
+ Assert.assertEquals(IllegalArgumentException.class, ex.getCause().getClass());
+ }
+
+ // Helper method to reduce boilerplate in mocking successful GCS downloads
+ private void mockSuccessfulDownload(MockedStatic mockedFileSystems, String gcsPath) {
+ MatchResult.Metadata metadata = Mockito.mock(MatchResult.Metadata.class);
+ ResourceId resourceId = Mockito.mock(ResourceId.class);
+ Mockito.when(metadata.resourceId()).thenReturn(resourceId);
+ mockedFileSystems
+ .when(() -> FileSystems.matchSingleFileSpec(ArgumentMatchers.eq(gcsPath)))
+ .thenReturn(metadata);
+
+ ReadableByteChannel channel = Channels.newChannel(new ByteArrayInputStream(new byte[0]));
+ mockedFileSystems
+ .when(() -> FileSystems.open(ArgumentMatchers.eq(resourceId)))
+ .thenReturn(channel);
+ }
+}
diff --git a/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java
new file mode 100644
index 000000000000..503b2f8f10c0
--- /dev/null
+++ b/sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java
@@ -0,0 +1,242 @@
+/*
+ * 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.kafka.factories;
+
+import static org.mockito.Mockito.spy;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Stream;
+import javax.security.auth.login.Configuration;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatchers;
+import org.mockito.MockedConstruction;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class KerberosConsumerFactoryFnTest {
+
+ private KerberosConsumerFactoryFn factory;
+ private String originalKrb5Conf;
+ private static final String KRB5_GCS_PATH = "gs://sec-bucket/kerberos/krb5.conf";
+ private static final String KRB5_S3_PATH = "s3://sec-bucket/kerberos/krb5.conf";
+ private static final String LOCAL_FACTORY_TYPE = "kerberos";
+
+ @Before
+ public void setup() {
+ try {
+ java.lang.reflect.Field field =
+ KerberosConsumerFactoryFn.class.getDeclaredField("localKrb5ConfPath");
+ field.setAccessible(true);
+ field.set(null, "");
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ originalKrb5Conf = System.getProperty("java.security.krb5.conf");
+ }
+
+ @After
+ public void tearDown() throws IOException {
+ // Clean up system property to avoid affecting other tests
+ if (originalKrb5Conf != null) {
+ System.setProperty("java.security.krb5.conf", originalKrb5Conf);
+ } else {
+ System.clearProperty("java.security.krb5.conf");
+ }
+
+ // Clean up the directory created outside of the JUnit TemporaryFolder rule.
+ Path pathToDelete = Paths.get(FileAwareFactoryFn.DIRECTORY_PREFIX, LOCAL_FACTORY_TYPE);
+ if (Files.exists(pathToDelete)) {
+ try (Stream walk = Files.walk(pathToDelete)) {
+ walk.sorted(Comparator.reverseOrder()).map(Path::toFile).forEach(File::delete);
+ }
+ }
+ }
+
+ @Test
+ @SuppressWarnings("rawtypes")
+ public void testHappyGcsPath() {
+ String keytabGcsPath = "gs://sec-bucket/keytabs/my.keytab";
+ String expectedKrb5LocalPath = "/tmp/kerberos/sec-bucket/kerberos/krb5.conf";
+ String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab";
+
+ Map config = new HashMap<>();
+ config.put(
+ "sasl.jaas.config",
+ "com.sun.security.auth.module.Krb5LoginModule required keyTab=\""
+ + keytabGcsPath
+ + "\" principal=\"user@REALM\";");
+
+ factory = spy(new KerberosConsumerFactoryFn(KRB5_GCS_PATH));
+ // This mock prevents the spy from calling the real createObject method,
+ // which would otherwise crash.
+ Mockito.doReturn(null).when(factory).createObject(ArgumentMatchers.anyMap());
+
+ try (MockedStatic mockedStaticFactory =
+ Mockito.mockStatic(FileAwareFactoryFn.class, Mockito.CALLS_REAL_METHODS);
+ MockedStatic mockedConfiguration = Mockito.mockStatic(Configuration.class);
+ MockedStatic mockedFiles = Mockito.mockStatic(Files.class);
+ MockedConstruction mockedConsumer =
+ Mockito.mockConstruction(KafkaConsumer.class)) {
+
+ Assert.assertNotNull(mockedConsumer);
+ // Mock the static downloadExternalFile method to prevent any GCS interaction
+ mockedStaticFactory
+ .when(
+ () ->
+ FileAwareFactoryFn.downloadExternalFile(
+ ArgumentMatchers.eq(KRB5_GCS_PATH), ArgumentMatchers.anyString()))
+ .thenReturn(expectedKrb5LocalPath);
+ mockedStaticFactory
+ .when(
+ () ->
+ FileAwareFactoryFn.downloadExternalFile(
+ ArgumentMatchers.eq(keytabGcsPath), ArgumentMatchers.anyString()))
+ .thenReturn(expectedKeytabLocalPath);
+
+ Configuration mockConf = Mockito.mock(Configuration.class);
+ mockedConfiguration.when(Configuration::getConfiguration).thenReturn(mockConf);
+ mockedFiles
+ .when(
+ () ->
+ Files.setPosixFilePermissions(
+ ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class)))
+ .thenReturn(null);
+ mockedFiles
+ .when(() -> Files.createDirectories(ArgumentMatchers.any(Path.class)))
+ .thenReturn(null);
+
+ // Act
+ factory.apply(config);
+
+ // Assert
+ // 1. Verify that the krb5.conf system property was set correctly.
+ Assert.assertEquals(expectedKrb5LocalPath, System.getProperty("java.security.krb5.conf"));
+
+ // 2. Capture the config passed to createObject and verify the keytab path was replaced.
+ ArgumentCaptor> configCaptor = ArgumentCaptor.forClass(Map.class);
+ Mockito.verify(factory).createObject(configCaptor.capture());
+ Map capturedConfig = configCaptor.getValue();
+ String processedJaasConfig = (String) capturedConfig.get("sasl.jaas.config");
+ Assert.assertTrue(processedJaasConfig.contains("keyTab=\"" + expectedKeytabLocalPath + "\""));
+
+ // 3. Verify that the JAAS configuration was refreshed.
+ Mockito.verify(mockConf).refresh();
+ }
+ }
+
+ @Test
+ @SuppressWarnings("rawtypes")
+ public void testHappyS3Path() {
+ String keytabPath = "s3://sec-bucket/keytabs/my.keytab";
+ String expectedKrb5LocalPath = "/tmp/kerberos/sec-bucket/kerberos/krb5.conf";
+ String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab";
+
+ Map config = new HashMap<>();
+ config.put(
+ "sasl.jaas.config",
+ "com.sun.security.auth.module.Krb5LoginModule required keyTab=\""
+ + keytabPath
+ + "\" principal=\"user@REALM\";");
+ factory = spy(new KerberosConsumerFactoryFn(KRB5_S3_PATH));
+ // This mock prevents the spy from calling the real createObject method,
+ // which would otherwise crash.
+ Mockito.doReturn(null).when(factory).createObject(ArgumentMatchers.anyMap());
+
+ try (MockedStatic mockedStaticFactory =
+ Mockito.mockStatic(FileAwareFactoryFn.class, Mockito.CALLS_REAL_METHODS);
+ MockedStatic mockedConfiguration = Mockito.mockStatic(Configuration.class);
+ MockedStatic mockedFiles = Mockito.mockStatic(Files.class);
+ MockedConstruction mockedConsumer =
+ Mockito.mockConstruction(KafkaConsumer.class)) {
+
+ Assert.assertNotNull(mockedConsumer);
+ // Mock the static downloadExternalFile method to prevent any interaction
+ mockedStaticFactory
+ .when(
+ () ->
+ FileAwareFactoryFn.downloadExternalFile(
+ ArgumentMatchers.eq(KRB5_S3_PATH), ArgumentMatchers.anyString()))
+ .thenReturn(expectedKrb5LocalPath);
+ mockedStaticFactory
+ .when(
+ () ->
+ FileAwareFactoryFn.downloadExternalFile(
+ ArgumentMatchers.eq(keytabPath), ArgumentMatchers.anyString()))
+ .thenReturn(expectedKeytabLocalPath);
+
+ Configuration mockConf = Mockito.mock(Configuration.class);
+ mockedConfiguration.when(Configuration::getConfiguration).thenReturn(mockConf);
+ mockedFiles
+ .when(
+ () ->
+ Files.setPosixFilePermissions(
+ ArgumentMatchers.any(Path.class), ArgumentMatchers.any(Set.class)))
+ .thenReturn(null);
+ mockedFiles
+ .when(() -> Files.createDirectories(ArgumentMatchers.any(Path.class)))
+ .thenReturn(null);
+
+ // Act
+ factory.apply(config);
+
+ // Assert
+ // 1. Verify that the krb5.conf system property was set correctly.
+ Assert.assertEquals(expectedKrb5LocalPath, System.getProperty("java.security.krb5.conf"));
+
+ // 2. Capture the config passed to createObject and verify the keytab path was replaced.
+ ArgumentCaptor> configCaptor = ArgumentCaptor.forClass(Map.class);
+ Mockito.verify(factory).createObject(configCaptor.capture());
+ Map capturedConfig = configCaptor.getValue();
+ String processedJaasConfig = (String) capturedConfig.get("sasl.jaas.config");
+ Assert.assertTrue(processedJaasConfig.contains("keyTab=\"" + expectedKeytabLocalPath + "\""));
+
+ // 3. Verify that the JAAS configuration was refreshed.
+ Mockito.verify(mockConf).refresh();
+ }
+ }
+
+ @Test
+ public void testInvalidKrb5ConfPathThrowsException() {
+ // Arrange
+ String invalidPath = "not-a-gcs-path"; // This path is missing the "gs://" prefix
+ factory = new KerberosConsumerFactoryFn(invalidPath);
+ Map config = new HashMap<>();
+
+ // Act & Assert
+ RuntimeException ex = Assert.assertThrows(RuntimeException.class, () -> factory.apply(config));
+
+ Assert.assertTrue(ex.getMessage().contains("Failed trying to process extra files"));
+ Assert.assertTrue(ex.getCause() instanceof IOException);
+ }
+}
diff --git a/settings.gradle.kts b/settings.gradle.kts
index c867e7ae2314..d887c8ad9d41 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -186,6 +186,7 @@ include(":sdks:java:extensions:kryo")
include(":sdks:java:extensions:google-cloud-platform-core")
include(":sdks:java:extensions:jackson")
include(":sdks:java:extensions:join-library")
+include(":sdks:java:extensions:kafka-factories")
include(":sdks:java:extensions:ml")
include(":sdks:java:extensions:ordered")
include(":sdks:java:extensions:protobuf")