[WIP] implement multi range query in single request

Cargo.toml

@@ -60,6 +60,7 @@ serde = { version = "1.0", default-features = false, features = ["derive"], opti
 serde_json = { version = "1.0", default-features = false, features = ["std"], optional = true }
 serde_urlencoded = { version = "0.7", optional = true }
 tokio = { version = "1.29.0", features = ["sync", "macros", "rt", "time", "io-util"] }
+mpart-async = "0.7.0"
 
 [target.'cfg(target_family="unix")'.dev-dependencies]
 nix = { version = "0.29.0", features = ["fs"] }
@@ -82,6 +83,7 @@ integration = ["rand"]
 [dev-dependencies] # In alphabetical order
 hyper = { version = "1.2", features = ["server"] }
 hyper-util = "0.1"
+mockito = "1.7"
 rand = "0.9"
 tempfile = "3.1.0"
 regex = "1.11.1"

src/aws/client.rs

@@ -35,6 +35,7 @@ use crate::client::s3::{
 use crate::client::{GetOptionsExt, HttpClient, HttpError, HttpResponse};
 use crate::multipart::PartId;
 use crate::path::DELIMITER;
+use crate::util::RangeValue;
 use crate::{
     Attribute, Attributes, ClientOptions, GetOptions, ListResult, MultipartId, Path,
     PutMultipartOpts, PutPayload, PutResult, Result, RetryConfig, TagSet,
@@ -838,14 +839,28 @@ impl GetClient for S3Client {
     };
 
     /// Make an S3 GET request <https://docs.aws.amazon.com/AmazonS3/latest/API/API_GetObject.html>
-    async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse> {
+    async fn get_request<R: RangeValue>(
+        &self,
+        path: &Path,
+        options: GetOptions<R>,
+    ) -> Result<HttpResponse> {
         let credential = self.config.get_session_credential().await?;
         let url = self.config.path_url(path);
         let method = match options.head {
             true => Method::HEAD,
             false => Method::GET,
         };
 
+        match options.range.as_ref() {
+            Some(inner) if inner.as_single().is_none() => {
+                // https://docs.aws.amazon.com/AmazonS3/latest/API/API_GetObject.html#API_GetObject_RequestSyntax
+                return Err(crate::Error::NotSupported {
+                    source: "AWS does not support multiple range requests".into(),
+                });
+            }
+            _ => {}
+        }
+
         let mut builder = self.client.request(method, url);
         if self
             .config

src/aws/mod.rs

@@ -500,6 +500,7 @@ mod tests {
     use crate::integration::*;
     use crate::tests::*;
     use crate::ClientOptions;
+    use crate::GetRange;
     use base64::prelude::BASE64_STANDARD;
     use base64::Engine;
     use http::HeaderMap;
@@ -743,7 +744,7 @@ mod tests {
         for location in &locations {
             let res = store
                 .client
-                .get_request(location, GetOptions::default())
+                .get_request::<GetRange>(location, GetOptions::default())
                 .await
                 .unwrap();
             let headers = res.headers();
@@ -799,7 +800,7 @@ mod tests {
         for location in &locations {
             let res = store
                 .client
-                .get_request(location, GetOptions::default())
+                .get_request::<GetRange>(location, GetOptions::default())
                 .await
                 .unwrap();
             let headers = res.headers();

src/azure/client.rs

@@ -26,7 +26,7 @@ use crate::client::retry::RetryExt;
 use crate::client::{GetOptionsExt, HttpClient, HttpError, HttpRequest, HttpResponse};
 use crate::multipart::PartId;
 use crate::path::DELIMITER;
-use crate::util::{deserialize_rfc1123, GetRange};
+use crate::util::{deserialize_rfc1123, GetRange, RangeValue};
 use crate::{
     Attribute, Attributes, ClientOptions, GetOptions, ListResult, ObjectMeta, Path, PutMode,
     PutMultipartOpts, PutOptions, PutPayload, PutResult, Result, RetryConfig, TagSet,
@@ -899,13 +899,27 @@ impl GetClient for AzureClient {
     /// Make an Azure GET request
     /// <https://docs.microsoft.com/en-us/rest/api/storageservices/get-blob>
     /// <https://docs.microsoft.com/en-us/rest/api/storageservices/get-blob-properties>
-    async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse> {
+    async fn get_request<R: RangeValue>(
+        &self,
+        path: &Path,
+        options: GetOptions<R>,
+    ) -> Result<HttpResponse> {
         // As of 2024-01-02, Azure does not support suffix requests,
         // so we should fail fast here rather than sending one
-        if let Some(GetRange::Suffix(_)) = options.range.as_ref() {
-            return Err(crate::Error::NotSupported {
-                source: "Azure does not support suffix range requests".into(),
-            });
+        if let Some(range) = options.range.as_ref() {
+            match range.as_single() {
+                Some(&GetRange::Suffix(_)) => {
+                    return Err(crate::Error::NotSupported {
+                        source: "Azure does not support suffix range requests".into(),
+                    });
+                }
+                None => {
+                    return Err(crate::Error::NotSupported {
+                        source: "Azure does not support multiple range requests".into(),
+                    });
+                }
+                _ => {}
+            }
         }
 
         let credential = self.get_credential().await?;

src/client/get.rs

@@ -20,14 +20,17 @@ use std::ops::Range;
 use crate::client::header::{header_meta, HeaderConfig};
 use crate::client::HttpResponse;
 use crate::path::Path;
+use crate::util::{GetManyRanges, RangeValue};
 use crate::{Attribute, Attributes, GetOptions, GetRange, GetResult, GetResultPayload, Result};
 use async_trait::async_trait;
+use bytes::{Bytes, BytesMut};
 use futures::{StreamExt, TryStreamExt};
 use http::header::{
     CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_RANGE,
     CONTENT_TYPE,
 };
 use http::StatusCode;
+use mpart_async::server::MultipartStream;
 use reqwest::header::ToStrError;
 
 /// A client that can perform a get request
@@ -38,13 +41,22 @@ pub(crate) trait GetClient: Send + Sync + 'static {
     /// Configure the [`HeaderConfig`] for this client
     const HEADER_CONFIG: HeaderConfig;
 
-    async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse>;
+    async fn get_request<R: RangeValue>(
+        &self,
+        path: &Path,
+        options: GetOptions<R>,
+    ) -> Result<HttpResponse>;
 }
 
 /// Extension trait for [`GetClient`] that adds common retrieval functionality
 #[async_trait]
 pub(crate) trait GetClientExt {
-    async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult>;
+    async fn get_opts(&self, location: &Path, options: GetOptions<GetRange>) -> Result<GetResult>;
+    async fn get_ranges_opts(
+        &self,
+        location: &Path,
+        options: GetOptions<GetManyRanges>,
+    ) -> Result<Vec<Bytes>>;
 }
 
 #[async_trait]
@@ -63,6 +75,29 @@ impl<T: GetClient> GetClientExt for T {
             source: Box::new(e),
         })
     }
+
+    async fn get_ranges_opts(
+        &self,
+        location: &Path,
+        options: GetOptions<GetManyRanges>,
+    ) -> Result<Vec<Bytes>> {
+        let range = options.range.clone();
+        if let Some(ranges) = range.as_ref() {
+            for r in ranges.as_ref() {
+                r.is_valid().map_err(|e| crate::Error::Generic {
+                    store: T::STORE,
+                    source: Box::new(e),
+                })?;
+            }
+        }
+        let response = self.get_request(location, options).await?;
+        get_manyranges::<T>(location, range.unwrap_or_default(), response)
+            .await
+            .map_err(|e| crate::Error::Generic {
+                store: T::STORE,
+                source: Box::new(e),
+            })
+    }
 }
 
 struct ContentRange {
@@ -143,6 +178,15 @@ enum GetResultError {
         expected: Range<u64>,
         actual: Range<u64>,
     },
+
+    #[error("Missing Content-Type header")]
+    MissingContentType,
+
+    #[error("No boundary found in Content-Type header, received {value:?}")]
+    NoBoundaryFound { value: String },
+
+    #[error(transparent)]
+    HttpError(#[from] super::HttpError),
 }
 
 fn get_result<T: GetClient>(
@@ -257,6 +301,125 @@ fn get_result<T: GetClient>(
     })
 }
 
+async fn get_manyranges<T: GetClient>(
+    _location: &Path,
+    ranges: GetManyRanges,
+    response: HttpResponse,
+) -> Result<Vec<bytes::Bytes>, GetResultError> {
+    let ranges = ranges.into_inner();
+    if ranges.is_empty() {
+        return Ok(Vec::new());
+    }
+
+    // ensure that we receive the range we asked for
+    if response.status() != StatusCode::PARTIAL_CONTENT {
+        return Err(GetResultError::NotPartial);
+    }
+
+    macro_rules! parse_attributes {
+        ($headers:expr, $(($header:expr, $attr:expr, $map_err:expr)),*) => {{
+            let mut attributes = Attributes::new();
+            $(
+            if let Some(x) = $headers.get($header) {
+                let x = x.to_str().map_err($map_err)?;
+                attributes.insert($attr, x.to_string().into());
+            }
+            )*
+            attributes
+        }}
+    }
+
+    let mut attributes = parse_attributes!(
+        response.headers(),
+        (CACHE_CONTROL, Attribute::CacheControl, |source| {
+            GetResultError::InvalidCacheControl { source }
+        }),
+        (
+            CONTENT_DISPOSITION,
+            Attribute::ContentDisposition,
+            |source| GetResultError::InvalidContentDisposition { source }
+        ),
+        (CONTENT_ENCODING, Attribute::ContentEncoding, |source| {
+            GetResultError::InvalidContentEncoding { source }
+        }),
+        (CONTENT_LANGUAGE, Attribute::ContentLanguage, |source| {
+            GetResultError::InvalidContentLanguage { source }
+        }),
+        (CONTENT_TYPE, Attribute::ContentType, |source| {
+            GetResultError::InvalidContentType { source }
+        })
+    );
+
+    // Add attributes that match the user-defined metadata prefix (e.g. x-amz-meta-)
+    if let Some(prefix) = T::HEADER_CONFIG.user_defined_metadata_prefix {
+        for (key, val) in response.headers() {
+            if let Some(suffix) = key.as_str().strip_prefix(prefix) {
+                if let Ok(val_str) = val.to_str() {
+                    attributes.insert(
+                        Attribute::Metadata(suffix.to_string().into()),
+                        val_str.to_string().into(),
+                    );
+                } else {
+                    return Err(GetResultError::InvalidMetadata {
+                        key: key.to_string(),
+                    });
+                }
+            }
+        }
+    }
+
+    let content_type = response.headers().get(CONTENT_TYPE);
+    let is_multiranges = content_type
+        .and_then(|v| v.to_str().ok())
+        .map(|v| v.starts_with("multipart/byteranges"))
+        .unwrap_or(false);
+
+    if ranges.len() == 1 && !is_multiranges {
+        let body = response.into_body().bytes().await?;
+        return Ok(vec![body]);
+    }
+
+    parse_multipart_byteranges(response).await
+}
+
+async fn parse_multipart_byteranges(response: HttpResponse) -> Result<Vec<Bytes>, GetResultError> {
+    let content_type = response
+        .headers()
+        .get(CONTENT_TYPE)
+        .ok_or(GetResultError::MissingContentType)?
+        .to_str()
+        .map_err(|source| GetResultError::InvalidContentType { source })?;
+
+    // extract boundary from content-type
+    let boundary = content_type
+        .split(";")
+        .find_map(|part| {
+            let trimmed = part.trim();
+            trimmed
+                .strip_prefix("boundary=")
+                .map(|b| b.trim_matches('"'))
+        })
+        .map(String::from)
+        .ok_or_else(|| GetResultError::NoBoundaryFound {
+            value: content_type.to_owned(),
+        })?;
+
+    let body = response.into_body().bytes_stream();
+
+    let mut parts = Vec::new();
+    let mut stream = MultipartStream::new(boundary, body);
+
+    while let Ok(Some(mut chunk)) = stream.try_next().await {
+        let mut buffer = BytesMut::new();
+        while let Ok(Some(part)) = chunk.try_next().await {
+            buffer.extend_from_slice(&part);
+        }
+        parts.push(buffer.freeze());
+    }
+
+    Ok(parts)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -275,7 +438,11 @@ mod tests {
             user_defined_metadata_prefix: Some("x-test-meta-"),
         };
 
-        async fn get_request(&self, _: &Path, _: GetOptions) -> Result<HttpResponse> {
+        async fn get_request<R: ToString + Send>(
+            &self,
+            _: &Path,
+            _: GetOptions<R>,
+        ) -> Result<HttpResponse> {
             unimplemented!()
         }
     }