You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@atlas.apache.org by ni...@apache.org on 2026/09/10 09:00:13 UTC
(atlas) branch master updated: ATLAS-5401: Atlas client configuration to support custom token supplier (#752)
This is an automated email from the ASF dual-hosted git repository.
nixonrodrigues pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/atlas.git
The following commit(s) were added to refs/heads/master by this push:
new fb78b7dc4 ATLAS-5401: Atlas client configuration to support custom token supplier (#752)
fb78b7dc4 is described below
commit fb78b7dc41bc74e20ee6b477df39b8bafeb84ffb
Author: Madhan Neethiraj <ma...@apache.org>
AuthorDate: Thu Sep 10 02:00:06 2026 -0700
ATLAS-5401: Atlas client configuration to support custom token supplier (#752)
* ATLAS-5401: Atlas client configuration to support custom token supplier
---
.../java/org/apache/atlas/AtlasClientV2Test.java | 108 +++++++++++++++++++++
.../java/org/apache/atlas/AtlasBaseClient.java | 86 ++++++++++++----
.../token/retriever/JwTokenRetrieverDefault.java | 27 ++++--
.../atlas/token/retriever/TokenRetriever.java | 24 -----
4 files changed, 196 insertions(+), 49 deletions(-)
diff --git a/client/client-v2/src/test/java/org/apache/atlas/AtlasClientV2Test.java b/client/client-v2/src/test/java/org/apache/atlas/AtlasClientV2Test.java
index c2d2e24d1..4a9707293 100644
--- a/client/client-v2/src/test/java/org/apache/atlas/AtlasClientV2Test.java
+++ b/client/client-v2/src/test/java/org/apache/atlas/AtlasClientV2Test.java
@@ -55,12 +55,14 @@ import org.apache.atlas.model.typedef.AtlasEnumDef;
import org.apache.atlas.model.typedef.AtlasRelationshipDef;
import org.apache.atlas.model.typedef.AtlasStructDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
+import org.apache.atlas.token.retriever.JwTokenRetrieverDefault;
import org.apache.atlas.utils.AtlasJson;
import org.apache.commons.configuration2.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.http.HttpStatus;
import org.mockito.ArgumentMatchers;
import org.mockito.Mock;
+import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
@@ -81,7 +83,10 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.function.Supplier;
+import static org.apache.atlas.AtlasBaseClient.PROP_REST_AUTH_TOKEN_SUPPLIER;
+import static org.apache.atlas.token.retriever.JwTokenRetrieverDefault.JWT_SOURCE;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
@@ -2723,4 +2728,107 @@ public class AtlasClientV2Test {
client.addClassifications("TestType", attributes, classifications);
assertTrue(true);
}
+
+ @Test
+ public void testDefaultTokenSupplier() throws Exception {
+ Configuration configuration = Mockito.mock(Configuration.class);
+
+ when(configuration.getString(JWT_SOURCE, "")).thenReturn("env");
+
+ AtlasClientV2 client = new AtlasClientV2(service, configuration);
+
+ assertTokenSupplierClass(client, JwTokenRetrieverDefault.class);
+ }
+
+ @Test
+ public void testCustomTokenSupplier() throws Exception {
+ Configuration configuration = Mockito.mock(Configuration.class);
+
+ when(configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER)).thenReturn(TestTokenSupplier.class.getName());
+
+ AtlasClientV2 client = new AtlasClientV2(service, configuration);
+
+ assertTokenSupplierClass(client, TestTokenSupplier.class);
+ }
+
+ @Test
+ public void testCustomTokenSupplierWithConfig() throws Exception {
+ Configuration configuration = Mockito.mock(Configuration.class);
+
+ when(configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER)).thenReturn(TestTokenSupplierWithConfig.class.getName());
+
+ AtlasClientV2 client = new AtlasClientV2(service, configuration);
+
+ assertTokenSupplierClass(client, TestTokenSupplierWithConfig.class);
+ }
+
+ @Test
+ public void testNonExistingTokenSupplier() {
+ Configuration configuration = Mockito.mock(Configuration.class);
+
+ when(configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER)).thenReturn("NonExistingTokenSupplier");
+
+ try {
+ AtlasClientV2 client = new AtlasClientV2(service, configuration);
+
+ fail("AtlasClientV2 instantiation should have failed. Token supplier class: NonExistingTokenSupplier");
+ } catch (IllegalArgumentException ignore) {
+ // ignored
+ }
+ }
+
+ @Test
+ public void testInvalidTokenSupplier() {
+ Configuration configuration = Mockito.mock(Configuration.class);
+
+ when(configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER)).thenReturn(TestInvalidTokenSupplier.class.getName());
+
+ try {
+ AtlasClientV2 client = new AtlasClientV2(service, configuration);
+
+ fail("AtlasClientV2 instantiation should have failed. Token supplier class: " + TestInvalidTokenSupplier.class.getName());
+ } catch (IllegalArgumentException ignore) {
+ // ignored
+ }
+ }
+
+ private void assertTokenSupplierClass(AtlasClientV2 client, Class clz) throws Exception {
+ Field fieldTokenSupplier = AtlasBaseClient.class.getDeclaredField("tokenSupplier");
+
+ fieldTokenSupplier.setAccessible(true);
+
+ Object tokenSupplier = fieldTokenSupplier.get(client);
+
+ assertNotNull(tokenSupplier);
+ assertEquals(tokenSupplier.getClass(), clz);
+ }
+
+ private static class TestTokenSupplier implements Supplier<String> {
+ public TestTokenSupplier() {
+ }
+
+ @Override
+ public String get() {
+ return "testToken";
+ }
+ }
+
+ private static class TestTokenSupplierWithConfig implements Supplier<String> {
+ public TestTokenSupplierWithConfig(Configuration ignored) {
+ }
+
+ @Override
+ public String get() {
+ return "testTokenWithConfig";
+ }
+ }
+
+ private static class TestInvalidTokenSupplier {
+ public TestInvalidTokenSupplier() {
+ }
+
+ public String get() {
+ return "testToken";
+ }
+ }
}
diff --git a/client/common/src/main/java/org/apache/atlas/AtlasBaseClient.java b/client/common/src/main/java/org/apache/atlas/AtlasBaseClient.java
index 9b62e82d9..111f7f046 100644
--- a/client/common/src/main/java/org/apache/atlas/AtlasBaseClient.java
+++ b/client/common/src/main/java/org/apache/atlas/AtlasBaseClient.java
@@ -44,7 +44,6 @@ import org.apache.atlas.model.impexp.AtlasServer;
import org.apache.atlas.model.metrics.AtlasMetrics;
import org.apache.atlas.security.SecureClientUtils;
import org.apache.atlas.token.retriever.JwTokenRetrieverDefault;
-import org.apache.atlas.token.retriever.TokenRetriever;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.utils.AtlasJson;
import org.apache.atlas.utils.AuthenticationUtil;
@@ -67,11 +66,12 @@ import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
+import java.lang.reflect.Constructor;
import java.net.ConnectException;
import java.net.URI;
import java.util.List;
import java.util.Map;
-import java.util.Optional;
+import java.util.function.Supplier;
import static org.apache.atlas.security.SecurityProperties.TLS_ENABLED;
import static org.apache.atlas.token.retriever.JwTokenRetrieverDefault.JWT_SOURCE;
@@ -97,6 +97,9 @@ public abstract class AtlasBaseClient {
public static final API API_VERSION = new API(BASE_URI + ADMIN_VERSION, HttpMethod.GET, Response.Status.OK);
public static final API API_METRICS = new API(BASE_URI + ADMIN_METRICS, HttpMethod.GET, Response.Status.OK);
+ public static final String PROP_REST_AUTH_TOKEN_SUPPLIER = "atlas.rest.auth.token.supplier";
+ public static final String DEFAULT_REST_AUTH_TOKEN_SUPPLIER = JwTokenRetrieverDefault.class.getName();
+
static final String JSON_MEDIA_TYPE = MediaType.APPLICATION_JSON + "; charset=UTF-8";
static final String UNKNOWN_STATUS = "Unknown status";
static final String ATLAS_CLIENT_HA_RETRIES_KEY = "atlas.client.ha.retries";
@@ -124,7 +127,7 @@ public abstract class AtlasBaseClient {
private boolean retryEnabled;
private Cookie cookie;
private boolean useJwtAuth;
- private TokenRetriever<String> tokenRetriever;
+ private Supplier<String> tokenSupplier;
private SecureClientUtils clientUtils;
@@ -167,6 +170,8 @@ public abstract class AtlasBaseClient {
protected AtlasBaseClient(WebResource service, Configuration configuration) {
this.service = service;
this.configuration = configuration;
+
+ initializeTokenSupplier(configuration);
}
@VisibleForTesting
@@ -550,14 +555,19 @@ public abstract class AtlasBaseClient {
protected abstract API formatPathParameters(API api, String... params);
+ void initializeTokenSupplier(Configuration configuration) {
+ useJwtAuth = isJwtConfigured(configuration);
+ tokenSupplier = useJwtAuth ? getTokenSupplier(configuration) : null;
+ }
+
void initializeState(String[] baseUrls, UserGroupInformation ugi, String doAsUser) {
initializeState(getClientProperties(), baseUrls, ugi, doAsUser);
}
void initializeState(Configuration configuration, String[] baseUrls, UserGroupInformation ugi, String doAsUser) {
+ initializeTokenSupplier(configuration);
+
this.configuration = configuration;
- useJwtAuth = isJwtSourceConfigured(configuration);
- tokenRetriever = useJwtAuth ? getJwtTokenRetriever(configuration) : null;
Client client = getClient(configuration, ugi, doAsUser);
@@ -573,33 +583,73 @@ public abstract class AtlasBaseClient {
service = client.resource(UriBuilder.fromUri(activeServiceUrl).build());
}
- private TokenRetriever<String> getJwtTokenRetriever(Configuration configuration) {
- return new JwTokenRetrieverDefault(configuration);
+ private boolean isJwtConfigured(Configuration configuration) {
+ boolean ret = false;
+
+ if (configuration != null) {
+ String tokenSupplier = configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER);
+
+ ret = StringUtils.isNotBlank(tokenSupplier);
+
+ if (!ret) {
+ String jwtSource = configuration.getString(JWT_SOURCE, "");
+
+ ret = StringUtils.isNotBlank(jwtSource);
+ }
+ }
+
+ return ret;
}
- private boolean isJwtSourceConfigured(Configuration configuration) {
- if (configuration == null) {
- return false;
+ private Supplier<String> getTokenSupplier(Configuration configuration) {
+ Supplier<String> ret = null;
+
+ if (configuration != null) {
+ String clzName = configuration.getString(PROP_REST_AUTH_TOKEN_SUPPLIER);
+
+ if (StringUtils.isBlank(clzName)) {
+ clzName = DEFAULT_REST_AUTH_TOKEN_SUPPLIER;
+ }
+
+ try {
+ Class<?> clz = Class.forName(clzName);
+
+ if (!Supplier.class.isAssignableFrom(clz)) {
+ throw new IllegalArgumentException(clzName + " does not implement Supplier");
+ }
+
+ try {
+ // look for a constructor taking Configuration as its argument
+ Constructor<?> c = clz.getDeclaredConstructor(Configuration.class);
+
+ ret = (Supplier<String>) c.newInstance(configuration);
+ } catch (NoSuchMethodException excp) {
+ // use the default constructor
+ ret = (Supplier<String>) clz.getDeclaredConstructor().newInstance();
+ }
+ } catch (ReflectiveOperationException excp) {
+ throw new IllegalArgumentException("Failed to instantiate token supplier: " + clzName, excp);
+ }
}
- String jwtSource = configuration.getString(JWT_SOURCE, "");
- return StringUtils.isNotBlank(jwtSource);
+ return ret;
}
private void handleJwt(com.sun.jersey.api.client.WebResource.Builder requestBuilder) {
if (!useJwtAuth) {
return;
}
- if (tokenRetriever == null) {
- LOG.warn("AtlasBaseClient.handleJwt(): tokenRetriever is null. Skipping JWT header injection.");
+ if (tokenSupplier == null) {
+ LOG.warn("AtlasBaseClient.handleJwt(): tokenSupplier is null. Skipping JWT header injection.");
return;
}
- Optional<String> jwtOptional = tokenRetriever.retrieve();
- if (jwtOptional.isPresent()) {
- requestBuilder.header(AUTHORIZATION_HEADER, JWT_AUTHZ_PREFIX + jwtOptional.get());
+ String jwtOptional = tokenSupplier.get();
+
+ if (StringUtils.isNotBlank(jwtOptional)) {
+ requestBuilder.header(AUTHORIZATION_HEADER, JWT_AUTHZ_PREFIX + jwtOptional);
} else {
- LOG.warn("AtlasBaseClient.handleJwt(): JWT token not available from configured retriever. Authorization header not set.");
+ LOG.warn("AtlasBaseClient.handleJwt(): JWT token not available from configured supplier. Authorization header not set.");
}
}
diff --git a/client/common/src/main/java/org/apache/atlas/token/retriever/JwTokenRetrieverDefault.java b/client/common/src/main/java/org/apache/atlas/token/retriever/JwTokenRetrieverDefault.java
index c9c858f7d..01a569ca3 100644
--- a/client/common/src/main/java/org/apache/atlas/token/retriever/JwTokenRetrieverDefault.java
+++ b/client/common/src/main/java/org/apache/atlas/token/retriever/JwTokenRetrieverDefault.java
@@ -28,8 +28,9 @@ import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.Optional;
+import java.util.function.Supplier;
-public class JwTokenRetrieverDefault implements TokenRetriever<String> {
+public class JwTokenRetrieverDefault implements Supplier<String> {
private static final Logger LOG = LoggerFactory.getLogger(JwTokenRetrieverDefault.class);
public static final String JWT_SOURCE = "atlas.jwt.source";
@@ -66,21 +67,33 @@ public class JwTokenRetrieverDefault implements TokenRetriever<String> {
}
@Override
- public synchronized Optional<String> retrieve() {
+ public synchronized String get() {
String source = StringUtils.lowerCase(jwtSource);
+
+ final Optional<String> ret;
+
switch (source) {
case SOURCE_ENV:
- return getJwtFromEnv();
+ ret = getJwtFromEnv();
+ break;
+
case SOURCE_FILE:
- return getJwtFromFile();
+ ret = getJwtFromFile();
+ break;
+
case SOURCE_CRED:
- return getJwtFromCredProvider();
+ ret = getJwtFromCredProvider();
+ break;
+
default:
if (StringUtils.isNotBlank(source)) {
- LOG.warn("JwTokenRetrieverDefault.retrieve(): unsupported source='{}'", source);
+ LOG.warn("JwTokenRetrieverDefault.get(): unsupported source='{}'", source);
}
- return Optional.empty();
+
+ ret = Optional.empty();
}
+
+ return ret.orElse(null);
}
private Optional<String> getJwtFromEnv() {
diff --git a/client/common/src/main/java/org/apache/atlas/token/retriever/TokenRetriever.java b/client/common/src/main/java/org/apache/atlas/token/retriever/TokenRetriever.java
deleted file mode 100644
index cdcd84a42..000000000
--- a/client/common/src/main/java/org/apache/atlas/token/retriever/TokenRetriever.java
+++ /dev/null
@@ -1,24 +0,0 @@
-/**
- * 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.atlas.token.retriever;
-
-import java.util.Optional;
-
-public interface TokenRetriever<T> {
- Optional<T> retrieve();
-}