Skip to content

Commit 13398fe

Browse files
committed
Revert "Revert "PLT-372: Add Heracles client in Metastore for fetching users efficiently""
1 parent 6412e6d commit 13398fe

File tree

26 files changed

+550
-199
lines changed

26 files changed

+550
-199
lines changed

auth-agents-common/pom.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151

5252
<dependency>
5353
<groupId>org.apache.atlas</groupId>
54-
<artifactId>client-keycloak</artifactId>
54+
<artifactId>client-auth</artifactId>
5555
<version>${project.version}</version>
5656
</dependency>
5757

auth-agents-common/src/main/java/org/apache/atlas/plugin/util/KeycloakUserStore.java

+5-7
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@
3939
import java.util.concurrent.TimeUnit;
4040
import java.util.stream.Collectors;
4141

42-
import static org.apache.atlas.keycloak.client.AtlasKeycloakClient.getKeycloakClient;
42+
import static org.apache.atlas.auth.client.keycloak.AtlasKeycloakClient.getKeycloakClient;
43+
import static org.apache.atlas.auth.client.heracles.AtlasHeraclesClient.getHeraclesClient;
4344
import static org.apache.atlas.repository.Constants.*;
4445
import static org.apache.atlas.repository.util.AccessControlUtils.ARGO_SERVICE_USER_NAME;
4546
import static org.apache.atlas.repository.util.AccessControlUtils.BACKEND_SERVICE_USER_NAME;
@@ -277,10 +278,7 @@ public RangerUserStore loadUserStoreIfUpdated(long lastUpdatedTime) throws Atlas
277278
}
278279

279280
Map<String, Set<String>> userGroupMapping = new HashMap<>();
280-
281-
List<UserRepresentation> kUsers = getKeycloakClient().getAllUsers();
282-
LOG.info("Found {} keycloak users", kUsers.size());
283-
281+
List<UserRepresentation> kUsers = getHeraclesClient().getAllUsers();
284282
List<Callable<Object>> callables = new ArrayList<>();
285283
kUsers.forEach(x -> callables.add(new UserGroupsFetcher(x, userGroupMapping)));
286284

