feat(kinesis) Cross-account Kinesis source support via STS assume-role in JsonKinesisSource

hudi-utilities/pom.xml

@@ -537,6 +537,12 @@
       <artifactId>kinesis</artifactId>
       <version>${aws.sdk.version}</version>
     </dependency>
+    <!-- STS: assume-role credentials for reading a Kinesis stream in a different AWS account. -->
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>sts</artifactId>
+      <version>${aws.sdk.version}</version>
+    </dependency>
     <!-- KPL de-aggregation: extracts user records from Kinesis Producer Library aggregated records -->
     <dependency>
       <groupId>com.amazonaws</groupId>

hudi-utilities/src/main/java/org/apache/hudi/utilities/config/KinesisSourceConfig.java

@@ -76,6 +76,17 @@ public class KinesisSourceConfig extends HoodieConfig {
       .withDocumentation("AWS secret key for Kinesis. Used when connecting to custom endpoints (e.g., LocalStack). "
           + "If not set with endpoint, uses the default AWS credential chain.");
 
+  public static final ConfigProperty<String> KINESIS_ROLE_ARN = ConfigProperty
+      .key(PREFIX + "role.arn")
+      .noDefaultValue()
+      .sinceVersion("1.2.0")
+      .markAdvanced()
+      .withDocumentation("IAM role ARN to assume (via STS) when the Kinesis stream lives in a different "
+          + "AWS account than the application. When set, the Kinesis client uses an auto-refreshing "
+          + "StsAssumeRoleCredentialsProvider whose base credentials come from the default credential chain. "
+          + "When empty/absent, the stream is read from the application's own AWS account using the default "
+          + "credential chain (legacy behavior). No external ID is used.");
+
   public static final ConfigProperty<Long> MAX_EVENTS_FROM_KINESIS_SOURCE = ConfigProperty
       .key(PREFIX + "max.events")
       .defaultValue(5000000L)

hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/JsonKinesisSource.java

@@ -117,6 +117,7 @@ protected JavaRDD<String> toBatch(KinesisOffsetGen.KinesisShardRange[] shardRang
         offsetGen.getEndpointUrl().orElse(null),
         getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_ACCESS_KEY, null),
         getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_SECRET_KEY, null),
+        getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_ROLE_ARN, null),
         offsetGen.getStartingPositionStrategy(),
         shouldAddMetaFields,
         getBooleanWithAltKeys(props, KinesisSourceConfig.KINESIS_ENABLE_DEAGGREGATION),
@@ -134,7 +135,8 @@ protected JavaRDD<String> toBatch(KinesisOffsetGen.KinesisShardRange[] shardRang
           List<ShardFetchResult> results = new ArrayList<>();
           try (KinesisClient client = KinesisOffsetGen.createKinesisClient(
               readConfig.getRegion(), readConfig.getEndpointUrl(),
-              readConfig.getAccessKey(), readConfig.getSecretKey())) {
+              readConfig.getAccessKey(), readConfig.getSecretKey(),
+              readConfig.getRoleArn())) {
             while (shardRangeIt.hasNext()) {
               KinesisOffsetGen.KinesisShardRange range = shardRangeIt.next();
               // Lazy iterator: fetches one GetRecords page at a time, keeping only one page in

hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/helpers/KinesisOffsetGen.java

@@ -41,6 +41,9 @@
 import software.amazon.awssdk.services.kinesis.model.ProvisionedThroughputExceededException;
 import software.amazon.awssdk.services.kinesis.model.ResourceNotFoundException;
 import software.amazon.awssdk.services.kinesis.model.Shard;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
 
 import java.math.BigInteger;
 import java.net.URI;
@@ -50,6 +53,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -304,28 +308,58 @@ public KinesisOffsetGen(TypedProperties props) {
         getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_STARTING_POSITION, true));
   }
 
