Better header insertion handling (#690)

* Add a request.insert_header method * AddAsHeader for Option<T> where T: AddAsHeader * Two more usages * Another usage
Azure · Mar 18, 2022 · 102ce3f · 102ce3f
1 parent 0ed3e65
commit 102ce3f
Show file tree

Hide file tree

Showing 20 changed files with 97 additions and 117 deletions.
diff --git a/sdk/core/src/headers/mod.rs b/sdk/core/src/headers/mod.rs
@@ -25,6 +25,29 @@ pub trait AddAsHeader {
     ) -> Result<(), crate::errors::HttpHeaderError>;
 }
 
+impl<T> AddAsHeader for Option<T>
+where
+    T: AddAsHeader,
+{
+    fn add_as_header(&self, builder: Builder) -> Builder {
+        if let Some(h) = self {
+            h.add_as_header(builder)
+        } else {
+            builder
+        }
+    }
+
+    fn add_as_header2(
+        &self,
+        request: &mut crate::Request,
+    ) -> Result<(), crate::errors::HttpHeaderError> {
+        if let Some(h) = self {
+            h.add_as_header2(request)?;
+        }
+        Ok(())
+    }
+}
+
 #[must_use]
 pub fn add_optional_header_ref<T: AddAsHeader>(item: &Option<&T>, mut builder: Builder) -> Builder {
     if let Some(item) = item {

diff --git a/sdk/core/src/pageable.rs b/sdk/core/src/pageable.rs
@@ -2,6 +2,8 @@ use futures::stream::unfold;
 use futures::Stream;
 use pin_project::pin_project;
 
+use crate::prelude::Continuation;
+
 macro_rules! r#try {
     ($expr:expr $(,)?) => {
         match $expr {
@@ -24,7 +26,7 @@ pub struct Pageable<T, E> {
 }
 
 impl<T: Continuable, E> Pageable<T, E> {
-    pub fn new<F>(make_request: impl Fn(Option<String>) -> F + Clone + 'static) -> Self
+    pub fn new<F>(make_request: impl Fn(Option<Continuation>) -> F + Clone + 'static) -> Self
     where
         F: std::future::Future<Output = Result<T, E>> + 'static,
     {
@@ -34,7 +36,7 @@ impl<T: Continuable, E> Pageable<T, E> {
                 let response = match state {
                     State::Init => r#try!(make_request(None).await),
                     State::Continuation(token) => {
-                        r#try!(make_request(Some(token)).await)
+                        r#try!(make_request(Some(Continuation::new(token))).await)
                     }
                     State::Done => return None,
                 };

diff --git a/sdk/core/src/request.rs b/sdk/core/src/request.rs
@@ -1,4 +1,7 @@
-use crate::SeekableStream;
+use crate::{
+    error::{ErrorKind, ResultExt},
+    AddAsHeader, SeekableStream,
+};
 use http::{HeaderMap, Method, Uri};
 use std::fmt::Debug;
 
@@ -54,6 +57,14 @@ impl Request {
         self.method.clone()
     }
 
+    pub fn insert_header<T: AddAsHeader + Debug>(&mut self, h: &T) -> crate::error::Result<()> {
+        h.add_as_header2(self)
+            .with_context(ErrorKind::DataConversion, || {
+                format!("could not encode '{:?}' as an http header", h)
+            })?;
+        Ok(())
+    }
+
     pub fn headers(&self) -> &HeaderMap {
         &self.headers
     }

diff --git a/sdk/core/src/request_options/continuation.rs b/sdk/core/src/request_options/continuation.rs
@@ -1,18 +1,22 @@
 use crate::{headers, AddAsHeader};
 use http::request::Builder;
 
-#[derive(Debug, Clone, Copy)]
-pub struct Continuation<'a>(&'a str);
+#[derive(Debug, Clone)]
+pub struct Continuation(String);
 
-impl<'a> Continuation<'a> {
-    pub fn new(c: &'a str) -> Self {
+impl Continuation {
+    pub fn new(c: String) -> Self {
         Self(c)
     }
+
+    pub fn into_raw(self) -> String {
+        self.0
+    }
 }
 
-impl AddAsHeader for Continuation<'_> {
+impl AddAsHeader for Continuation {
     fn add_as_header(&self, builder: Builder) -> Builder {
-        builder.header(headers::CONTINUATION, self.0)
+        builder.header(headers::CONTINUATION, &self.0)
     }
 
     fn add_as_header2(
@@ -21,7 +25,7 @@ impl AddAsHeader for Continuation<'_> {
     ) -> Result<(), crate::errors::HttpHeaderError> {
         request
             .headers_mut()
-            .append(headers::CONTINUATION, http::HeaderValue::from_str(self.0)?);
+            .append(headers::CONTINUATION, http::HeaderValue::from_str(&self.0)?);
 
         Ok(())
     }

diff --git a/sdk/core/src/request_options/next_marker.rs b/sdk/core/src/request_options/next_marker.rs
@@ -1,5 +1,7 @@
 use crate::{AppendToUrlQuery, Error, HttpHeaderError};
 
+use super::Continuation;
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct NextMarker(String);
 
@@ -47,6 +49,12 @@ impl AppendToUrlQuery for NextMarker {
     }
 }
 
+impl From<Continuation> for NextMarker {
+    fn from(next_marker: Continuation) -> Self {
+        Self::new(next_marker.into_raw())
+    }
+}
+
 impl From<String> for NextMarker {
     fn from(next_marker: String) -> Self {
         Self::new(next_marker)

diff --git a/sdk/data_cosmos/src/operations/create_database.rs b/sdk/data_cosmos/src/operations/create_database.rs
@@ -45,13 +45,7 @@ impl CreateDatabaseBuilder {
                 id: self.database_name.as_str(),
             };
 
-            azure_core::headers::add_optional_header2(&self.consistency_level, &mut request)
-                .with_context(ErrorKind::DataConversion, || {
-                    format!(
-                        "could not encode '{:?}' as an http header",
-                        self.consistency_level
-                    )
-                })?;
+            request.insert_header(&self.consistency_level)?;
             request.set_body(bytes::Bytes::from(serde_json::to_string(&body)?).into());
 
             let response = self

diff --git a/sdk/data_cosmos/src/operations/list_attachments.rs b/sdk/data_cosmos/src/operations/list_attachments.rs
@@ -4,7 +4,6 @@ use crate::resources::Attachment;
 use crate::resources::ResourceType;
 use crate::ResourceQuota;
 use azure_core::collect_pinned_stream;
-use azure_core::headers;
 use azure_core::headers::{
     continuation_token_from_headers_optional, item_count_from_headers, session_token_from_headers,
 };
@@ -43,7 +42,7 @@ impl ListAttachmentsBuilder {
     }
 
     pub fn into_stream(self) -> ListAttachments {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
@@ -66,11 +65,7 @@ impl ListAttachmentsBuilder {
                     &mut request,
                 );
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client

diff --git a/sdk/data_cosmos/src/operations/list_collections.rs b/sdk/data_cosmos/src/operations/list_collections.rs
@@ -5,7 +5,7 @@ use crate::ResourceQuota;
 use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::prelude::*;
 use azure_core::Response as HttpResponse;
-use azure_core::{collect_pinned_stream, headers, Pageable};
+use azure_core::{collect_pinned_stream, Pageable};
 use chrono::{DateTime, Utc};
 
 #[derive(Debug, Clone)]
@@ -33,19 +33,15 @@ impl ListCollectionsBuilder {
     }
 
     pub fn into_stream(self) -> ListCollections {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
                 let mut request = this.client.prepare_collections_pipeline(http::Method::GET);
                 azure_core::headers::add_optional_header2(&this.consistency_level, &mut request)?;
                 azure_core::headers::add_mandatory_header2(&this.max_item_count, &mut request)?;
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client

diff --git a/sdk/data_cosmos/src/operations/list_databases.rs b/sdk/data_cosmos/src/operations/list_databases.rs
@@ -5,9 +5,7 @@ use crate::ResourceQuota;
 
 use azure_core::error::ErrorKind;
 use azure_core::error::ResultExt;
-use azure_core::headers::{
-    self, continuation_token_from_headers_optional, session_token_from_headers,
-};
+use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::{collect_pinned_stream, prelude::*, Pageable, Response};
 use chrono::{DateTime, Utc};
 
@@ -36,36 +34,16 @@ impl ListDatabasesBuilder {
     }
 
     pub fn into_stream(self) -> ListDatabases {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
                 let mut request = this
                     .client
                     .prepare_request_pipeline("dbs", http::Method::GET);
-
-                azure_core::headers::add_optional_header2(&this.consistency_level, &mut request)
-                    .with_context(ErrorKind::DataConversion, || {
-                        format!(
-                            "could not encode '{:?}' as an http header",
-                            this.consistency_level
-                        )
-                    })?;
-                azure_core::headers::add_mandatory_header2(&this.max_item_count, &mut request)
-                    .with_context(ErrorKind::DataConversion, || {
-                        format!(
-                            "could not encode '{:?}' as an http header",
-                            this.max_item_count
-                        )
-                    })?;
-
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .with_context(ErrorKind::DataConversion, || {
-                            format!("could not encode '{:?}' as an http header", c)
-                        })?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&this.consistency_level)?;
+                request.insert_header(&this.max_item_count)?;
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client

diff --git a/sdk/data_cosmos/src/operations/list_documents.rs b/sdk/data_cosmos/src/operations/list_documents.rs
@@ -4,8 +4,7 @@ use crate::resources::document::{Document, DocumentAttributes};
 use crate::resources::ResourceType;
 use crate::ResourceQuota;
 use azure_core::headers::{
-    self, continuation_token_from_headers_optional, item_count_from_headers,
-    session_token_from_headers,
+    continuation_token_from_headers_optional, item_count_from_headers, session_token_from_headers,
 };
 use azure_core::{collect_pinned_stream, Response, SessionToken};
 use azure_core::{prelude::*, Pageable};
@@ -45,7 +44,7 @@ impl ListDocumentsBuilder {
     }
 
     pub fn into_stream<T: DeserializeOwned>(self) -> ListDocuments<T> {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
@@ -64,10 +63,8 @@ impl ListDocumentsBuilder {
                 azure_core::headers::add_mandatory_header2(&this.a_im, &mut req)?;
                 azure_core::headers::add_optional_header2(&this.partition_range_id, &mut req)?;
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    req.headers_mut().append(headers::CONTINUATION, h);
+                if let Some(ref c) = continuation {
+                    req.insert_header(c)?;
                 }
 
                 let response = this

diff --git a/sdk/data_cosmos/src/operations/list_permissions.rs b/sdk/data_cosmos/src/operations/list_permissions.rs
@@ -3,9 +3,7 @@ use crate::prelude::*;
 use crate::resources::Permission;
 use crate::resources::ResourceType;
 use azure_core::collect_pinned_stream;
-use azure_core::headers::{
-    self, continuation_token_from_headers_optional, session_token_from_headers,
-};
+use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::prelude::*;
 use azure_core::{Pageable, Response as HttpResponse};
 
@@ -34,7 +32,7 @@ impl ListPermissionsBuilder {
     }
 
     pub fn into_stream(self) -> ListPermissions {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
@@ -50,11 +48,7 @@ impl ListPermissionsBuilder {
                 azure_core::headers::add_optional_header2(&this.consistency_level, &mut request)?;
                 azure_core::headers::add_mandatory_header2(&this.max_item_count, &mut request)?;
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client

diff --git a/sdk/data_cosmos/src/operations/list_stored_procedures.rs b/sdk/data_cosmos/src/operations/list_stored_procedures.rs
@@ -4,7 +4,6 @@ use crate::resources::ResourceType;
 use crate::resources::StoredProcedure;
 use crate::ResourceQuota;
 use azure_core::collect_pinned_stream;
-use azure_core::headers;
 use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::prelude::*;
 use azure_core::{Pageable, Response as HttpResponse};
@@ -35,7 +34,7 @@ impl ListStoredProceduresBuilder {
     }
 
     pub fn into_stream(self) -> ListStoredProcedures {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
@@ -51,11 +50,7 @@ impl ListStoredProceduresBuilder {
                 azure_core::headers::add_optional_header2(&this.consistency_level, &mut request)?;
                 azure_core::headers::add_mandatory_header2(&this.max_item_count, &mut request)?;
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client

diff --git a/sdk/data_cosmos/src/operations/list_triggers.rs b/sdk/data_cosmos/src/operations/list_triggers.rs
@@ -3,7 +3,6 @@ use crate::prelude::*;
 use crate::resources::ResourceType;
 use crate::ResourceQuota;
 use azure_core::collect_pinned_stream;
-use azure_core::headers;
 use azure_core::headers::item_count_from_headers;
 use azure_core::headers::{continuation_token_from_headers_optional, session_token_from_headers};
 use azure_core::prelude::*;
@@ -38,7 +37,7 @@ impl ListTriggersBuilder {
     }
 
     pub fn into_stream(self) -> ListTriggers {
-        let make_request = move |continuation: Option<String>| {
+        let make_request = move |continuation: Option<Continuation>| {
             let this = self.clone();
             let ctx = self.context.clone();
             async move {
@@ -55,11 +54,7 @@ impl ListTriggersBuilder {
                 azure_core::headers::add_optional_header2(&this.consistency_level, &mut request)?;
                 azure_core::headers::add_mandatory_header2(&this.max_item_count, &mut request)?;
 
-                if let Some(c) = continuation {
-                    let h = http::HeaderValue::from_str(c.as_str())
-                        .map_err(azure_core::HttpHeaderError::InvalidHeaderValue)?;
-                    request.headers_mut().append(headers::CONTINUATION, h);
-                }
+                request.insert_header(&continuation)?;
 
                 let response = this
                     .client