TEZ-4102: Let session credentials be merged before merging am launch credentials (László Bodor reviewed by Ashutosh Chauhan)
Signed-off-by: Laszlo Bodor <bodorlaszlo0202@gmail.com>
diff --git a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java b/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
index 2b21024..2aeabe4 100644
--- a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
+++ b/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
@@ -470,19 +470,8 @@
// Setup required Credentials for the AM launch. DAG specific credentials
// are handled separately.
ByteBuffer securityTokens = null;
- // Setup security tokens
- Credentials amLaunchCredentials = new Credentials();
- if (amConfig.getCredentials() != null) {
- amLaunchCredentials.addAll(amConfig.getCredentials());
- }
-
- // Add Staging dir creds to the list of session credentials.
- TokenCache.obtainTokensForFileSystems(sessionCreds, new Path[]{binaryConfPath}, conf);
-
- populateTokenCache(conf, sessionCreds);
-
- // Add session specific credentials to the AM credentials.
- amLaunchCredentials.mergeAll(sessionCreds);
+ Credentials amLaunchCredentials =
+ prepareAmLaunchCredentials(amConfig, sessionCreds, conf, binaryConfPath);
DataOutputBuffer dob = new DataOutputBuffer();
amLaunchCredentials.writeTokenStorageToStream(dob);
@@ -703,6 +692,25 @@
}
+ static Credentials prepareAmLaunchCredentials(AMConfiguration amConfig, Credentials sessionCreds,
+ TezConfiguration conf, Path binaryConfPath) throws IOException {
+ // Setup security tokens
+ Credentials amLaunchCredentials = new Credentials();
+
+ // Add Staging dir creds to the list of session credentials.
+ TokenCache.obtainTokensForFileSystems(sessionCreds, new Path[] {binaryConfPath }, conf);
+
+ populateTokenCache(conf, sessionCreds);
+
+ // Add session specific credentials to the AM credentials.
+ amLaunchCredentials.mergeAll(sessionCreds);
+
+ if (amConfig.getCredentials() != null) {
+ amLaunchCredentials.mergeAll(amConfig.getCredentials());
+ }
+ return amLaunchCredentials;
+ }
+
//get secret keys and tokens and store them into TokenCache
private static void populateTokenCache(TezConfiguration conf, Credentials credentials)
throws IOException{
diff --git a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java b/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
index 581d722..edcec49 100644
--- a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
+++ b/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
@@ -46,8 +46,10 @@
import org.apache.hadoop.fs.permission.FsPermission;
import org.apache.hadoop.hdfs.DistributedFileSystem;
import org.apache.hadoop.io.DataInputByteBuffer;
+import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.token.Token;
+import org.apache.hadoop.security.token.TokenIdentifier;
import org.apache.hadoop.yarn.api.ApplicationConstants;
import org.apache.hadoop.yarn.api.ApplicationConstants.Environment;
import org.apache.hadoop.yarn.api.records.ApplicationId;
@@ -890,5 +892,29 @@
}
+ @Test
+ public void testSessionCredentialsMergedBeforeAmConfigCredentials() throws Exception {
+ TezConfiguration conf = new TezConfiguration();
+ Text tokenType = new Text("TEST_TOKEN_TYPE");
+ Text tokenKind = new Text("TEST_TOKEN_KIND");
+ Text tokenService = new Text("TEST_TOKEN_SERVICE");
+ Credentials amConfigCredentials = new Credentials();
+ amConfigCredentials.addToken(tokenType,
+ new Token<>("id1".getBytes(), null, tokenKind, tokenService));
+
+ Credentials sessionCredentials = new Credentials();
+ Token<TokenIdentifier> sessionToken =
+ new Token<>("id2".getBytes(), null, tokenKind, tokenService);
+ sessionCredentials.addToken(tokenType, sessionToken);
+
+ AMConfiguration amConfig = new AMConfiguration(conf, null, amConfigCredentials);
+
+ Credentials amLaunchCredentials =
+ TezClientUtils.prepareAmLaunchCredentials(amConfig, sessionCredentials, conf, null);
+
+ // if there is another token in am conf creds of the same token type,
+ // session token should be applied while creating ContainerLaunchContext
+ Assert.assertEquals(sessionToken, amLaunchCredentials.getToken(tokenType));
+ }
}
