From 8d691ca33805cf7158d07feab5f37596597a46a1 Mon Sep 17 00:00:00 2001 From: Hongyi Wang Date: Mon, 25 Mar 2024 10:56:04 -0700 Subject: [PATCH 1/3] Support customized header in Rest catalog client --- crates/catalog/rest/src/catalog.rs | 93 +++++++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 2 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 33a2ea2e0c..3472364d55 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -18,6 +18,7 @@ //! This module contains rest catalog implementation. use std::collections::HashMap; +use std::str::FromStr; use async_trait::async_trait; use reqwest::header::{self, HeaderMap, HeaderName, HeaderValue}; @@ -100,8 +101,7 @@ impl RestCatalogConfig { [&self.uri, PATH_V1, "oauth", "tokens"].join("/") } - fn try_create_rest_client(&self) -> Result { - // TODO: We will add ssl config, sigv4 later + fn get_default_headers(&self) -> Result { let mut headers = HeaderMap::from_iter([ ( header::CONTENT_TYPE, @@ -130,6 +130,33 @@ impl RestCatalogConfig { ); } + for (key, value) in self.props.iter() { + if let Some(stripped_key) = key.strip_prefix("header.") { + headers.insert( + HeaderName::from_str(stripped_key).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid header name: {stripped_key}!"), + ) + .with_source(e) + })?, + HeaderValue::from_str(value).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid header value: {value}!"), + ) + .with_source(e) + })?, + ); + } + } + Ok(headers) + } + + fn try_create_rest_client(&self) -> Result { + // TODO: We will add ssl config, sigv4 later + let headers = self.get_default_headers()?; + Ok(HttpClient( Client::builder().default_headers(headers).build()?, )) @@ -956,6 +983,68 @@ mod tests { ); } + #[tokio::test] + async fn test_get_default_headers() { + let server = Server::new_async().await; + let mut props = HashMap::new(); + props.insert("credential".to_string(), "client1:secret1".to_string()); + + let config = RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(); + let headers: HeaderMap = config.get_default_headers().unwrap(); + + let expected_headers = HeaderMap::from_iter([ + ( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ), + ( + HeaderName::from_static("x-client-version"), + HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), + ), + ( + header::USER_AGENT, + HeaderValue::from_str(&format!("iceberg-rs/{}", CARGO_PKG_VERSION)).unwrap(), + ), + ]); + assert_eq!(headers, expected_headers); + } + + #[tokio::test] + async fn test_get_default_headers_with_custom_headers() { + let server = Server::new_async().await; + let mut props = HashMap::new(); + props.insert("credential".to_string(), "client1:secret1".to_string()); + props.insert( + "header.content-type".to_string(), + "application/yaml".to_string(), + ); + + let config = RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(); + let headers: HeaderMap = config.get_default_headers().unwrap(); + + let expected_headers = HeaderMap::from_iter([ + ( + header::CONTENT_TYPE, + HeaderValue::from_static("application/yaml"), + ), + ( + HeaderName::from_static("x-client-version"), + HeaderValue::from_static(ICEBERG_REST_SPEC_VERSION), + ), + ( + header::USER_AGENT, + HeaderValue::from_str(&format!("iceberg-rs/{}", CARGO_PKG_VERSION)).unwrap(), + ), + ]); + assert_eq!(headers, expected_headers); + } + #[tokio::test] async fn test_list_namespace() { let mut server = Server::new_async().await; From f16f1a992ccc2d32f3a55fec23d75fb50a755640 Mon Sep 17 00:00:00 2001 From: Hongyi Wang Date: Tue, 26 Mar 2024 14:49:28 -0700 Subject: [PATCH 2/3] Avoid overwriting default headers --- crates/catalog/rest/src/catalog.rs | 45 +++++++++++++++++++----------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index dbbcd5e45b..c42268203f 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -131,22 +131,25 @@ impl RestCatalogConfig { for (key, value) in self.props.iter() { if let Some(stripped_key) = key.strip_prefix("header.") { - headers.insert( - HeaderName::from_str(stripped_key).map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!("Invalid header name: {stripped_key}!"), - ) - .with_source(e) - })?, - HeaderValue::from_str(value).map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!("Invalid header value: {value}!"), - ) - .with_source(e) - })?, - ); + // Avoid overwriting default headers + if !headers.contains_key(stripped_key) { + headers.insert( + HeaderName::from_str(stripped_key).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid header name: {stripped_key}!"), + ) + .with_source(e) + })?, + HeaderValue::from_str(value).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid header value: {value}!"), + ) + .with_source(e) + })?, + ); + } } } Ok(headers) @@ -1020,6 +1023,10 @@ mod tests { "header.content-type".to_string(), "application/yaml".to_string(), ); + props.insert( + "header.customized-header".to_string(), + "some/value".to_string(), + ); let config = RestCatalogConfig::builder() .uri(server.url()) @@ -1030,7 +1037,7 @@ mod tests { let expected_headers = HeaderMap::from_iter([ ( header::CONTENT_TYPE, - HeaderValue::from_static("application/yaml"), + HeaderValue::from_static("application/json"), ), ( HeaderName::from_static("x-client-version"), @@ -1040,6 +1047,10 @@ mod tests { header::USER_AGENT, HeaderValue::from_str(&format!("iceberg-rs/{}", CARGO_PKG_VERSION)).unwrap(), ), + ( + HeaderName::from_static("customized-header"), + HeaderValue::from_static("some/value"), + ), ]); assert_eq!(headers, expected_headers); } From 9ab76786ce65c73f6ee691a5542e57a54bdc420e Mon Sep 17 00:00:00 2001 From: Hongyi Wang Date: Wed, 27 Mar 2024 10:00:37 -0700 Subject: [PATCH 3/3] Rename function --- crates/catalog/rest/src/catalog.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 90577886f7..85996f8047 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -104,7 +104,7 @@ impl RestCatalogConfig { ]) } - fn get_default_headers(&self) -> Result { + fn http_headers(&self) -> Result { let mut headers = HeaderMap::from_iter([ ( header::CONTENT_TYPE, @@ -161,7 +161,7 @@ impl RestCatalogConfig { fn try_create_rest_client(&self) -> Result { // TODO: We will add ssl config, sigv4 later - let headers = self.get_default_headers()?; + let headers = self.http_headers()?; Ok(HttpClient( Client::builder().default_headers(headers).build()?, @@ -994,7 +994,7 @@ mod tests { } #[tokio::test] - async fn test_get_default_headers() { + async fn test_http_headers() { let server = Server::new_async().await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1003,7 +1003,7 @@ mod tests { .uri(server.url()) .props(props) .build(); - let headers: HeaderMap = config.get_default_headers().unwrap(); + let headers: HeaderMap = config.http_headers().unwrap(); let expected_headers = HeaderMap::from_iter([ ( @@ -1023,7 +1023,7 @@ mod tests { } #[tokio::test] - async fn test_get_default_headers_with_custom_headers() { + async fn test_http_headers_with_custom_headers() { let server = Server::new_async().await; let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); @@ -1040,7 +1040,7 @@ mod tests { .uri(server.url()) .props(props) .build(); - let headers: HeaderMap = config.get_default_headers().unwrap(); + let headers: HeaderMap = config.http_headers().unwrap(); let expected_headers = HeaderMap::from_iter([ (