+  /**
+   * Assume-role credentials providers cached by {@code region|roleArn}. Kinesis clients are built
+   * per {@code mapPartitions} task, so creating a fresh STS client each time would leak one (with its
+   * own HTTP connection pool) per micro-batch on a long-lived streaming executor — AWS SDK v2 does not
+   * close a user-supplied credentials provider (nor its injected StsClient) when the KinesisClient
+   * closes. Caching means at most one provider/StsClient per distinct role per executor JVM, reused
+   * across partitions and micro-batches. Providers auto-refresh and live for the JVM lifetime, so they
+   * are intentionally never closed.
+   */
+  private static final Map<String, StsAssumeRoleCredentialsProvider> ASSUME_ROLE_PROVIDERS =
+      new ConcurrentHashMap<>();
+
+  private static StsAssumeRoleCredentialsProvider assumeRoleProvider(String region, String roleArn) {
+    return ASSUME_ROLE_PROVIDERS.computeIfAbsent(region + "|" + roleArn, ignored ->
+        StsAssumeRoleCredentialsProvider.builder()
+            .stsClient(StsClient.builder().region(Region.of(region)).build())
+            .refreshRequest(AssumeRoleRequest.builder()
+                .roleArn(roleArn)
+                .roleSessionName("hudi-kinesis-source")
+                .build())
+            .build());
+  }
+
   /**
    * Builds a Kinesis client from explicit parameters. Used by both the instance method
    * {@link #createKinesisClient()} and by {@link org.apache.hudi.utilities.sources.JsonKinesisSource}
    * from serializable {@link KinesisReadConfig} in Spark closures.
    */
   public static KinesisClient createKinesisClient(String region, String endpointUrl,
-      String accessKey, String secretKey) {
+      String accessKey, String secretKey, String roleArn) {
     KinesisClientBuilder builder = KinesisClient.builder().region(Region.of(region));
     if (endpointUrl != null && !endpointUrl.isEmpty()) {
       builder = builder.endpointOverride(URI.create(endpointUrl));
     }
     if (accessKey != null && !accessKey.isEmpty() && secretKey != null && !secretKey.isEmpty()) {
+      // Static credentials (e.g. LocalStack / custom endpoint) take precedence.
       builder = builder.credentialsProvider(
           StaticCredentialsProvider.create(AwsBasicCredentials.create(accessKey, secretKey)));
+    } else if (roleArn != null && !roleArn.isEmpty()) {
+      // Cross-account stream: assume the customer-provided role via STS, using a per-executor cached
+      // provider (see ASSUME_ROLE_PROVIDERS) so we don't leak an StsClient per micro-batch. The base
+      // STS client uses the default credential chain, which must be granted sts:AssumeRole on this ARN.
+      builder = builder.credentialsProvider(assumeRoleProvider(region, roleArn));
     }
     return builder.build();
   }
 
   public KinesisClient createKinesisClient() {
     String accessKey = getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_ACCESS_KEY, null);
     String secretKey = getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_SECRET_KEY, null);
