This is an automated email from the ASF dual-hosted git repository.
xuanwo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
The following commit(s) were added to refs/heads/main by this push:
new a9614834 feat(iceberg-catalog-rest): expose invalidate_token,
regenerate_token APIs (#1465)
a9614834 is described below
commit a96148349a141436efb93d94bdce306873536744
Author: Connor McArthur <[email protected]>
AuthorDate: Wed Jun 25 02:41:17 2025 -0500
feat(iceberg-catalog-rest): expose invalidate_token, regenerate_token APIs
(#1465)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Partially implements #437
## What changes are included in this PR?
<!--
Provide a summary of the modifications in this PR. List the main changes
such as new features, bug fixes, refactoring, or any other updates.
-->
Per discussion here:
https://github.com/apache/iceberg-rust/issues/437#issuecomment-2997408926
This branch exposes two public methods on the rest catalog for managing
tokens: `invalidate_token` and `regenerate_token`.
- `invalidate_token`: takes the token lock and sets the token to `None`,
forcing re-auth on the next request
- `regenerate_token`: calls the oauth endpoint to generate a new token.
if successful, takes the token lock and writes the new token.
## Are these changes tested?
<!--
Specify what test covers (unit test, integration test, etc.).
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
Yes -- I've added new unit tests to cover both methods.
---
crates/catalog/rest/src/catalog.rs | 169 ++++++++++++++++++++++++++++++++++---
crates/catalog/rest/src/client.rs | 89 +++++++++++--------
2 files changed, 214 insertions(+), 44 deletions(-)
diff --git a/crates/catalog/rest/src/catalog.rs
b/crates/catalog/rest/src/catalog.rs
index 3d1f7cf8..a3c24db2 100644
--- a/crates/catalog/rest/src/catalog.rs
+++ b/crates/catalog/rest/src/catalog.rs
@@ -318,6 +318,22 @@ impl RestCatalog {
Ok(file_io)
}
+
+ /// Invalidate the current token without generating a new one. On the next
request, the client
+ /// will attempt to generate a new token.
+ pub async fn invalidate_token(&self) -> Result<()> {
+ self.context().await?.client.invalidate_token().await
+ }
+
+ /// Invalidate the current token and set a new one. Generates a new token
before invalidating
+ /// the current token, meaning the old token will be used until this
function acquires the lock
+ /// and overwrites the token.
+ ///
+ /// If credential is invalid, or the request fails, this method will
return an error and leave
+ /// the current token unchanged.
+ pub async fn regenerate_token(&self) -> Result<()> {
+ self.context().await?.client.regenerate_token().await
+ }
}
/// All requests and expected responses are derived from the REST catalog API
spec:
@@ -860,21 +876,27 @@ mod tests {
}
async fn create_oauth_mock(server: &mut ServerGuard) -> Mock {
- create_oauth_mock_with_path(server, "/v1/oauth/tokens").await
+ create_oauth_mock_with_path(server, "/v1/oauth/tokens",
"ey000000000000", 200).await
}
- async fn create_oauth_mock_with_path(server: &mut ServerGuard, path: &str)
-> Mock {
- server
- .mock("POST", path)
- .with_status(200)
- .with_body(
- r#"{
- "access_token": "ey000000000000",
+ async fn create_oauth_mock_with_path(
+ server: &mut ServerGuard,
+ path: &str,
+ token: &str,
+ status: usize,
+ ) -> Mock {
+ let body = format!(
+ r#"{{
+ "access_token": "{token}",
"token_type": "Bearer",
"issued_token_type":
"urn:ietf:params:oauth:token-type:access_token",
"expires_in": 86400
- }"#,
- )
+ }}"#
+ );
+ server
+ .mock("POST", path)
+ .with_status(status)
+ .with_body(body)
.expect(1)
.create_async()
.await
@@ -949,6 +971,129 @@ mod tests {
assert_eq!(token, Some("ey000000000000".to_string()));
}
+ #[tokio::test]
+ async fn test_invalidate_token() {
+ let mut server = Server::new_async().await;
+ let oauth_mock = create_oauth_mock(&mut server).await;
+ let config_mock = create_config_mock(&mut server).await;
+
+ let mut props = HashMap::new();
+ props.insert("credential".to_string(), "client1:secret1".to_string());
+
+ let catalog = RestCatalog::new(
+ RestCatalogConfig::builder()
+ .uri(server.url())
+ .props(props)
+ .build(),
+ );
+
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ config_mock.assert_async().await;
+ assert_eq!(token, Some("ey000000000000".to_string()));
+
+ let oauth_mock =
+ create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens",
"ey000000000001", 200)
+ .await;
+ catalog.invalidate_token().await.unwrap();
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ assert_eq!(token, Some("ey000000000001".to_string()));
+ }
+
+ #[tokio::test]
+ async fn test_invalidate_token_failing_request() {
+ let mut server = Server::new_async().await;
+ let oauth_mock = create_oauth_mock(&mut server).await;
+ let config_mock = create_config_mock(&mut server).await;
+
+ let mut props = HashMap::new();
+ props.insert("credential".to_string(), "client1:secret1".to_string());
+
+ let catalog = RestCatalog::new(
+ RestCatalogConfig::builder()
+ .uri(server.url())
+ .props(props)
+ .build(),
+ );
+
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ config_mock.assert_async().await;
+ assert_eq!(token, Some("ey000000000000".to_string()));
+
+ let oauth_mock =
+ create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens",
"ey000000000001", 500)
+ .await;
+ catalog.invalidate_token().await.unwrap();
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ assert_eq!(token, None);
+ }
+
+ #[tokio::test]
+ async fn test_regenerate_token() {
+ let mut server = Server::new_async().await;
+ let oauth_mock = create_oauth_mock(&mut server).await;
+ let config_mock = create_config_mock(&mut server).await;
+
+ let mut props = HashMap::new();
+ props.insert("credential".to_string(), "client1:secret1".to_string());
+
+ let catalog = RestCatalog::new(
+ RestCatalogConfig::builder()
+ .uri(server.url())
+ .props(props)
+ .build(),
+ );
+
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ config_mock.assert_async().await;
+ assert_eq!(token, Some("ey000000000000".to_string()));
+
+ let oauth_mock =
+ create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens",
"ey000000000001", 200)
+ .await;
+ catalog.regenerate_token().await.unwrap();
+ oauth_mock.assert_async().await;
+ let token = catalog.context().await.unwrap().client.token().await;
+ assert_eq!(token, Some("ey000000000001".to_string()));
+ }
+
+ #[tokio::test]
+ async fn test_regenerate_token_failing_request() {
+ let mut server = Server::new_async().await;
+ let oauth_mock = create_oauth_mock(&mut server).await;
+ let config_mock = create_config_mock(&mut server).await;
+
+ let mut props = HashMap::new();
+ props.insert("credential".to_string(), "client1:secret1".to_string());
+
+ let catalog = RestCatalog::new(
+ RestCatalogConfig::builder()
+ .uri(server.url())
+ .props(props)
+ .build(),
+ );
+
+ let token = catalog.context().await.unwrap().client.token().await;
+ oauth_mock.assert_async().await;
+ config_mock.assert_async().await;
+ assert_eq!(token, Some("ey000000000000".to_string()));
+
+ let oauth_mock =
+ create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens",
"ey000000000001", 500)
+ .await;
+ let invalidate_result = catalog.regenerate_token().await;
+ assert!(invalidate_result.is_err());
+ oauth_mock.assert_async().await;
+ let token = catalog.context().await.unwrap().client.token().await;
+
+ // original token is left intact
+ assert_eq!(token, Some("ey000000000000".to_string()));
+ }
+
#[tokio::test]
async fn test_http_headers() {
let server = Server::new_async().await;
@@ -1026,7 +1171,9 @@ mod tests {
let mut auth_server = Server::new_async().await;
let auth_server_path = "/some/path";
- let oauth_mock = create_oauth_mock_with_path(&mut auth_server,
auth_server_path).await;
+ let oauth_mock =
+ create_oauth_mock_with_path(&mut auth_server, auth_server_path,
"ey000000000000", 200)
+ .await;
let mut props = HashMap::new();
props.insert("credential".to_string(), "client1:secret1".to_string());
diff --git a/crates/catalog/rest/src/client.rs
b/crates/catalog/rest/src/client.rs
index fe311a71..0b9af9b5 100644
--- a/crates/catalog/rest/src/client.rs
+++ b/crates/catalog/rest/src/client.rs
@@ -106,38 +106,7 @@ impl HttpClient {
self.token.lock().await.clone()
}
- /// Authenticate the request by filling token.
- ///
- /// - If neither token nor credential is provided, this method will do
nothing.
- /// - If only credential is provided, this method will try to fetch token
from the server.
- /// - If token is provided, this method will use the token directly.
- ///
- /// # TODO
- ///
- /// Support refreshing token while needed.
- async fn authenticate(&self, req: &mut Request) -> Result<()> {
- // Clone the token from lock without holding the lock for entire
function.
- let token = self.token.lock().await.clone();
-
- if self.credential.is_none() && token.is_none() {
- return Ok(());
- }
-
- // Use token if provided.
- if let Some(token) = &token {
- req.headers_mut().insert(
- http::header::AUTHORIZATION,
- format!("Bearer {token}").parse().map_err(|e| {
- Error::new(
- ErrorKind::DataInvalid,
- "Invalid token received from catalog server!",
- )
- .with_source(e)
- })?,
- );
- return Ok(());
- }
-
+ async fn exchange_credential_for_token(&self) -> Result<String> {
// Credential must exist here.
let (client_id, client_secret) =
self.credential.as_ref().ok_or_else(|| {
Error::new(
@@ -202,7 +171,61 @@ impl HttpClient {
})?;
Err(Error::from(e))
}?;
- let token = auth_res.access_token;
+ Ok(auth_res.access_token)
+ }
+
+ /// Invalidate the current token without generating a new one. On the next
request, the client
+ /// will attempt to generate a new token.
+ pub(crate) async fn invalidate_token(&self) -> Result<()> {
+ *self.token.lock().await = None;
+ Ok(())
+ }
+
+ /// Invalidate the current token and set a new one. Generates a new token
before invalidating
+ /// the current token, meaning the old token will be used until this
function acquires the lock
+ /// and overwrites the token.
+ ///
+ /// If credential is invalid, or the request fails, this method will
return an error and leave
+ /// the current token unchanged.
+ pub(crate) async fn regenerate_token(&self) -> Result<()> {
+ let new_token = self.exchange_credential_for_token().await?;
+ *self.token.lock().await = Some(new_token.clone());
+ Ok(())
+ }
+
+ /// Authenticate the request by filling token.
+ ///
+ /// - If neither token nor credential is provided, this method will do
nothing.
+ /// - If only credential is provided, this method will try to fetch token
from the server.
+ /// - If token is provided, this method will use the token directly.
+ ///
+ /// # TODO
+ ///
+ /// Support refreshing token while needed.
+ async fn authenticate(&self, req: &mut Request) -> Result<()> {
+ // Clone the token from lock without holding the lock for entire
function.
+ let token = self.token.lock().await.clone();
+
+ if self.credential.is_none() && token.is_none() {
+ return Ok(());
+ }
+
+ // Use token if provided.
+ if let Some(token) = &token {
+ req.headers_mut().insert(
+ http::header::AUTHORIZATION,
+ format!("Bearer {token}").parse().map_err(|e| {
+ Error::new(
+ ErrorKind::DataInvalid,
+ "Invalid token received from catalog server!",
+ )
+ .with_source(e)
+ })?,
+ );
+ return Ok(());
+ }
+
+ let token = self.exchange_credential_for_token().await?;
// Update token.
*self.token.lock().await = Some(token.clone());
// Insert token in request.