@@ -417,13 +415,13 @@ public RangerRole call() throws Exception {
417415
//get all users for Roles
418416
Thread usersFetcher = new Thread(() -> {
419417
int start = 0;
420-
int size = 500;
418+
int size = 100;
421419
boolean found = true;
422420
Set<UserRepresentation> ret = new HashSet<>();
423421

424422
do {
425423
try {
426-
Set<UserRepresentation> userRepresentations = getKeycloakClient().getRoleUserMembers(kRole.getName(), start, size);
424+
Set<UserRepresentation> userRepresentations = getHeraclesClient().getRoleUserMembers(kRole.getName(), start, size);
427425
if (CollectionUtils.isNotEmpty(userRepresentations)) {
428426
ret.addAll(userRepresentations);
429427
start += size;

client-keycloak/pom.xml client-auth/pom.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
</parent>
2828
<modelVersion>4.0.0</modelVersion>
2929

30-
<artifactId>client-keycloak</artifactId>
30+
<artifactId>client-auth</artifactId>
3131

3232
<properties>
3333
<maven.compiler.source>8</maven.compiler.source>

client-keycloak/src/main/java/org/apache/atlas/keycloak/client/AbstractKeycloakClient.java client-auth/src/main/java/org/apache/atlas/auth/client/auth/AbstractAuthClient.java

+30-23
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
1-
package org.apache.atlas.keycloak.client;
1+
package org.apache.atlas.auth.client.auth;
22

3+
import com.fasterxml.jackson.databind.DeserializationFeature;
34
import com.fasterxml.jackson.databind.ObjectMapper;
45
import io.micrometer.core.instrument.Timer;
6+
import org.apache.atlas.auth.client.config.AuthConfig;
7+
import org.apache.atlas.auth.client.heracles.RetrofitHeraclesClient;
8+
import org.apache.atlas.auth.client.keycloak.RetrofitKeycloakClient;
59
import okhttp3.*;
610
import okhttp3.logging.HttpLoggingInterceptor;
711
import org.apache.atlas.AtlasErrorCode;
812
import org.apache.atlas.exception.AtlasBaseException;
9-
import org.apache.atlas.keycloak.client.config.KeycloakConfig;
10-
import org.apache.atlas.keycloak.client.service.AtlasKeycloakAuthService;
1113
import org.apache.atlas.service.metrics.MetricUtils;
1214
import org.slf4j.Logger;
1315
import org.slf4j.LoggerFactory;
@@ -26,31 +28,32 @@
2628
import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;
2729
import static org.apache.atlas.AtlasErrorCode.RESOURCE_NOT_FOUND;
2830

29-
abstract class AbstractKeycloakClient {
31+
public class AbstractAuthClient {
3032

31-
private final static Logger LOG = LoggerFactory.getLogger(AbstractKeycloakClient.class);
33+
private final static Logger LOG = LoggerFactory.getLogger(AbstractAuthClient.class);
3234
private static final Map<Integer, AtlasErrorCode> ERROR_CODE_MAP = new HashMap<>();
3335

34-
private static final int DEFAULT_KEYCLOAK_RETRY = 3;
36+
private static final int DEFAULT_RETRY = 3;
3537
private static final String AUTHORIZATION = "Authorization";
3638
private static final String BEARER = "Bearer ";
3739
private static final int TIMEOUT_IN_SEC = 60;
3840
private static final String INTEGRATION = "integration";
3941
private static final String KEYCLOAK = "keycloak";
4042

41-
protected final KeycloakConfig keycloakConfig;
42-
protected final RetrofitKeycloakClient retrofit;
43+
protected final AuthConfig authConfig;
44+
protected final RetrofitKeycloakClient retrofitKeycloakClient;
45+
protected final RetrofitHeraclesClient retrofitHeraclesClient;
4346

44-
private final AtlasKeycloakAuthService authService;
47+
private final KeycloakAuthenticationService authService;
4548
private MetricUtils metricUtils = null;
4649

4750
static {
4851
ERROR_CODE_MAP.put(HTTP_NOT_FOUND, RESOURCE_NOT_FOUND);
4952
ERROR_CODE_MAP.put(HTTP_BAD_REQUEST, BAD_REQUEST);
5053
}
5154

52-
public AbstractKeycloakClient(KeycloakConfig keycloakConfig) {
53-
this.keycloakConfig = keycloakConfig;
55+
public AbstractAuthClient(AuthConfig authConfig) {
56+
this.authConfig = authConfig;
5457
this.metricUtils = new MetricUtils();
5558
HttpLoggingInterceptor httpInterceptor = new HttpLoggingInterceptor();
5659
httpInterceptor.setLevel(HttpLoggingInterceptor.Level.BODY);
@@ -64,11 +67,15 @@ public AbstractKeycloakClient(KeycloakConfig keycloakConfig) {
6467
.writeTimeout(TIMEOUT_IN_SEC, TimeUnit.SECONDS)
6568
.readTimeout(TIMEOUT_IN_SEC, TimeUnit.SECONDS)
6669
.build();
67-
this.retrofit = new Retrofit.Builder().client(okHttpClient)
68-
.baseUrl(this.keycloakConfig.getAuthServerUrl())
70+
this.retrofitKeycloakClient = new Retrofit.Builder().client(okHttpClient)
71+
.baseUrl(this.authConfig.getAuthServerUrl())
6972
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper())).build()
7073
.create(RetrofitKeycloakClient.class);
71-
authService = new AtlasKeycloakAuthService(keycloakConfig);
74+
this.retrofitHeraclesClient = new Retrofit.Builder().client(okHttpClient)
75+
.baseUrl(this.authConfig.getHeraclesApiServerUrl())
76+
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper().disable(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES))).build()
77+
.create(RetrofitHeraclesClient.class);
78+
authService = new KeycloakAuthenticationService(authConfig);
7279
}
7380

7481
/**
@@ -97,21 +104,20 @@ public Response intercept(@NonNull Chain chain) throws IOException {
97104
return chain.proceed(request);
98105
}
99106
};
100-
101107
/**
102108
* Called only during auth failures.
103109
*/
104110
Authenticator authInterceptor = new Authenticator() {
105111
@Override
106112
public Request authenticate(Route route, @NonNull Response response) {
107-
if (responseCount(response) > DEFAULT_KEYCLOAK_RETRY) {
108-
LOG.warn("Keycloak: Falling back, retried {} times", DEFAULT_KEYCLOAK_RETRY);
113+
if (responseCount(response) > DEFAULT_RETRY) {
114+
LOG.warn("Auth Client: Falling back, retried {} times", DEFAULT_RETRY);
109115
return null;
110116
}
111-
LOG.info("Keycloak: Current keycloak token status, Expired: {}", authService.isTokenExpired());
117+
LOG.info("Auth Client: Current keycloak token status, Expired: {}", authService.isTokenExpired());
112118
return response.request().newBuilder()
113-
.addHeader(AUTHORIZATION, BEARER + authService.getAuthToken())
114-
.build();
119+
.addHeader(AUTHORIZATION, BEARER + authService.getAuthToken())
120+
.build();
115121
}
116122

117123
private int responseCount(Response response) {
@@ -134,13 +140,14 @@ protected <T> retrofit2.Response<T> processResponse(retrofit2.Call<T> req) throw
134140
return response;
135141
}
136142
String errMsg = response.errorBody().string();
137-
LOG.error("Keycloak: Client request processing failed code {} message:{}, request: {} {}",
143+
LOG.error("Auth Client: Client request processing failed code {} message:{}, request: {} {}",
138144
response.code(), errMsg, req.request().method(), req.request().url());
139145
throw new AtlasBaseException(ERROR_CODE_MAP.getOrDefault(response.code(), BAD_REQUEST), errMsg);
140146
} catch (Exception e) {
141-
LOG.error("Keycloak: request failed, request: {} {}, Exception: {}", req.request().method(), req.request().url(), e);
142-
throw new AtlasBaseException(BAD_REQUEST, "Keycloak request failed");
147+
LOG.error("Auth Client: request failed, request: {} {}, Exception: {}", req.request().method(), req.request().url(), e);
148+
throw new AtlasBaseException(BAD_REQUEST, "Auth request failed");
143149
}
144150
}
145151

152+
146153
}

client-keycloak/src/main/java/org/apache/atlas/keycloak/client/service/AtlasKeycloakAuthService.java client-auth/src/main/java/org/apache/atlas/auth/client/auth/KeycloakAuthenticationService.java

+14-14
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
package org.apache.atlas.keycloak.client.service;
1+
package org.apache.atlas.auth.client.auth;
22

33
import com.fasterxml.jackson.databind.ObjectMapper;
44
import okhttp3.*;
55
import okhttp3.logging.HttpLoggingInterceptor;
6+
import org.apache.atlas.auth.client.config.AuthConfig;
67
import org.apache.atlas.exception.AtlasBaseException;
7-
import org.apache.atlas.keycloak.client.RetrofitKeycloakClient;
8-
import org.apache.atlas.keycloak.client.config.KeycloakConfig;
8+
import org.apache.atlas.auth.client.keycloak.RetrofitKeycloakClient;
99
import org.jetbrains.annotations.NotNull;
1010
import org.keycloak.representations.AccessTokenResponse;
1111
import org.slf4j.Logger;
@@ -19,9 +19,9 @@
1919

2020
import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;
2121

22-
public final class AtlasKeycloakAuthService {
22+
public final class KeycloakAuthenticationService {
2323

24-
public final static Logger LOG = LoggerFactory.getLogger(AtlasKeycloakAuthService.class);
24+
public final static Logger LOG = LoggerFactory.getLogger(KeycloakAuthenticationService.class);
2525

2626
private final static String GRANT_TYPE = "grant_type";
2727
private static final String CLIENT_ID = "client_id";
@@ -30,14 +30,14 @@ public final class AtlasKeycloakAuthService {
3030
private static final int TIMEOUT_IN_SECS = 60;
3131

3232
private final RetrofitKeycloakClient retrofit;
33-
private final KeycloakConfig keycloakConfig;
33+
private final AuthConfig authConfig;
3434
private AccessTokenResponse currentAccessToken;
3535
private long expirationTime = -1;
3636

37-
public AtlasKeycloakAuthService(KeycloakConfig keycloakConfig) {
38-
this.keycloakConfig = keycloakConfig;
37+
public KeycloakAuthenticationService(AuthConfig authConfig) {
38+
this.authConfig = authConfig;
3939
this.retrofit = new Retrofit.Builder().client(getOkHttpClient())
40-
.baseUrl(this.keycloakConfig.getAuthServerUrl())
40+
.baseUrl(this.authConfig.getAuthServerUrl())
4141
.addConverterFactory(JacksonConverterFactory.create(new ObjectMapper())).build()
4242
.create(RetrofitKeycloakClient.class);
4343
}
@@ -59,7 +59,7 @@ private OkHttpClient getOkHttpClient() {
5959
Interceptor responseLoggingInterceptor = chain -> {
6060
Request request = chain.request();
6161
okhttp3.Response response = chain.proceed(request);
62-
LOG.info("Keycloak: Auth Request for url {} Status: {}", request.url(), response.code());
62+
LOG.info("Auth Client: Auth Request for url {} Status: {}", request.url(), response.code());
6363
return response;
6464
};
6565

@@ -70,16 +70,16 @@ public String getAuthToken() {
7070
synchronized (this) {
7171
if (isTokenExpired()) {
7272
try {
73-
retrofit2.Response<AccessTokenResponse> resp = this.retrofit.grantToken(this.keycloakConfig.getRealmId(), getTokenRequest()).execute();
73+
retrofit2.Response<AccessTokenResponse> resp = this.retrofit.grantToken(this.authConfig.getRealmId(), getTokenRequest()).execute();
7474
if (resp.isSuccessful()) {
7575
currentAccessToken = resp.body();
7676
expirationTime = currentTime() + currentAccessToken.getExpiresIn() - EXPIRY_OFFSET_SEC;
77-
LOG.info("Keycloak: Auth token fetched with expiry:{} sec", expirationTime);
77+
LOG.info("Auth Client: Auth token fetched with expiry:{} sec", expirationTime);
7878
} else {
7979
throw new AtlasBaseException(BAD_REQUEST, resp.errorBody().string());
8080
}
8181
} catch (Exception e) {
82-
LOG.error("Keycloak: Error while fetching access token for keycloak client.", e);
82+
LOG.error("Auth Client: Error while fetching access token for keycloak client.", e);
8383
throw new RuntimeException(e);
8484
}
8585
}
@@ -97,7 +97,7 @@ public boolean isTokenExpired() {
9797
}
9898

9999
private RequestBody getTokenRequest() {
100-
return new FormBody.Builder().addEncoded(CLIENT_ID, this.keycloakConfig.getClientId()).addEncoded(CLIENT_SECRET, this.keycloakConfig.getClientSecret()).addEncoded(GRANT_TYPE, this.keycloakConfig.getGrantType()).build();
100+
return new FormBody.Builder().addEncoded(CLIENT_ID, this.authConfig.getClientId()).addEncoded(CLIENT_SECRET, this.authConfig.getClientSecret()).addEncoded(GRANT_TYPE, this.authConfig.getGrantType()).build();
101101
}
102102

103103
private long currentTime() {
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package org.apache.atlas.auth.client.config;
2+
3+
import org.apache.atlas.AtlasErrorCode;
4+
import org.apache.atlas.exception.AtlasBaseException;
5+
import org.apache.commons.lang.StringUtils;
6+
import org.codehaus.jettison.json.JSONObject;
7+
import org.slf4j.Logger;
8+
import org.slf4j.LoggerFactory;
9+
10+
import java.io.File;
11+
import java.nio.charset.StandardCharsets;
12+
import java.nio.file.Files;
13+
import java.util.Optional;
14+
15+
import static org.apache.atlas.ApplicationProperties.ATLAS_CONFIGURATION_DIRECTORY_PROPERTY;
16+
17+
public class AuthConfig {
18+
private static final Logger LOG = LoggerFactory.getLogger(AuthConfig.class);
19+
20+
public String authServerUrl;
21+
public String realmId;
22+
public String clientId;
23+
public String clientSecret;
24+
public String grantType;
25+
public String heraclesApiServerUrl;
26+
27+
private static final String KEYCLOAK_PROPERTIES = "keycloak.json";
28+
private static final String DEFAULT_GRANT_TYPE = "client_credentials";
29+
private static final String KEY_REALM_ID = "realm";
30+
private static final String KEY_AUTH_SERVER_URL = "auth-server-url";
31+
private static final String KEY_CLIENT_ID = "resource";
32+
private static final String KEY_CREDENTIALS = "credentials";
33+
private static final String KEY_SECRET = "secret";
34+
35+
public String getAuthServerUrl() {
36+
return authServerUrl;
37+
}
38+
39+
public String getRealmId() {
40+
return realmId;
41+
}
42+
43+
public String getClientId() {
44+
return clientId;
45+
}
46+
47+
public String getClientSecret() {
48+
return clientSecret;
49+
}
50+
51+
public String getGrantType() {
52+
return grantType;
53+
}
54+
55+
public String getHeraclesApiServerUrl() {
56+
return heraclesApiServerUrl;
57+
}
58+
59+
public static AuthConfig getConfig() throws AtlasBaseException {
60+
String confLocation = System.getProperty(ATLAS_CONFIGURATION_DIRECTORY_PROPERTY);
61+
Optional<File> confFile = getConfigurationFile(confLocation);
62+
63+
if (confFile.isPresent()) {
64+
try {
65+
JSONObject object = new JSONObject(readFileToString(confFile.get()));
66+
return buildAuthConfigFromJson(object);
67+
} catch (Exception e) {
68+
LOG.error("Error parsing Keycloak configuration: ", e);
69+
throw new AtlasBaseException(AtlasErrorCode.KEYCLOAK_INIT_FAILED, "Error parsing Keycloak configuration");
70+
}
71+
} else {
72+
throw new AtlasBaseException(AtlasErrorCode.KEYCLOAK_INIT_FAILED, "Keycloak configuration file not found in location " + confLocation);
73+
}
74+
}
75+
76+
private static Optional<File> getConfigurationFile(String confLocation) {
77+
if (StringUtils.isNotEmpty(confLocation)) {
78+
File confFile = new File(confLocation, KEYCLOAK_PROPERTIES);
79+
if (confFile.exists()) {
80+
return Optional.of(confFile);
81+
}
82+
}
83+
return Optional.empty();
84+
}
85+
86+
private static String readFileToString(File file) throws Exception {
87+
return new String(Files.readAllBytes(file.toPath()), StandardCharsets.UTF_8);
88+
}
89+
90+
private static AuthConfig buildAuthConfigFromJson(JSONObject object) throws Exception {
91+
String realmId = object.getString(KEY_REALM_ID);
92+
String authServerUrl = object.getString(KEY_AUTH_SERVER_URL) + "/";
93+
String clientId = object.getString(KEY_CLIENT_ID);
94+
String grantType = DEFAULT_GRANT_TYPE;
95+
String clientSecret = object.getJSONObject(KEY_CREDENTIALS).getString(KEY_SECRET);
96+
97+
LOG.info("Keycloak configuration: REALM_ID:{}, AUTH_SERVER_URL:{}", realmId, authServerUrl);
98+
return AuthConfigBuilder.builder().realId(realmId).authServerUrl(authServerUrl).clientId(clientId).grantType(grantType).clientSecret(clientSecret).build();
99+
}
100+
}

0 commit comments

Comments
 (0)