when CredentialsProvider is soon expire force refresh

apache · Dec 6, 2024 · 9aafd86 · 9aafd86
1 parent ec074c2
commit 9aafd86
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 6 deletions.
diff --git a/paimon-core/src/main/java/org/apache/paimon/rest/auth/AuthSession.java b/paimon-core/src/main/java/org/apache/paimon/rest/auth/AuthSession.java
@@ -79,6 +79,9 @@ public static AuthSession fromRefreshCredentialsProvider(
     }
 
     public Map<String, String> getHeaders() {
+        if (this.credentialsProvider.keepRefreshed() && this.credentialsProvider.willSoonExpire()) {
+            refresh();
+        }
         return headers;
     }
 

diff --git a/paimon-core/src/main/java/org/apache/paimon/rest/auth/BearTokenFileCredentialsProvider.java b/paimon-core/src/main/java/org/apache/paimon/rest/auth/BearTokenFileCredentialsProvider.java
@@ -28,6 +28,9 @@
 
 /** credentials provider for get bear token from file. */
 public class BearTokenFileCredentialsProvider extends BaseBearTokenCredentialsProvider {
+
+    public static final double EXPIRED_FACTOR = 0.4;
+
     private final String tokenFilePath;
     private String token;
     private boolean keepRefreshed = false;
@@ -76,6 +79,16 @@ public boolean keepRefreshed() {
         return this.keepRefreshed;
     }
 
+    @Override
+    public boolean willSoonExpire() {
+        if (keepRefreshed()) {
+            return expiresAtMillis().get() - System.currentTimeMillis()
+                    < expiresInMills().get() * EXPIRED_FACTOR;
+        } else {
+            return false;
+        }
+    }
+
     @Override
     public Optional<Long> expiresAtMillis() {
         return Optional.ofNullable(this.expiresAtMillis);

diff --git a/paimon-core/src/main/java/org/apache/paimon/rest/auth/CredentialsProvider.java b/paimon-core/src/main/java/org/apache/paimon/rest/auth/CredentialsProvider.java
@@ -35,6 +35,10 @@ default boolean keepRefreshed() {
         return false;
     }
 
+    default boolean willSoonExpire() {
+        return false;
+    }
+
     default Optional<Long> expiresAtMillis() {
         return Optional.empty();
     }

diff --git a/paimon-core/src/test/java/org/apache/paimon/rest/auth/AuthSessionTest.java b/paimon-core/src/test/java/org/apache/paimon/rest/auth/AuthSessionTest.java
@@ -18,6 +18,7 @@
 
 package org.apache.paimon.rest.auth;
 
+import org.apache.paimon.utils.Pair;
 import org.apache.paimon.utils.ThreadPoolUtils;
 
 import org.apache.commons.io.FileUtils;
@@ -43,9 +44,9 @@ public class AuthSessionTest {
     public void testRefreshBearTokenFileCredentialsProvider()
             throws IOException, InterruptedException {
         String fileName = "token";
-        File tokenFile = folder.newFile(fileName);
-        String token = UUID.randomUUID().toString();
-        FileUtils.writeStringToFile(tokenFile, token);
+        Pair<File, String> tokenFile2Token = generateTokenAndWriteToFile(fileName);
+        String token = tokenFile2Token.getRight();
+        File tokenFile = tokenFile2Token.getLeft();
         Map<String, String> initialHeaders = new HashMap<>();
         long expiresInMillis = 1000L;
         CredentialsProvider credentialsProvider =
@@ -59,11 +60,46 @@ public void testRefreshBearTokenFileCredentialsProvider()
         Map<String, String> header = session.getHeaders();
         assertEquals(header.get("Authorization"), "Bearer " + token);
         tokenFile.delete();
-        tokenFile = folder.newFile(fileName);
-        token = UUID.randomUUID().toString();
-        FileUtils.writeStringToFile(tokenFile, token);
+        tokenFile2Token = generateTokenAndWriteToFile(fileName);
+        token = tokenFile2Token.getRight();
         Thread.sleep(expiresInMillis + 500L);
         header = session.getHeaders();
         assertEquals(header.get("Authorization"), "Bearer " + token);
     }
+
+    @Test
+    public void testRefreshCredentialsProviderIsSoonExpire()
+            throws IOException, InterruptedException {
+        String fileName = "token";
+        Pair<File, String> tokenFile2Token = generateTokenAndWriteToFile(fileName);
+        String token = tokenFile2Token.getRight();
+        File tokenFile = tokenFile2Token.getLeft();
+        Map<String, String> initialHeaders = new HashMap<>();
+        long expiresInMillis = 1000L;
+        CredentialsProvider credentialsProvider =
+                new BearTokenFileCredentialsProvider(
+                        tokenFile.getPath(), true, -1L, expiresInMillis);
+        AuthSession session =
+                AuthSession.fromRefreshCredentialsProvider(
+                        null, initialHeaders, credentialsProvider);
+        Map<String, String> header = session.getHeaders();
+        assertEquals(header.get("Authorization"), "Bearer " + token);
+        tokenFile.delete();
+        tokenFile2Token = generateTokenAndWriteToFile(fileName);
+        token = tokenFile2Token.getRight();
+        tokenFile = tokenFile2Token.getLeft();
+        FileUtils.writeStringToFile(tokenFile, token);
+        Thread.sleep(
+                (long) (expiresInMillis * (1 - BearTokenFileCredentialsProvider.EXPIRED_FACTOR))
+                        + 10L);
+        header = session.getHeaders();
+        assertEquals(header.get("Authorization"), "Bearer " + token);
+    }
+
+    private Pair<File, String> generateTokenAndWriteToFile(String fileName) throws IOException {
+        File tokenFile = folder.newFile(fileName);
+        String token = UUID.randomUUID().toString();
+        FileUtils.writeStringToFile(tokenFile, token);
+        return Pair.of(tokenFile, token);
+    }
 }