-    return createKinesisClient(region, endpointUrl.orElse(null), accessKey, secretKey);
+    String roleArn = getStringWithAltKeys(props, KinesisSourceConfig.KINESIS_ROLE_ARN, null);
+    return createKinesisClient(region, endpointUrl.orElse(null), accessKey, secretKey, roleArn);
   }
 
   /**

hudi-utilities/src/main/java/org/apache/hudi/utilities/sources/helpers/KinesisReadConfig.java

@@ -40,6 +40,7 @@ public class KinesisReadConfig implements Serializable {
   private final String endpointUrl; // null if not set
   private final String accessKey; // null if not set
   private final String secretKey; // null if not set
+  private final String roleArn; // null if not set; cross-account stream reads via STS assume-role
   private final KinesisSourceConfig.KinesisStartingPositionStrategy startingPosition;
   private final boolean metaFieldsEnabled;
   private final boolean deaggregationEnabled;

hudi-utilities/src/test/java/org/apache/hudi/utilities/sources/helpers/TestKinesisOffsetGenClient.java

@@ -0,0 +1,73 @@
+/*
+ * 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 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.hudi.utilities.sources.helpers;
+
+import org.apache.hudi.common.config.TypedProperties;
+import org.apache.hudi.utilities.config.KinesisSourceConfig;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.kinesis.KinesisClient;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+/**
+ * Covers the credential-provider branches of {@link KinesisOffsetGen#createKinesisClient}. AWS SDK v2
+ * clients resolve credentials lazily (only on the first request), so a client can be built for each
+ * branch without any AWS environment or network access.
+ */
+class TestKinesisOffsetGenClient {
+
+  private static final String REGION = "us-west-2";
+  private static final String ROLE_ARN = "arn:aws:iam::123456789012:role/kinesis-cross-account-read";
+
+  @Test
+  void buildsClientWithAssumeRoleProviderWhenArnPresent() {
+    try (KinesisClient client = KinesisOffsetGen.createKinesisClient(REGION, null, null, null, ROLE_ARN)) {
+      assertNotNull(client);
+    }
+  }
+
+  @Test
+  void buildsClientWithDefaultChainWhenNoArnOrKeys() {
+    try (KinesisClient client = KinesisOffsetGen.createKinesisClient(REGION, null, null, null, null)) {
+      assertNotNull(client);
+    }
+  }
+
+  @Test
+  void staticKeysTakePrecedenceOverAssumeRole() {
+    // Both static keys and an ARN set: the static-credentials branch wins (no STS assume-role).
+    try (KinesisClient client =
+        KinesisOffsetGen.createKinesisClient(REGION, null, "access", "secret", ROLE_ARN)) {
+      assertNotNull(client);
+    }
+  }
+
+  @Test
+  void instanceClientReadsRoleArnFromProps() {
+    TypedProperties props = new TypedProperties();
+    props.setProperty(KinesisSourceConfig.KINESIS_STREAM_NAME.key(), "test-stream");
+    props.setProperty(KinesisSourceConfig.KINESIS_REGION.key(), REGION);
+    props.setProperty(KinesisSourceConfig.KINESIS_ROLE_ARN.key(), ROLE_ARN);
+
+    try (KinesisClient client = new KinesisOffsetGen(props).createKinesisClient()) {
+      assertNotNull(client);
+    }
+  }
+}

hudi-utilities/src/test/java/org/apache/hudi/utilities/sources/helpers/TestKinesisReadConfig.java

@@ -0,0 +1,82 @@
+/*
+ * 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 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.hudi.utilities.sources.helpers;
+
+import org.apache.hudi.utilities.config.KinesisSourceConfig;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+/**
+ * Config-plumbing tests for the cross-account Kinesis role ARN: the config-key contract and the ARN
+ * surviving transport to Spark executors via the serializable {@link KinesisReadConfig}.
+ */
+class TestKinesisReadConfig {
+
+  private static final String ROLE_ARN = "arn:aws:iam::123456789012:role/kinesis-cross-account-read";
+
+  @Test
+  void roleArnKeyMatchesExpectedContract() {
+    assertEquals("hoodie.streamer.source.kinesis.role.arn",
+        KinesisSourceConfig.KINESIS_ROLE_ARN.key());
+    // No default: an absent key must read as null so the client falls back to the default chain.
+    assertFalse(KinesisSourceConfig.KINESIS_ROLE_ARN.hasDefaultValue());
+  }
+
+  @Test
+  void readConfigRoundTripsRoleArn() {
+    assertEquals(ROLE_ARN, newReadConfig(ROLE_ARN).getRoleArn());
+    // Legacy same-account path: null ARN is preserved (not coerced to empty).
+    assertNull(newReadConfig(null).getRoleArn());
+  }
+
+  /**
+   * The client is rebuilt on executors from a deserialized KinesisReadConfig, so the ARN must
+   * survive Java serialization for cross-account reads to work outside the driver.
+   */
+  @Test
+  void roleArnSurvivesSerialization() throws Exception {
+    KinesisReadConfig deserialized = serializeRoundTrip(newReadConfig(ROLE_ARN));
+    assertEquals(ROLE_ARN, deserialized.getRoleArn());
+  }
+
+  private static KinesisReadConfig newReadConfig(String roleArn) {
+    return new KinesisReadConfig("test-stream", "us-west-2", null, null, null, roleArn,
+        KinesisSourceConfig.KinesisStartingPositionStrategy.LATEST, false, true,
+        10000, 200L, 5000L, 1000L, 10000L, 600000L);
+  }
+
+  private static KinesisReadConfig serializeRoundTrip(KinesisReadConfig config) throws Exception {
+    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+    try (ObjectOutputStream out = new ObjectOutputStream(bytes)) {
+      out.writeObject(config);
+    }
+    try (ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
+      return (KinesisReadConfig) in.readObject();
+    }
+  }
+}