Xuanwo commented on code in PR #431:
URL: https://github.com/apache/iceberg-rust/pull/431#discussion_r1665083079


##########
crates/catalog/rest/src/client.rs:
##########
@@ -15,25 +15,165 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::types::{ErrorResponse, TokenResponse, OK};
+use crate::RestCatalogConfig;
 use iceberg::Result;
 use iceberg::{Error, ErrorKind};
 use reqwest::header::HeaderMap;
 use reqwest::{Client, IntoUrl, Method, Request, RequestBuilder, Response};
 use serde::de::DeserializeOwned;
+use std::collections::HashMap;
+use std::fmt::{Debug, Formatter};
+use std::sync::Mutex;
 
-#[derive(Debug)]
-pub(crate) struct HttpClient(Client);
+pub(crate) struct HttpClient {
+    client: Client,
+
+    /// The token to be used for authentication.
+    ///
+    /// It's possible to fetch the token from the server while needed.
+    token: Mutex<Option<String>>,
+    /// The token endpoint to be used for authentication.
+    token_endpoint: String,
+    /// The credential to be used for authentication.
+    credential: Option<(Option<String>, String)>,
+    /// Extra headers to be added to each request.
+    extra_headers: HeaderMap,
+    /// Extra oauth parameters to be added to each authentication request.
+    extra_oauth_params: HashMap<String, String>,
+}
+
+impl Debug for HttpClient {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HttpClient")
+            .field("client", &self.client)
+            .field("extra_headers", &self.extra_headers)
+            .finish_non_exhaustive()
+    }
+}
 
 impl HttpClient {
-    pub fn try_create(default_headers: HeaderMap) -> Result<Self> {
-        Ok(HttpClient(
-            Client::builder().default_headers(default_headers).build()?,
-        ))
+    pub fn new(cfg: &RestCatalogConfig) -> Result<Self> {
+        Ok(HttpClient {
+            client: Client::new(),
+
+            token: Mutex::new(cfg.token()),
+            token_endpoint: cfg.get_token_endpoint(),
+            credential: cfg.credential(),
+            extra_headers: cfg.extra_headers()?,
+            extra_oauth_params: cfg.extra_oauth_params(),
+        })
+    }
+
+    /// This API is testing only to assert the token.
+    #[cfg(test)]
+    pub(crate) async fn token(&self) -> Option<String> {
+        let mut req = self
+            .request(Method::GET, &self.token_endpoint)
+            .build()
+            .unwrap();
+        self.authenticate(&mut req).await.ok();
+        self.token.lock().unwrap().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().expect("lock poison").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)
+                })?,
+            );
+        }
+
+        // Credential must exist here.
+        let (client_id, client_secret) = self.credential.as_ref().unwrap();

Review Comment:
   Credential is impossible to be `None`. But it's fine for me to return an 
error instead.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For additional commands, e-mail: issues-h...@iceberg.apache.org

Reply via email to