Runtime update of total blob size using content_range header (#253)

* Runtime update of total blob size using content_range header

* typos

* simplified code

Author: MindFlavor

sdk/core/src/errors.rs

@@ -18,8 +18,10 @@ pub enum ParsingError {
 
 #[derive(Debug, Clone, PartialEq, thiserror::Error)]
 pub enum ParseError {
-    #[error("Split not found")]
-    SplitNotFound,
+    #[error("Expected token \"{}\" not found", 0)]
+    TokenNotFound(String),
+    #[error("Expected split char \'{}\' not found", 0)]
+    SplitNotFound(char),
     #[error("Parse int error {}", 0)]
     ParseIntError(ParseIntError),
 }

sdk/core/src/request_options/content_range.rs

@@ -0,0 +1,127 @@
+use crate::errors::ParseError;
+use std::fmt;
+use std::str::FromStr;
+
+const PREFIX: &str = "bytes ";
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct ContentRange {
+    start: u64,
+    end: u64,
+    total_length: u64,
+}
+
+impl ContentRange {
+    pub fn new(start: u64, end: u64, total_length: u64) -> ContentRange {
+        ContentRange {
+            start,
+            end,
+            total_length,
+        }
+    }
+
+    pub fn start(&self) -> u64 {
+        self.start
+    }
+
+    pub fn end(&self) -> u64 {
+        self.end
+    }
+
+    pub fn total_length(&self) -> u64 {
+        self.total_length
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.end == self.start
+    }
+}
+
+impl FromStr for ContentRange {
+    type Err = ParseError;
+    fn from_str(s: &str) -> Result<ContentRange, ParseError> {
+        let remaining = s
+            .strip_prefix(PREFIX)
+            .ok_or(ParseError::TokenNotFound(PREFIX.to_owned()))?;
+
+        let mut split_at_dash = remaining.split('-');
+        let start = split_at_dash.next().unwrap().parse()?;
+
+        let mut split_at_slash = split_at_dash
+            .next()
+            .ok_or(ParseError::SplitNotFound('-'))?
+            .split('/');
+
+        let end = split_at_slash.next().unwrap().parse()?;
+        let total_length = split_at_slash
+            .next()
+            .ok_or(ParseError::SplitNotFound('/'))?
+            .parse()?;
+
+        Ok(ContentRange {
+            start,
+            end,
+            total_length,
+        })
+    }
+}
+
+impl fmt::Display for ContentRange {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}{}-{}/{}",
+            PREFIX,
+            self.start(),
+            self.end(),
+            self.total_length()
+        )
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_parse() {
+        let range = "bytes 172032-172489/172490"
+            .parse::<ContentRange>()
+            .unwrap();
+
+        assert_eq!(range.start(), 172032);
+        assert_eq!(range.end(), 172489);
+        assert_eq!(range.total_length(), 172490);
+    }
+
+    #[test]
+    #[should_panic(expected = "TokenNotFound(\"bytes \")")]
+    fn test_parse_no_starting_token() {
+        "something else".parse::<ContentRange>().unwrap();
+    }
+
+    #[test]
+    #[should_panic(expected = "SplitNotFound('-')")]
+    fn test_parse_no_dash() {
+        "bytes 100".parse::<ContentRange>().unwrap();
+    }
+
+    #[test]
+    #[should_panic(expected = "SplitNotFound('/')")]
+    fn test_parse_no_slash() {
+        "bytes 100-500".parse::<ContentRange>().unwrap();
+    }
+
+    #[test]
+    fn test_display() {
+        let range = ContentRange {
+            start: 100,
+            end: 501,
+            total_length: 5000,
+        };
+
+        let txt = format!("{}", range);
+
+        assert_eq!(txt, "bytes 100-501/5000");
+    }
+}

sdk/core/src/request_options/mod.rs

@@ -4,6 +4,7 @@ mod client_request_id;
 mod content_disposition;
 mod content_encoding;
 mod content_language;
