fix(dlf): refresh expiring credentials

crates/paimon/src/api/auth/dlf_provider.rs

@@ -167,10 +167,24 @@ impl DLFECSTokenLoader {
     /// Get the token from ECS metadata service.
     async fn get_token(&self, url: &str) -> Result<DLFToken> {
         let token_json = self.http_client.get(url).await?;
-        serde_json::from_str(&token_json).map_err(|e| Error::DataInvalid {
-            message: format!("Failed to parse token JSON: {e}"),
-            source: None,
-        })
+        let mut token: DLFToken =
+            serde_json::from_str(&token_json).map_err(|e| Error::DataInvalid {
+                message: format!("Failed to parse token JSON: {e}"),
+                source: None,
+            })?;
+        if token.expiration_at_millis.is_none() {
+            if let Some(expiration) = token.expiration.as_deref() {
+                token.expiration_at_millis = Some(
+                    DLFToken::parse_expiration_to_millis(expiration).ok_or_else(|| {
+                        Error::DataInvalid {
+                            message: format!("Failed to parse token Expiration: {expiration}"),
+                            source: None,
+                        }
+                    })?,
+                );
+            }
+        }
+        Ok(token)
     }
 
     /// Build the token URL from base URL and role name.
@@ -473,4 +487,64 @@ mod tests {
         let millis = DLFToken::parse_expiration_to_millis(expiration);
         assert!(millis.is_some());
     }
+
+    struct RotatingTokenLoader {
+        requests: std::sync::atomic::AtomicUsize,
+    }
+
+    #[async_trait]
+    impl DLFTokenLoader for RotatingTokenLoader {
+        async fn load_token(&self) -> Result<DLFToken> {
+            use std::sync::atomic::Ordering;
+
+            let request = self.requests.fetch_add(1, Ordering::SeqCst);
+            let lifetime = if request == 0 {
+                TOKEN_EXPIRATION_SAFE_TIME_MILLIS / 2
+            } else {
+                TOKEN_EXPIRATION_SAFE_TIME_MILLIS * 2
+            };
+            Ok(DLFToken::new(
+                format!("key-{request}"),
+                "secret",
+                None,
+                Some(Utc::now().timestamp_millis() + lifetime),
+                None,
+            ))
+        }
+
+        fn description(&self) -> &str {
+            "test"
+        }
+    }
+
+    #[tokio::test]
+    async fn test_refreshes_expiring_loaded_token() {
+        use std::sync::atomic::Ordering;
+
+        let loader = Arc::new(RotatingTokenLoader {
+            requests: std::sync::atomic::AtomicUsize::new(0),
+        });
+        let provider = DLFAuthProvider::new(
+            "https://dlf.cn-hangzhou.aliyuncs.com",
+            "cn-hangzhou",
+            "default",
+            None,
+            Some(loader.clone()),
+        )
+        .unwrap();
+
+        assert_eq!(
+            provider.get_or_refresh_token().await.unwrap().access_key_id,
+            "key-0"
+        );
+        assert_eq!(
+            provider.get_or_refresh_token().await.unwrap().access_key_id,
+            "key-1"
+        );
+        assert_eq!(
+            provider.get_or_refresh_token().await.unwrap().access_key_id,
+            "key-1"
+        );
+        assert_eq!(loader.requests.load(Ordering::SeqCst), 2);
+    }
 }