From 983dd70145665a176a0b6a8909544a02fd2135b8 Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Thu, 4 Jan 2024 16:11:18 +0000 Subject: [PATCH] Fix Azure and Memory --- object_store/src/lib.rs | 30 ++++++++++++++------- object_store/src/memory.rs | 53 ++++++++++---------------------------- object_store/src/util.rs | 2 +- 3 files changed, 34 insertions(+), 51 deletions(-) diff --git a/object_store/src/lib.rs b/object_store/src/lib.rs index d32aa493269f..dbd6cce4c4c9 100644 --- a/object_store/src/lib.rs +++ b/object_store/src/lib.rs @@ -1340,21 +1340,31 @@ mod tests { range: Some(GetRange::Suffix(2)), ..Default::default() }; - let result = storage.get_opts(&location, opts).await.unwrap(); - assert_eq!(result.range, 12..14); - assert_eq!(result.meta.size, 14); - let bytes = result.bytes().await.unwrap(); - assert_eq!(bytes, b"ta".as_ref()); + match storage.get_opts(&location, opts).await { + Ok(result) => { + assert_eq!(result.range, 12..14); + assert_eq!(result.meta.size, 14); + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes, b"ta".as_ref()); + } + Err(Error::NotImplemented { .. }) => {} + Err(e) => panic!("{e}"), + } let opts = GetOptions { range: Some(GetRange::Suffix(100)), ..Default::default() }; - let result = storage.get_opts(&location, opts).await.unwrap(); - assert_eq!(result.range, 0..14); - assert_eq!(result.meta.size, 14); - let bytes = result.bytes().await.unwrap(); - assert_eq!(bytes, b"arbitrary data".as_ref()); + match storage.get_opts(&location, opts).await { + Ok(result) => { + assert_eq!(result.range, 0..14); + assert_eq!(result.meta.size, 14); + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes, b"arbitrary data".as_ref()); + } + Err(Error::NotImplemented { .. }) => {} + Err(e) => panic!("{e}"), + } let opts = GetOptions { range: Some(GetRange::Offset(3)), diff --git a/object_store/src/memory.rs b/object_store/src/memory.rs index 703e052d2883..41cfcc490da6 100644 --- a/object_store/src/memory.rs +++ b/object_store/src/memory.rs @@ -16,9 +16,10 @@ // under the License. //! An in-memory object store implementation +use crate::util::InvalidGetRange; use crate::{ - path::Path, GetResult, GetResultPayload, ListResult, ObjectMeta, ObjectStore, PutMode, - PutOptions, PutResult, Result, UpdateVersion, + path::Path, GetRange, GetResult, GetResultPayload, ListResult, ObjectMeta, ObjectStore, + PutMode, PutOptions, PutResult, Result, UpdateVersion, }; use crate::{GetOptions, MultipartId}; use async_trait::async_trait; @@ -26,7 +27,7 @@ use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::{stream::BoxStream, StreamExt}; use parking_lot::RwLock; -use snafu::{ensure, OptionExt, Snafu}; +use snafu::{OptionExt, ResultExt, Snafu}; use std::collections::BTreeMap; use std::collections::BTreeSet; use std::io; @@ -43,16 +44,8 @@ enum Error { #[snafu(display("No data in memory found. Location: {path}"))] NoDataInMemory { path: String }, - #[snafu(display( - "Requested range {}..{} is out of bounds for object with length {}", range.start, range.end, len - ))] - OutOfRange { range: Range, len: usize }, - - #[snafu(display("Invalid range: {}..{}", range.start, range.end))] - BadRange { range: Range }, - - #[snafu(display("Invalid suffix: {} bytes", nbytes))] - BadSuffix { nbytes: usize }, + #[snafu(display("Invalid range: {source}"))] + Range { source: InvalidGetRange }, #[snafu(display("Object already exists at that location: {path}"))] AlreadyExists { path: String }, @@ -209,8 +202,6 @@ impl ObjectStore for InMemory { } async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { - use crate::util::GetRange::*; - let entry = self.entry(location).await?; let e_tag = entry.e_tag.to_string(); @@ -225,23 +216,8 @@ impl ObjectStore for InMemory { let (range, data) = match options.range { Some(range) => { - let len = entry.data.len(); - match range { - Bounded(r) => { - ensure!(r.end <= len, OutOfRangeSnafu { range: r, len }); - ensure!(r.start <= r.end, BadRangeSnafu { range: r }); - (r.clone(), entry.data.slice(r)) - } - Offset(o) => { - ensure!(o < len, OutOfRangeSnafu { range: o..len, len }); - (o..len, entry.data.slice(o..len)) - } - Suffix(n) => { - ensure!(n < len, BadSuffixSnafu { nbytes: n }); - let start = len - n; - (start..len, entry.data.slice(start..len)) - } - } + let r = range.as_range(entry.data.len()).context(RangeSnafu)?; + (r.clone(), entry.data.slice(r)) } None => (0..entry.data.len(), entry.data), }; @@ -259,14 +235,11 @@ impl ObjectStore for InMemory { ranges .iter() .map(|range| { - let range = range.clone(); - let len = entry.data.len(); - ensure!( - range.end <= entry.data.len(), - OutOfRangeSnafu { range, len } - ); - ensure!(range.start <= range.end, BadRangeSnafu { range }); - Ok(entry.data.slice(range)) + let r = GetRange::Bounded(range.clone()) + .as_range(entry.data.len()) + .context(RangeSnafu)?; + + Ok(entry.data.slice(r)) }) .collect() } diff --git a/object_store/src/util.rs b/object_store/src/util.rs index 156a476cc1e0..6c6e53723c61 100644 --- a/object_store/src/util.rs +++ b/object_store/src/util.rs @@ -207,7 +207,7 @@ pub enum GetRange { #[derive(Debug, Snafu)] pub(crate) enum InvalidGetRange { #[snafu(display( - "Wanted range starting at {requested}, but resource was only {length} bytes long" + "Wanted range starting at {requested}, but object was only {length} bytes long" ))] StartTooLarge { requested: usize, length: usize },