Skip to content

Commit

Permalink
fix gccs
Browse files Browse the repository at this point in the history
  • Loading branch information
FANNG1 committed Dec 26, 2024
1 parent 422dd7f commit 0a485b4
Show file tree
Hide file tree
Showing 8 changed files with 39 additions and 31 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
public class ADLSTokenCredential implements Credential {

/** ADLS SAS token credential type. */
public static final String ADLS_SAS_TOKEN_CREDENTIAL_TYPE = "adls-token";
public static final String ADLS_SAS_TOKEN_CREDENTIAL_TYPE = "adls-sas-token";
/** ADLS base domain */
public static final String ADLS_DOMAIN = "dfs.core.windows.net";
/** ADLS storage account name */
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,10 @@

package org.apache.gravitino.credential;

import com.google.common.base.Preconditions;
import java.io.Closeable;
import java.io.IOException;
import java.util.Map;
import org.apache.gravitino.exceptions.NoSuchCredentialException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand All @@ -34,21 +34,21 @@ public class CatalogCredentialManager implements Closeable {

private static final Logger LOG = LoggerFactory.getLogger(CatalogCredentialManager.class);

private CredentialCache<CredentialCacheKey> cacheManager;
private CredentialCache<CredentialCacheKey> credentialCache;

private final String catalogName;
private final Map<String, CredentialProvider> credentialProviders;

public CatalogCredentialManager(String catalogName, Map<String, String> catalogProperties) {
this.catalogName = catalogName;
this.credentialProviders = CredentialUtils.loadCredentialProviders(catalogProperties);
this.cacheManager = new CredentialCache();
cacheManager.initialize(catalogProperties);
this.credentialCache = new CredentialCache();
credentialCache.initialize(catalogProperties);
}

public Credential getCredential(String credentialType, CredentialContext context) {
public Credential doGetCredential(String credentialType, CredentialContext context) {
CredentialCacheKey credentialCacheKey = new CredentialCacheKey(credentialType, context);
return cacheManager.getCredential(credentialCacheKey, cacheKey -> getCredential(cacheKey));
return credentialCache.getCredential(credentialCacheKey, cacheKey -> doGetCredential(cacheKey));
}

@Override
Expand All @@ -69,14 +69,13 @@ public void close() {
});
}

private Credential getCredential(CredentialCacheKey credentialCacheKey) {
private Credential doGetCredential(CredentialCacheKey credentialCacheKey) {
String credentialType = credentialCacheKey.getCredentialType();
CredentialContext context = credentialCacheKey.getCredentialContext();
LOG.info("try get credential, credential type: {}, context: {}", credentialType, context);
CredentialProvider credentialProvider = credentialProviders.get(credentialType);
if (credentialProvider == null) {
throw new NoSuchCredentialException("No such credential: %s", credentialType);
}
return credentialProvider.getCredential(context);
LOG.debug("try get credential, credential type: {}, context: {}", credentialType, context);
Preconditions.checkState(
credentialProviders.containsKey(credentialType),
String.format("Credential %s not found", credentialType));
return credentialProviders.get(credentialType).getCredential(context);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.Expiry;
import java.io.Closeable;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
Expand All @@ -31,7 +33,7 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CredentialCache<T> {
public class CredentialCache<T> implements Closeable {

private static final Logger LOG = LoggerFactory.getLogger(CredentialCache.class);

Expand Down Expand Up @@ -84,21 +86,29 @@ public long expireAfterRead(

public void initialize(Map<String, String> catalogProperties) {
CredentialConfig credentialConfig = new CredentialConfig(catalogProperties);
long cache_size = credentialConfig.get(CredentialConfig.CREDENTIAL_MAX_CACHE_SIZE);
long cache_time = credentialConfig.get(CredentialConfig.CREDENTIAL_MAX_CACHE_TIME);
long cache_size = credentialConfig.get(CredentialConfig.CREDENTIAL_CACHE_MAZ_SIZE);
long cache_time = credentialConfig.get(CredentialConfig.CREDENTIAL_CACHE_MAX_TIME_IN_SECS);

this.credentialCache =
Caffeine.newBuilder()
.expireAfter(new CredentialExpireTimeCalculator(cache_time))
.maximumSize(cache_size)
.removalListener(
(cacheKey, credential, c) -> {
LOG.info("credential expire {}.", cacheKey);
LOG.debug("credential expire {}.", cacheKey);
})
.build();
}

public Credential getCredential(T cacheKey, Function<T, Credential> credentialSupplier) {
return credentialCache.get(cacheKey, key -> credentialSupplier.apply(cacheKey));
}

@Override
public void close() throws IOException {
if (credentialCache != null) {
credentialCache.invalidateAll();
credentialCache = null;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@
@Getter
public class CredentialCacheKey {

private String credentialType;
private CredentialContext credentialContext;
private final String credentialType;
private final CredentialContext credentialContext;

public CredentialCacheKey(String credentialType, CredentialContext credentialContext) {
this.credentialType = credentialType;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@

/** Contains credential context information to get credential from a credential provider. */
public interface CredentialContext {

/**
* Providing the username.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ private List<Credential> getCredentials(
entry ->
baseCatalog
.catalogCredentialManager()
.getCredential(entry.getKey(), entry.getValue()))
.doGetCredential(entry.getKey(), entry.getValue()))
.filter(Objects::nonNull)
.collect(Collectors.toList());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@

public class CredentialConfig extends Config {

private static final long DEFAULT_MAX_CREDENTIAL_CACHE_SIZE = 10_000L;
private static final long DEFAULT_MAX_CREDENTIAL_CACHE_TIME = 300;
private static final long DEFAULT_CREDENTIAL_CACHE_MAX_SIZE = 10_000L;
private static final long DEFAULT_CREDENTIAL_CACHE_MAX_TIME = 300;

public static final Map<String, PropertyEntry<?>> CREDENTIAL_PROPERTY_ENTRIES =
new ImmutableMap.Builder<String, PropertyEntry<?>>()
Expand All @@ -52,7 +52,7 @@ public class CredentialConfig extends Config {
"Max cache time for the credential.",
false /* required */,
false /* immutable */,
DEFAULT_MAX_CREDENTIAL_CACHE_TIME /* default value */,
DEFAULT_CREDENTIAL_CACHE_MAX_TIME /* default value */,
false /* hidden */,
false /* reserved */))
.put(
Expand All @@ -62,24 +62,24 @@ public class CredentialConfig extends Config {
"Max size for the credential cache.",
false /* required */,
false /* immutable */,
DEFAULT_MAX_CREDENTIAL_CACHE_SIZE /* default value */,
DEFAULT_CREDENTIAL_CACHE_MAX_SIZE /* default value */,
false /* hidden */,
false /* reserved */))
.build();

public static final ConfigEntry<Long> CREDENTIAL_MAX_CACHE_TIME =
public static final ConfigEntry<Long> CREDENTIAL_CACHE_MAX_TIME_IN_SECS =
new ConfigBuilder(CredentialConstants.CREDENTIAL_CACHE_MAX_TIME_IN_SECS)
.doc("Max cache time for the credential.")
.version(ConfigConstants.VERSION_0_8_0)
.longConf()
.createWithDefault(DEFAULT_MAX_CREDENTIAL_CACHE_TIME);
.createWithDefault(DEFAULT_CREDENTIAL_CACHE_MAX_TIME);

public static final ConfigEntry<Long> CREDENTIAL_MAX_CACHE_SIZE =
public static final ConfigEntry<Long> CREDENTIAL_CACHE_MAZ_SIZE =
new ConfigBuilder(CredentialConstants.CREDENTIAL_CACHE_MAX_SIZE)
.doc("Max cache size for the credential.")
.version(ConfigConstants.VERSION_0_8_0)
.longConf()
.createWithDefault(DEFAULT_MAX_CREDENTIAL_CACHE_SIZE);
.createWithDefault(DEFAULT_CREDENTIAL_CACHE_MAX_SIZE);

public CredentialConfig(Map<String, String> properties) {
super(false);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
public class TestCredentialCacheKey {

@Test
void testCredentialCache() {
void testCredentialCacheKey() {

PathBasedCredentialContext context1 =
new PathBasedCredentialContext("user1", ImmutableSet.of("path1"), ImmutableSet.of("path2"));
Expand Down

0 comments on commit 0a485b4

Please sign in to comment.