+mod content_range;
 mod content_type;
 mod continuation;
 mod delimiter;
@@ -33,6 +34,7 @@ pub use client_request_id::ClientRequestId;
 pub use content_disposition::ContentDisposition;
 pub use content_encoding::ContentEncoding;
 pub use content_language::ContentLanguage;
+pub use content_range::ContentRange;
 pub use content_type::ContentType;
 pub use continuation::Continuation;
 pub use delimiter::Delimiter;

sdk/core/src/request_options/range.rs

@@ -74,7 +74,7 @@ impl FromStr for Range {
     fn from_str(s: &str) -> Result<Range, ParseError> {
         let v = s.split('/').collect::<Vec<&str>>();
         if v.len() != 2 {
-            return Err(ParseError::SplitNotFound);
+            return Err(ParseError::SplitNotFound('/'));
         }
 
         let cp_start = v[0].parse::<u64>()?;

sdk/storage/examples/blob_00.rs

@@ -38,21 +38,28 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
 
     trace!("Requesting blob");
 
+    // this is a single call that retrieves the first 1KB of the blob (or less if the blob is
+    // smaller). The range(...) call is optional.
     let response = blob_client
         .get()
-        .range(Range::new(0, 128000))
+        .range(Range::new(0, 1024))
         .execute()
         .await?;
 
     println!("{:#?}", response);
 
-    let mut stream = Box::pin(blob_client.get().stream(128));
+    let mut complete_response = Vec::new();
+
+    // this is how you stream a blob. You can specify the range(...) value as above if necessary.
+    // In this case we are retrieving the whole blob in 8KB chunks.
+    let mut stream = Box::pin(blob_client.get().stream(1024 * 8));
     while let Some(value) = stream.next().await {
-        println!("received {:?} bytes", value?.data.len());
+        let data = value?.data;
+        println!("received {:?} bytes", data.len());
+        complete_response.extend(&data as &[u8]);
     }
 
-    let s_content = String::from_utf8(response.data.to_vec())?;
-    println!("blob == {:?}", blob);
+    let s_content = String::from_utf8(complete_response)?;
     println!("s_content == {}", s_content);
 
     Ok(())

sdk/storage/src/blob/blob/requests/get_blob_builder.rs

@@ -65,6 +65,8 @@ impl<'a> GetBlobBuilder<'a> {
             None,
         )?;
 
+        debug!("request == {:#?}", request);
+
         let expected_status_code = if self.range.is_some() {
             http::StatusCode::PARTIAL_CONTENT
         } else {
@@ -91,15 +93,21 @@ impl<'a> GetBlobBuilder<'a> {
             End,
         }
 
-        let complete_range = Range::new(0, u64::MAX);
+        // this can either be the range requested by the caller or the complete file.
+        let requested_range = self.range.unwrap_or(Range::new(0, u64::MAX));
 
         futures::stream::unfold(States::Init, move |state| async move {
-            let remaining = match state {
-                States::Init => self.range.unwrap_or(complete_range),
+            let mut remaining = match state {
+                States::Init => requested_range,
                 States::Progress(range) => range,
                 States::End => return None,
             };
 
+            debug!(
+                "remaining.start == {}, chunk_size == {}, remaining.end == {}",
+                remaining.start, chunk_size, remaining.end
+            );
+
             let range = if remaining.start + chunk_size > remaining.end {
                 Range::new(remaining.start, remaining.end)
             } else {
@@ -113,41 +121,16 @@ impl<'a> GetBlobBuilder<'a> {
                 Err(err) => return Some((Err(err), States::End)),
             };
 
+            debug!("response.content_range == {:?}", response.content_range);
+
+            // now that we know what the remote blob size is, let's update the
+            // boundary. We do this only if it's smaller than the requested size because the could
+            // have specified a smaller range.
+            remaining.end =
+                std::cmp::min(requested_range.end, response.content_range.total_length());
+
             let next_state = if remaining.end > range.end {
-                if self.range.is_some() {
-                    States::Progress(Range::new(range.end, remaining.end))
-                } else {
-                    // if we are here it means the user have not specified a
-                    // range and we didn't get the whole blob in one passing.
-                    // We specified u64::MAX as the first range but now
-                    // we need to find the correct size to avoid requesting data
-                    // outside the valid range.
-                    debug!("content-range == {:?}", response.content_range);
-                    // this unwrap should always be safe since we did not
-                    // get the whole blob in the previous call.
-                    let content_range = response.content_range.clone().unwrap();
-                    let ridx = match content_range.find('/') {
-                        Some(ridx) => ridx,
-                        None => {
-                            return Some((
-                                Err("The returned content-range is invalid: / is not present"
-                                    .into()),
-                                States::End,
-                            ))
-                        }
-                    };
-                    let total =
-                            match str::parse(&content_range[ridx + 1..]) {
-                                Ok(total) => total,
-                                Err(_err) => return Some((
-                                    Err("The returned content-range is invalid: after / there is a non valid number"
-                                        .into()),
-                                    States::End,
-                                )),
-                            };
-
-                    States::Progress(Range::new(range.end, total))
-                }
+                States::Progress(Range::new(range.end, remaining.end))
             } else {
                 States::End
             };

sdk/storage/src/blob/blob/responses/get_blob_response.rs

@@ -1,18 +1,21 @@
 use crate::{blob::blob::Blob, AzureStorageError};
+use azure_core::errors::AzureError;
 use azure_core::headers::{date_from_headers, request_id_from_headers};
+use azure_core::prelude::ContentRange;
 use azure_core::RequestId;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use http::Response;
 use std::convert::TryFrom;
+use std::str::FromStr;
 
 #[derive(Debug, Clone)]
 pub struct GetBlobResponse {
     pub request_id: RequestId,
     pub blob: Blob,
     pub data: Bytes,
     pub date: DateTime<Utc>,
-    pub content_range: Option<String>,
+    pub content_range: ContentRange,
 }
 
 impl TryFrom<(&str, Response<Bytes>)> for GetBlobResponse {
@@ -23,10 +26,13 @@ impl TryFrom<(&str, Response<Bytes>)> for GetBlobResponse {
         let request_id = request_id_from_headers(response.headers())?;
         let date = date_from_headers(response.headers())?;
 
-        let content_range = response
-            .headers()
-            .get(http::header::CONTENT_RANGE)
-            .map(|h| h.to_str().unwrap().to_owned());
+        let content_range = ContentRange::from_str(
+            response
+                .headers()
+                .get(http::header::CONTENT_RANGE)
+                .ok_or_else(|| AzureError::HeaderNotFound(http::header::CONTENT_RANGE.to_string()))?
+                .to_str()?,
+        )?;
 
         Ok(GetBlobResponse {
             request_id,

sdk/storage/src/core/errors.rs

@@ -2,6 +2,8 @@
 pub enum AzureStorageError {
     #[error(transparent)]
     CoreError(#[from] azure_core::errors::AzureError),
+    #[error("Parse error: {}", 0)]
+    ParseError(#[from] azure_core::errors::ParseError),
     #[error("Parsing error: {}", 0)]
     ParsingError(azure_core::errors::ParsingError),
     #[error("Permission error: {}", 0)]

sdk/storage/tests/stream_blob00.rs

@@ -92,6 +92,10 @@ async fn code() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         assert_eq!(expected_string, returned_string);
     }
 
+    // test streaming a blob smaller than the chunk size issue 239.
+    let mut stream = Box::pin(blob.get().stream(1024 * 8));
+    while let Some(_value) = stream.next().await {}
+
     blob.delete()
         .delete_snapshots_method(DeleteSnapshotsMethod::Include)
         .execute()