Merge pull request #1 from snazy/opa-authorizer-file-refresh

Make token-refresh asynchronous ...

Author: sungwy

extensions/auth/opa/impl/build.gradle.kts

@@ -31,6 +31,7 @@ dependencies {
   implementation(libs.guava)
   implementation(libs.slf4j.api)
   implementation(libs.auth0.jwt)
+  implementation(project(":polaris-async-api"))
 
   // Iceberg dependency for ForbiddenException
   implementation(platform(libs.iceberg.bom))
@@ -47,4 +48,7 @@ dependencies {
   testImplementation(libs.assertj.core)
   testImplementation(libs.mockito.core)
   testImplementation(libs.threeten.extra)
+  testImplementation(testFixtures(project(":polaris-async-api")))
+  testImplementation(project(":polaris-async-java"))
+  testImplementation(project(":polaris-idgen-mocks"))
 }

extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerFactory.java

@@ -36,6 +36,7 @@
 import org.apache.polaris.extension.auth.opa.token.BearerTokenProvider;
 import org.apache.polaris.extension.auth.opa.token.FileBearerTokenProvider;
 import org.apache.polaris.extension.auth.opa.token.StaticBearerTokenProvider;
+import org.apache.polaris.nosql.async.AsyncExec;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -48,14 +49,17 @@ class OpaPolarisAuthorizerFactory implements PolarisAuthorizerFactory {
 
   private final OpaAuthorizationConfig opaConfig;
   private final Clock clock;
+  private final ObjectMapper objectMapper;
+  private final AsyncExec asyncExec;
   private CloseableHttpClient httpClient;
   private BearerTokenProvider bearerTokenProvider;
-  private ObjectMapper objectMapper;
 
   @Inject
-  public OpaPolarisAuthorizerFactory(OpaAuthorizationConfig opaConfig, Clock clock) {
+  public OpaPolarisAuthorizerFactory(
+      OpaAuthorizationConfig opaConfig, Clock clock, AsyncExec asyncExec) {
     this.opaConfig = opaConfig;
     this.clock = clock;
+    this.asyncExec = asyncExec;
     this.objectMapper = new ObjectMapper();
   }
 
@@ -167,7 +171,13 @@ private BearerTokenProvider createBearerTokenProvider(
       Duration jwtExpirationBuffer = fileConfig.jwtExpirationBuffer().orElse(Duration.ofMinutes(1));
 
       return new FileBearerTokenProvider(
-          fileConfig.path(), refreshInterval, jwtExpirationRefresh, jwtExpirationBuffer, clock);
+          fileConfig.path(),
+          refreshInterval,
+          jwtExpirationRefresh,
+          jwtExpirationBuffer,
+          Duration.ofSeconds(5),
+          asyncExec,
+          clock::instant);
     } else {
       throw new IllegalStateException(
           "No bearer token configuration found. Must specify either 'static-token' or 'file-based'");

extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/token/FileBearerTokenProvider.java

@@ -18,21 +18,26 @@
  */
 package org.apache.polaris.extension.auth.opa.token;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import com.auth0.jwt.JWT;
 import com.auth0.jwt.exceptions.JWTDecodeException;
 import com.auth0.jwt.interfaces.DecodedJWT;
-import com.google.common.base.Strings;
 import jakarta.annotation.Nullable;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Date;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+import org.apache.polaris.nosql.async.AsyncExec;
+import org.apache.polaris.nosql.async.Cancelable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -58,12 +63,16 @@ public class FileBearerTokenProvider implements BearerTokenProvider {
   private final Duration refreshInterval;
   private final boolean jwtExpirationRefresh;
   private final Duration jwtExpirationBuffer;
-  private final Clock clock;
+  private final Supplier<Instant> clock;
   private final AtomicBoolean refreshLock = new AtomicBoolean();
+  private final AsyncExec asyncExec;
+  private final CompletableFuture<String> initialTokenFuture = new CompletableFuture<>();
+  private final long initialTokenWaitMillis;
 
   private volatile String cachedToken;
   private volatile Instant lastRefresh;
   private volatile Instant nextRefresh;
+  private volatile Cancelable<?> refreshTask;
 
   /**
    * Create a new file-based token provider with JWT expiration support.
@@ -80,24 +89,23 @@ public FileBearerTokenProvider(
       Duration refreshInterval,
       boolean jwtExpirationRefresh,
       Duration jwtExpirationBuffer,
-      Clock clock) {
+      Duration initialTokenWait,
+      AsyncExec asyncExec,
+      Supplier<Instant> clock) {
     this.tokenFilePath = tokenFilePath;
     this.refreshInterval = refreshInterval;
     this.jwtExpirationRefresh = jwtExpirationRefresh;
     this.jwtExpirationBuffer = jwtExpirationBuffer;
+    this.initialTokenWaitMillis = initialTokenWait.toMillis();
     this.clock = clock;
+    this.asyncExec = asyncExec;
 
-    // Load initial token eagerly to avoid race conditions during first getToken() calls
-    this.cachedToken = loadTokenFromFile();
-    if (Strings.isNullOrEmpty(this.cachedToken)) {
-      throw new IllegalStateException(
-          "Failed to load initial bearer token from file: "
-              + tokenFilePath
-              + ". This is required for OPA authorization.");
-    }
+    this.nextRefresh = Instant.MIN;
+    this.lastRefresh = Instant.MIN;
+    // start refreshing the token (immediately)
+    scheduleRefreshAttempt(Duration.ZERO);
 
-    this.lastRefresh = clock.instant();
-    this.nextRefresh = calculateNextRefresh(this.cachedToken);
+    checkState(Files.isReadable(tokenFilePath), "OPA token file does not exist or is not readable");
 
     logger.debug(
         "Created file token provider for path: {} with refresh interval: {}, JWT expiration refresh: {}, JWT buffer: {}, next refresh: {}",
@@ -110,56 +118,74 @@ public FileBearerTokenProvider(
 
   @Override
   public String getToken() {
-    // Check if we need to refresh
-    if (shouldRefresh()) {
-      refreshToken();
+    String token = cachedToken;
+    if (token != null) {
+      // Regular case, we have a cached token
+      return cachedToken;
     }
-
-    // Token is guaranteed to be present after construction, but check anyway for safety
-    if (Strings.isNullOrEmpty(cachedToken)) {
-      throw new RuntimeException(
-          "Bearer token is unexpectedly empty. This should not happen after successful construction.");
+    // We get here if the cached token is null, which means that the initial token
+    // has not been loaded yet.
+    // In this case we wait for the configured amount of time
+    // (5 seconds in production, much lower in tests).
+    try {
+      return initialTokenFuture.get(initialTokenWaitMillis, TimeUnit.MILLISECONDS);
+    } catch (Exception e) {
+      throw new IllegalStateException("Failed to read initial OPA bearer token", e);
     }
-    return cachedToken;
   }
 
   @Override
   public void close() {
     cachedToken = null;
+    Cancelable<?> task = refreshTask;
+    if (task != null) {
+      refreshTask.cancel();
+    }
+  }
+
+  private void refreshTokenAttempt() {
+    boolean isInitialRefresh = cachedToken == null;
+    Duration delay;
+    if (doRefreshToken()) {
+      delay = Duration.between(clock.get(), nextRefresh);
+      if (isInitialRefresh) {
+        // If we have never cached a token, complete the initial token-future to "unblock"
+        // getToken() call sites waiting for it.
+        initialTokenFuture.complete(cachedToken);
+      }
+    } else {
+      // Token refresh did not succeed, retry soon
+      delay = Duration.ofSeconds(1); // TODO configurable ?
+    }
+    scheduleRefreshAttempt(delay);
   }
 
-  private boolean shouldRefresh() {
-    return clock.instant().isAfter(nextRefresh);
+  private void scheduleRefreshAttempt(Duration delay) {
+    this.refreshTask = asyncExec.schedule(this::refreshTokenAttempt, delay);
   }
 
-  private void refreshToken() {
-    // Only one thread should refresh at a time. Other threads will use the cached token.
-    if (!refreshLock.compareAndSet(false, true)) {
-      return;
+  private boolean doRefreshToken() {
+    String newToken = loadTokenFromFile();
+
+    // Only update cached token if we successfully loaded a new one
+    if (newToken == null) {
+      logger.debug("Couldn't load new bearer token from {}, will retry.", tokenFilePath);
+      return false;
     }
-    try {
-      String newToken = loadTokenFromFile();
+    cachedToken = newToken;
 
-      // Only update cached token if we successfully loaded a new one
-      if (newToken == null) {
-        logger.debug("Couldn't load new bearer token from {}, will retry.", tokenFilePath);
-        return;
-      }
-      cachedToken = newToken;
+    lastRefresh = clock.get();
 
-      lastRefresh = clock.instant();
+    // Calculate next refresh time based on current token (may be cached)
+    nextRefresh = calculateNextRefresh(cachedToken);
 
-      // Calculate next refresh time based on current token (may be cached)
-      nextRefresh = calculateNextRefresh(cachedToken);
+    logger.debug(
+        "Token refreshed from file: {} (token present: {}), next refresh: {}",
+        tokenFilePath,
+        cachedToken != null && !cachedToken.isEmpty(),
+        nextRefresh);
 
-      logger.debug(
-          "Token refreshed from file: {} (token present: {}), next refresh: {}",
-          tokenFilePath,
-          cachedToken != null && !cachedToken.isEmpty(),
-          nextRefresh);
-    } finally {
-      refreshLock.set(false);
-    }
+    return true;
   }
 
   /** Calculate when the next refresh should occur based on JWT expiration or fixed interval. */
@@ -176,7 +202,7 @@ private Instant calculateNextRefresh(String token) {
       Instant refreshTime = expiration.get().minus(jwtExpirationBuffer);
 
       // Ensure refresh time is in the future and not too soon (at least 1 second)
-      Instant minRefreshTime = clock.instant().plus(Duration.ofSeconds(1));
+      Instant minRefreshTime = clock.get().plus(Duration.ofSeconds(1));
       if (refreshTime.isBefore(minRefreshTime)) {
         logger.warn(
             "JWT expires too soon ({}), using minimum refresh interval instead", expiration.get());

extensions/auth/opa/impl/src/test/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerFactoryTest.java

@@ -31,6 +31,7 @@
 import java.util.Optional;
 import org.apache.polaris.core.config.RealmConfig;
 import org.apache.polaris.extension.auth.opa.token.FileBearerTokenProvider;
+import org.apache.polaris.nosql.async.java.JavaPoolAsyncExec;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -63,14 +64,16 @@ public void testFactoryWithStaticTokenConfiguration() {
     when(opaConfig.auth()).thenReturn(authConfig);
     when(opaConfig.http()).thenReturn(httpConfig);
 
-    OpaPolarisAuthorizerFactory factory =
-        new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC());
+    try (JavaPoolAsyncExec asyncExec = new JavaPoolAsyncExec()) {
+      OpaPolarisAuthorizerFactory factory =
+          new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC(), asyncExec);
 
-    // Create authorizer
-    RealmConfig realmConfig = mock(RealmConfig.class);
-    OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
+      // Create authorizer
+      RealmConfig realmConfig = mock(RealmConfig.class);
+      OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
 
-    assertThat(authorizer).isNotNull();
+      assertThat(authorizer).isNotNull();
+    }
   }
 
   @Test
@@ -106,22 +109,30 @@ public void testFactoryWithFileBasedTokenConfiguration() throws IOException {
     when(opaConfig.auth()).thenReturn(authConfig);
     when(opaConfig.http()).thenReturn(httpConfig);
 
-    OpaPolarisAuthorizerFactory factory =
-        new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC());
-
-    // Create authorizer
-    RealmConfig realmConfig = mock(RealmConfig.class);
-    OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
-
-    assertThat(authorizer).isNotNull();
-
-    // Also verify that the token provider actually reads from the file
-    try (FileBearerTokenProvider provider =
-        new FileBearerTokenProvider(
-            tokenFile, Duration.ofMinutes(5), true, Duration.ofMinutes(1), Clock.systemUTC())) {
-
-      String actualToken = provider.getToken();
-      assertThat(actualToken).isEqualTo(tokenValue);
+    try (JavaPoolAsyncExec asyncExec = new JavaPoolAsyncExec()) {
+      OpaPolarisAuthorizerFactory factory =
+          new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC(), asyncExec);
+
+      // Create authorizer
+      RealmConfig realmConfig = mock(RealmConfig.class);
+      OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
+
+      assertThat(authorizer).isNotNull();
+
+      // Also verify that the token provider actually reads from the file
+      try (FileBearerTokenProvider provider =
+          new FileBearerTokenProvider(
+              tokenFile,
+              Duration.ofMinutes(5),
+              true,
+              Duration.ofMinutes(1),
+              Duration.ofSeconds(10),
+              asyncExec,
+              Clock.systemUTC()::instant)) {
+
+        String actualToken = provider.getToken();
+        assertThat(actualToken).isEqualTo(tokenValue);
+      }
     }
   }
 
@@ -141,14 +152,16 @@ public void testFactoryWithNoTokenConfiguration() {
     when(opaConfig.auth()).thenReturn(authConfig);
     when(opaConfig.http()).thenReturn(httpConfig);
 
-    OpaPolarisAuthorizerFactory factory =
-        new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC());
+    try (JavaPoolAsyncExec asyncExec = new JavaPoolAsyncExec()) {
+      OpaPolarisAuthorizerFactory factory =
+          new OpaPolarisAuthorizerFactory(opaConfig, Clock.systemUTC(), asyncExec);
 
-    // Create authorizer
-    RealmConfig realmConfig = mock(RealmConfig.class);
-    OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
+      // Create authorizer
+      RealmConfig realmConfig = mock(RealmConfig.class);
+      OpaPolarisAuthorizer authorizer = (OpaPolarisAuthorizer) factory.create(realmConfig);
 
-    assertThat(authorizer).isNotNull();
+      assertThat(authorizer).isNotNull();
+    }
   }
 
   private OpaAuthorizationConfig.HttpConfig createMockHttpConfig() {