From bc039235da386de822bf8edc786860e933a0f34f Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Sat, 14 Dec 2024 15:33:21 +0800
Subject: [PATCH 1/3] fix(catalog/rest): Ensure token been reused correctly

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 crates/catalog/rest/src/catalog.rs | 13 ++++++++-----
 crates/catalog/rest/src/client.rs  | 27 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs
index fce5fe2be0..b6c3b8dcf9 100644
--- a/crates/catalog/rest/src/catalog.rs
+++ b/crates/catalog/rest/src/catalog.rs
@@ -256,9 +256,10 @@ impl RestCatalog {
     async fn context(&self) -> Result<&RestContext> {
         self.ctx
             .get_or_try_init(|| async {
-                let catalog_config = RestCatalog::load_config(&self.user_config).await?;
+                let client = HttpClient::new(&self.user_config)?;
+                let catalog_config = RestCatalog::load_config(&client, &self.user_config).await?;
                 let config = self.user_config.clone().merge_with_config(catalog_config);
-                let client = HttpClient::new(&config)?;
+                let client = client.update_with(&config)?;
 
                 Ok(RestContext { config, client })
             })
@@ -268,9 +269,10 @@ impl RestCatalog {
     /// Load the runtime config from the server by user_config.
     ///
     /// It's required for a rest catalog to update it's config after creation.
-    async fn load_config(user_config: &RestCatalogConfig) -> Result<CatalogConfig> {
-        let client = HttpClient::new(user_config)?;
-
+    async fn load_config(
+        client: &HttpClient,
+        user_config: &RestCatalogConfig,
+    ) -> Result<CatalogConfig> {
         let mut request = client.request(Method::GET, user_config.config_endpoint());
 
         if let Some(warehouse_location) = &user_config.warehouse {
@@ -280,6 +282,7 @@ impl RestCatalog {
         let config = client
             .query::<CatalogConfig, ErrorResponse, OK>(request.build()?)
             .await?;
+
         Ok(config)
     }
 
diff --git a/crates/catalog/rest/src/client.rs b/crates/catalog/rest/src/client.rs
index 53dcd4ceec..1143008e27 100644
--- a/crates/catalog/rest/src/client.rs
+++ b/crates/catalog/rest/src/client.rs
@@ -54,6 +54,7 @@ impl Debug for HttpClient {
 }
 
 impl HttpClient {
+    /// Create a new http client.
     pub fn new(cfg: &RestCatalogConfig) -> Result<Self> {
         Ok(HttpClient {
             client: Client::new(),
@@ -66,6 +67,32 @@ impl HttpClient {
         })
     }
 
+    /// Update the http client with new configuration.
+    ///
+    /// If cfg carries new value, we will use cfg instead.
+    /// Otherwise, we will keep the old value.
+    pub fn update_with(self, cfg: &RestCatalogConfig) -> Result<Self> {
+        Ok(HttpClient {
+            client: self.client,
+
+            token: Mutex::new(
+                cfg.token()
+                    .or_else(|| self.token.into_inner().ok().flatten()),
+            ),
+            token_endpoint: (!cfg.get_token_endpoint().is_empty())
+                .then(|| cfg.get_token_endpoint())
+                .unwrap_or(self.token_endpoint),
+            credential: cfg.credential().or(self.credential),
+            extra_headers: (!cfg.extra_headers()?.is_empty())
+                .then(|| cfg.extra_headers())
+                .transpose()?
+                .unwrap_or(self.extra_headers),
+            extra_oauth_params: (!cfg.extra_oauth_params().is_empty())
+                .then(|| cfg.extra_oauth_params())
+                .unwrap_or(self.extra_oauth_params),
+        })
+    }
+
     /// This API is testing only to assert the token.
     #[cfg(test)]
     pub(crate) async fn token(&self) -> Option<String> {

From e06951ed6b47fcc42eda93c4229d211464d46aac Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Sat, 14 Dec 2024 16:15:22 +0800
Subject: [PATCH 2/3] Fix oauth test

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 crates/catalog/rest/src/catalog.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs
index b6c3b8dcf9..648b50a6c7 100644
--- a/crates/catalog/rest/src/catalog.rs
+++ b/crates/catalog/rest/src/catalog.rs
@@ -780,7 +780,7 @@ mod tests {
                 "expires_in": 86400
                 }"#,
             )
-            .expect(2)
+            .expect(1)
             .create_async()
             .await
     }

From 5bb16221c7d0f629770b9cfa3e5dcd235c457828 Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Sat, 14 Dec 2024 18:54:24 +0800
Subject: [PATCH 3/3] Fix tests

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 crates/catalog/rest/src/catalog.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs
index 648b50a6c7..96da5dc951 100644
--- a/crates/catalog/rest/src/catalog.rs
+++ b/crates/catalog/rest/src/catalog.rs
@@ -834,7 +834,7 @@ mod tests {
                 "expires_in": 86400
                 }"#,
             )
-            .expect(2)
+            .expect(1)
             .create_async()
             .await;