generalize the classes such that no external storage system is strictly specified.

Author: fozzie15

sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFn.java

@@ -45,11 +45,11 @@
 public abstract class FileAwareFactoryFn<T>
     implements SerializableFunction<Map<String, Object>, T> {
 
-  public static final String GCS_PATH_PREFIX = "gs://";
+  public static String externalBucketPrefix = "";
   public static final String SECRET_VALUE_PREFIX = "secretValue:";
   public static final String DIRECTORY_PREFIX = "/tmp";
   private static final Pattern PATH_PATTERN =
-      Pattern.compile("(gs://[^\"]+)|(secretValue:[^\"]+)|(secretFile:[^\"]+)");
+      Pattern.compile("([a-zA-Z0-9]+://[^\"]+)|(secretValue:[^\"]+)|(secretFile:[^\"]+)");
 
   private static final Map<String, byte[]> secretCache = new ConcurrentHashMap<>();
 
@@ -86,18 +86,18 @@ public T apply(Map<String, Object> config) {
             StringBuffer sb = new StringBuffer();
 
             while (matcher.find()) {
-              String gcsPath = matcher.group(1);
+              String externalPath = matcher.group(1);
               String secretValue = matcher.group(2);
               String secretFile = matcher.group(3);
 
-              if (gcsPath != null) {
+              if (externalPath != null) {
                 try {
-                  String tmpPath = replacePathWithLocal(gcsPath);
-                  String localPath = downloadGcsFile(gcsPath, tmpPath);
+                  String tmpPath = replacePathWithLocal(externalPath);
+                  String localPath = downloadExternalFile(externalPath, tmpPath);
                   matcher.appendReplacement(sb, Matcher.quoteReplacement(localPath));
-                  LOG.info("Downloaded {} to {}", gcsPath, localPath);
+                  LOG.info("Downloaded {} to {}", externalPath, localPath);
                 } catch (IOException io) {
-                  throw new IOException("Failed to download file : " + gcsPath, io);
+                  throw new IOException("Failed to download file : " + externalPath, io);
                 }
               } else if (secretValue != null) {
                 try {
@@ -131,16 +131,16 @@ public T apply(Map<String, Object> config) {
   }
 
   /**
-   * A function to download files from their specified gcs path and copy them to the provided local
-   * filepath. The local filepath is provided by the replacePathWithLocal.
+   * 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 gcsFilePath
+   * @param externalFilePath
    * @param outputFileString
    * @return
    * @throws IOException
    */
-  protected static synchronized String downloadGcsFile(String gcsFilePath, String outputFileString)
-      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;
@@ -150,14 +150,15 @@ protected static synchronized String downloadGcsFile(String gcsFilePath, String
     if (parentDir != null) {
       Files.createDirectories(parentDir);
     }
-    LOG.info("Staging GCS file [{}] to [{}]", gcsFilePath, outputFileString);
+    LOG.info("Staging external file [{}] to [{}]", externalFilePath, outputFileString);
     Set<StandardOpenOption> options = new HashSet<>(2);
     options.add(StandardOpenOption.CREATE);
     options.add(StandardOpenOption.WRITE);
 
-    // Copy the GCS file into a local file and will throw an I/O exception in case file not found.
+    // 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(gcsFilePath).resourceId())) {
+        FileSystems.open(FileSystems.matchSingleFileSpec(externalFilePath).resourceId())) {
       try (FileChannel writeChannel = FileChannel.open(outputFilePath, options)) {
         writeChannel.transferFrom(readerChannel, 0, Long.MAX_VALUE);
       }
@@ -170,16 +171,25 @@ protected byte[] getSecretWithCache(String secretId) {
   }
 
   /**
-   * A helper method to create a new string with the gcs 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
+   * 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 gcsPath
-   * @return a string with all instances of GCS paths converted to the local paths where the files
-   *     sit.
+   * @param externalPath
+   * @return a string with all instances of external paths converted to the local paths where the
+   *     files sit.
    */
-  private String replacePathWithLocal(String gcsPath) throws IOException {
-    return DIRECTORY_PREFIX + "/" + factoryType + "/" + gcsPath.substring(GCS_PATH_PREFIX.length());
+  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);
   }
 
   /**

sdks/java/extensions/kafka-factories/src/main/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFn.java

@@ -38,7 +38,7 @@
 
 public class KerberosConsumerFactoryFn extends FileAwareFactoryFn<Consumer<byte[], byte[]>> {
   private static final String LOCAL_FACTORY_TYPE = "kerberos";
-  private String krb5ConfigGcsPath = "";
+  private String krb5ConfigPath = "";
   private static volatile String localKrb5ConfPath = "";
 
   private static final Object lock = new Object();
@@ -51,16 +51,16 @@ public class KerberosConsumerFactoryFn extends FileAwareFactoryFn<Consumer<byte[
 
   private static final Logger LOG = LoggerFactory.getLogger(KerberosConsumerFactoryFn.class);
 
-  public KerberosConsumerFactoryFn(String krb5ConfigGcsPath) {
+  public KerberosConsumerFactoryFn(String krb5ConfigPath) {
     super("kerberos");
-    this.krb5ConfigGcsPath = krb5ConfigGcsPath;
+    this.krb5ConfigPath = krb5ConfigPath;
   }
 
   @Override
   protected Consumer<byte[], byte[]> createObject(Map<String, Object> 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 GCS prefix in this case.
+    // 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 = "";
@@ -90,10 +90,10 @@ 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.krb5ConfigGcsPath != null && !this.krb5ConfigGcsPath.isEmpty()) {
+        if (this.krb5ConfigPath != null && !this.krb5ConfigPath.isEmpty()) {
           String localPath =
               super.getBaseDirectory() + "/" + LOCAL_FACTORY_TYPE + "/" + "krb5.conf";
-          localKrb5ConfPath = downloadGcsFile(this.krb5ConfigGcsPath, localPath);
+          localKrb5ConfPath = downloadExternalFile(this.krb5ConfigPath, localPath);
 
           System.setProperty("java.security.krb5.conf", localKrb5ConfPath);
           Configuration.getConfiguration().refresh();

sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/FileAwareFactoryFnTest.java

@@ -72,7 +72,7 @@ public void setup() throws IOException {
   }
 
   @Test
-  public void testHappyPathReplacesGcsPath() {
+  public void testHappyPathReplacesExternalPath() {
     // Arrange
     String gcsPath = "gs://test-bucket/config-file.json";
     String expectedLocalPath =
@@ -118,7 +118,6 @@ public void testApplyFailurePathThrowsRuntimeExceptionOnDownloadFailure() {
       mockedFileSystems
           .when(() -> FileSystems.matchSingleFileSpec(gcsPath))
           .thenThrow(new IOException("GCS file not found"));
-
       // Act & Assert
       RuntimeException exception =
           Assert.assertThrows(RuntimeException.class, () -> factory.apply(config));
@@ -129,7 +128,7 @@ public void testApplyFailurePathThrowsRuntimeExceptionOnDownloadFailure() {
   }
 
   @Test
-  public void testApplyHappyPathIgnoresNonGcsValues() {
+  public void testApplyHappyPathIgnoresNonExternalValues() {
     // Arrange
     Map<String, Object> config = new HashMap<>();
     config.put("some.string", "/local/path/file.txt");
@@ -144,7 +143,7 @@ public void testApplyHappyPathIgnoresNonGcsValues() {
   }
 
   @Test
-  public void testApplyEdgeCaseMultipleGcsPathsInSingleValue() {
+  public void testApplyEdgeCaseMultipleExternalPathsInSingleValue() {
     // Arrange
     String gcsPath1 = "gs://bucket/keytab.keytab";
     String gcsPath2 = "gs://bucket/trust.jks";

sdks/java/extensions/kafka-factories/src/test/java/org/apache/beam/sdk/extensions/kafka/factories/KerberosConsumerFactoryFnTest.java

@@ -49,6 +49,7 @@ 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
@@ -61,8 +62,6 @@ public void setup() {
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
-
-    factory = spy(new KerberosConsumerFactoryFn(KRB5_GCS_PATH));
     originalKrb5Conf = System.getProperty("java.security.krb5.conf");
   }
 
@@ -86,8 +85,7 @@ public void tearDown() throws IOException {
 
   @Test
   @SuppressWarnings("rawtypes")
-  public void testHappyPath() {
-    // Arrange
+  public void testHappyGcsPath() {
     String keytabGcsPath = "gs://sec-bucket/keytabs/my.keytab";
     String expectedKrb5LocalPath = "/tmp/kerberos/krb5.conf";
     String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab";
@@ -99,6 +97,69 @@ public void testHappyPath() {
             + keytabGcsPath
             + "\" principal=\"user@REALM\";");
 
+    factory = spy(new KerberosConsumerFactoryFn(KRB5_GCS_PATH));
+    try (MockedStatic<FileAwareFactoryFn> mockedStaticFactory =
+            Mockito.mockStatic(FileAwareFactoryFn.class);
+        MockedStatic<Configuration> mockedConfiguration = Mockito.mockStatic(Configuration.class);
+        MockedStatic<Files> mockedFiles = Mockito.mockStatic(Files.class);
+        MockedConstruction<KafkaConsumer> mockedConsumer =
+            Mockito.mockConstruction(KafkaConsumer.class)) {
+
+      Assert.assertNotNull(mockedConsumer);
+      // Mock the static downloadExternalFile method to prevent any GCS interaction
+      mockedStaticFactory
+          .when(() -> FileAwareFactoryFn.downloadExternalFile(KRB5_GCS_PATH, expectedKrb5LocalPath))
+          .thenReturn(expectedKrb5LocalPath);
+      mockedStaticFactory
+          .when(
+              () -> FileAwareFactoryFn.downloadExternalFile(keytabGcsPath, expectedKeytabLocalPath))
+          .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<Map<String, Object>> configCaptor = ArgumentCaptor.forClass(Map.class);
+      Mockito.verify(factory).createObject(configCaptor.capture());
+      Map<String, Object> 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/krb5.conf";
+    String expectedKeytabLocalPath = "/tmp/kerberos/sec-bucket/keytabs/my.keytab";
+
+    Map<String, Object> 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));
     try (MockedStatic<FileAwareFactoryFn> mockedStaticFactory =
             Mockito.mockStatic(FileAwareFactoryFn.class);
         MockedStatic<Configuration> mockedConfiguration = Mockito.mockStatic(Configuration.class);
@@ -107,12 +168,12 @@ public void testHappyPath() {
             Mockito.mockConstruction(KafkaConsumer.class)) {
 
       Assert.assertNotNull(mockedConsumer);
-      // Mock the static downloadGcsFile method to prevent any GCS interaction
+      // Mock the static downloadExternalFile method to prevent any GCS interaction
       mockedStaticFactory
-          .when(() -> FileAwareFactoryFn.downloadGcsFile(KRB5_GCS_PATH, expectedKrb5LocalPath))
+          .when(() -> FileAwareFactoryFn.downloadExternalFile(KRB5_S3_PATH, expectedKrb5LocalPath))
           .thenReturn(expectedKrb5LocalPath);
       mockedStaticFactory
-          .when(() -> FileAwareFactoryFn.downloadGcsFile(keytabGcsPath, expectedKeytabLocalPath))
+          .when(() -> FileAwareFactoryFn.downloadExternalFile(keytabPath, expectedKeytabLocalPath))
           .thenReturn(expectedKeytabLocalPath);
 
       Configuration mockConf = Mockito.mock(Configuration.class);