Skip to main content

object_store/
prefix.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! An object store wrapper handling a constant path prefix
19use bytes::Bytes;
20use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
21use std::ops::Range;
22
23use crate::multipart::{MultipartStore, PartId};
24use crate::path::Path;
25#[cfg(feature = "cloud-base")]
26use crate::signer::Signer;
27use crate::{
28    CopyOptions, GetOptions, GetResult, ListResult, MultipartId, MultipartUpload, ObjectMeta,
29    ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result,
30};
31
32/// Store wrapper that applies a constant prefix to all paths handled by the store.
33#[derive(Debug, Clone)]
34pub struct PrefixStore<T> {
35    prefix: Path,
36    inner: T,
37}
38
39impl<T> std::fmt::Display for PrefixStore<T> {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        write!(f, "PrefixObjectStore({})", self.prefix.as_ref())
42    }
43}
44
45impl<T> PrefixStore<T> {
46    /// Create a new instance of [`PrefixStore`]
47    pub fn new(store: T, prefix: impl Into<Path>) -> Self {
48        Self {
49            prefix: prefix.into(),
50            inner: store,
51        }
52    }
53
54    /// Create the full path from a path relative to prefix
55    fn full_path(&self, location: &Path) -> Path {
56        full_path(&self.prefix, location)
57    }
58
59    /// Strip the constant prefix from a given path
60    fn strip_prefix(&self, path: Path) -> Path {
61        strip_prefix(&self.prefix, path)
62    }
63
64    /// Strip the constant prefix from a given ObjectMeta
65    fn strip_meta(&self, meta: ObjectMeta) -> ObjectMeta {
66        strip_meta(&self.prefix, meta)
67    }
68}
69
70// Note: This is a relative hack to move these functions to pure functions so they don't rely
71// on the `self` lifetime.
72
73/// Create the full path from a path relative to prefix
74fn full_path(prefix: &Path, path: &Path) -> Path {
75    prefix.parts().chain(path.parts()).collect()
76}
77
78/// Strip the constant prefix from a given path
79fn strip_prefix(prefix: &Path, path: Path) -> Path {
80    // Note cannot use match because of borrow checker
81    if let Some(suffix) = path.prefix_match(prefix) {
82        return suffix.collect();
83    }
84    path
85}
86
87/// Strip the constant prefix from a given ObjectMeta
88fn strip_meta(prefix: &Path, meta: ObjectMeta) -> ObjectMeta {
89    let ObjectMeta {
90        last_modified,
91        size,
92        location,
93        e_tag,
94        version,
95    } = meta;
96    ObjectMeta {
97        last_modified,
98        size,
99        location: strip_prefix(prefix, location),
100        e_tag,
101        version,
102    }
103}
104
105#[async_trait::async_trait]
106#[deny(clippy::missing_trait_methods)]
107impl<T: ObjectStore> ObjectStore for PrefixStore<T> {
108    async fn put_opts(
109        &self,
110        location: &Path,
111        payload: PutPayload,
112        opts: PutOptions,
113    ) -> Result<PutResult> {
114        let full_path = self.full_path(location);
115        self.inner.put_opts(&full_path, payload, opts).await
116    }
117
118    async fn put_multipart_opts(
119        &self,
120        location: &Path,
121        opts: PutMultipartOptions,
122    ) -> Result<Box<dyn MultipartUpload>> {
123        let full_path = self.full_path(location);
124        self.inner.put_multipart_opts(&full_path, opts).await
125    }
126
127    async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
128        let full_path = self.full_path(location);
129        let mut result = self.inner.get_opts(&full_path, options).await?;
130        result.meta = self.strip_meta(result.meta);
131        Ok(result)
132    }
133
134    async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
135        let full_path = self.full_path(location);
136        self.inner.get_ranges(&full_path, ranges).await
137    }
138
139    fn delete_stream(
140        &self,
141        locations: BoxStream<'static, Result<Path>>,
142    ) -> BoxStream<'static, Result<Path>> {
143        let prefix = self.prefix.clone();
144        let locations = locations
145            .map(move |location| location.map(|loc| full_path(&prefix, &loc)))
146            .boxed();
147        let prefix = self.prefix.clone();
148        self.inner
149            .delete_stream(locations)
150            .map(move |location| location.map(|loc| strip_prefix(&prefix, loc)))
151            .boxed()
152    }
153
154    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
155        let prefix = self.full_path(prefix.unwrap_or(&Path::default()));
156        let s = self.inner.list(Some(&prefix));
157        let slf_prefix = self.prefix.clone();
158        s.map_ok(move |meta| strip_meta(&slf_prefix, meta)).boxed()
159    }
160
161    fn list_with_offset(
162        &self,
163        prefix: Option<&Path>,
164        offset: &Path,
165    ) -> BoxStream<'static, Result<ObjectMeta>> {
166        let offset = self.full_path(offset);
167        let prefix = self.full_path(prefix.unwrap_or(&Path::default()));
168        let s = self.inner.list_with_offset(Some(&prefix), &offset);
169        let slf_prefix = self.prefix.clone();
170        s.map_ok(move |meta| strip_meta(&slf_prefix, meta)).boxed()
171    }
172
173    async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
174        let prefix = self.full_path(prefix.unwrap_or(&Path::default()));
175        self.inner
176            .list_with_delimiter(Some(&prefix))
177            .await
178            .map(|lst| ListResult {
179                common_prefixes: lst
180                    .common_prefixes
181                    .into_iter()
182                    .map(|p| self.strip_prefix(p))
183                    .collect(),
184                objects: lst
185                    .objects
186                    .into_iter()
187                    .map(|meta| self.strip_meta(meta))
188                    .collect(),
189                extensions: lst.extensions,
190            })
191    }
192
193    async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> {
194        let full_from = self.full_path(from);
195        let full_to = self.full_path(to);
196        self.inner.copy_opts(&full_from, &full_to, options).await
197    }
198
199    async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> Result<()> {
200        let full_from = self.full_path(from);
201        let full_to = self.full_path(to);
202        self.inner.rename_opts(&full_from, &full_to, options).await
203    }
204}
205
206#[async_trait::async_trait]
207impl<T: MultipartStore> MultipartStore for PrefixStore<T> {
208    async fn create_multipart(&self, path: &Path) -> Result<MultipartId> {
209        let full_path = self.full_path(path);
210        self.inner.create_multipart(&full_path).await
211    }
212
213    async fn create_multipart_opts(
214        &self,
215        path: &Path,
216        opts: PutMultipartOptions,
217    ) -> Result<MultipartId> {
218        let full_path = self.full_path(path);
219        self.inner.create_multipart_opts(&full_path, opts).await
220    }
221
222    async fn put_part(
223        &self,
224        path: &Path,
225        id: &MultipartId,
226        part_idx: usize,
227        data: PutPayload,
228    ) -> Result<PartId> {
229        let full_path = self.full_path(path);
230        self.inner.put_part(&full_path, id, part_idx, data).await
231    }
232
233    async fn complete_multipart(
234        &self,
235        path: &Path,
236        id: &MultipartId,
237        parts: Vec<PartId>,
238    ) -> Result<PutResult> {
239        let full_path = self.full_path(path);
240        self.inner.complete_multipart(&full_path, id, parts).await
241    }
242
243    async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()> {
244        let full_path = self.full_path(path);
245        self.inner.abort_multipart(&full_path, id).await
246    }
247}
248
249#[cfg(feature = "cloud-base")]
250#[async_trait::async_trait]
251impl<T: Signer> Signer for PrefixStore<T> {
252    async fn signed_url(
253        &self,
254        method: http::Method,
255        path: &Path,
256        expires_in: std::time::Duration,
257    ) -> Result<url::Url> {
258        self.inner
259            .signed_url(method, &self.full_path(path), expires_in)
260            .await
261    }
262
263    async fn signed_url_opts(
264        &self,
265        method: http::Method,
266        path: &Path,
267        expires_in: std::time::Duration,
268        options: &crate::signer::SignedUrlOptions,
269    ) -> Result<url::Url> {
270        self.inner
271            .signed_url_opts(method, &self.full_path(path), expires_in, options)
272            .await
273    }
274
275    async fn signed_urls(
276        &self,
277        method: http::Method,
278        paths: &[Path],
279        expires_in: std::time::Duration,
280    ) -> Result<Vec<url::Url>> {
281        self.inner
282            .signed_urls(
283                method,
284                &paths.iter().map(|p| self.full_path(p)).collect::<Vec<_>>(),
285                expires_in,
286            )
287            .await
288    }
289}
290
291#[cfg(not(target_arch = "wasm32"))]
292#[cfg(test)]
293mod tests {
294    use std::slice;
295
296    use super::*;
297    use crate::local::LocalFileSystem;
298    use crate::memory::InMemory;
299    use crate::{ObjectStoreExt, integration::*};
300
301    use tempfile::TempDir;
302
303    #[tokio::test]
304    async fn prefix_test() {
305        let root = TempDir::new().unwrap();
306        let inner = LocalFileSystem::new_with_prefix(root.path()).unwrap();
307        let integration = PrefixStore::new(inner, "prefix");
308
309        put_get_delete_list(&integration).await;
310        get_opts(&integration).await;
311        list_uses_directories_correctly(&integration).await;
312        list_with_delimiter(&integration).await;
313        rename_and_copy(&integration).await;
314        copy_if_not_exists(&integration).await;
315        stream_get(&integration).await;
316    }
317
318    #[tokio::test]
319    async fn prefix_test_applies_prefix() {
320        let tmpdir = TempDir::new().unwrap();
321        let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
322
323        let location = Path::from("prefix/test_file.json");
324        let data = Bytes::from("arbitrary data");
325
326        local.put(&location, data.clone().into()).await.unwrap();
327
328        let prefix = PrefixStore::new(local, "prefix");
329        let location_prefix = Path::from("test_file.json");
330
331        let content_list = flatten_list_stream(&prefix, None).await.unwrap();
332        assert_eq!(content_list, slice::from_ref(&location_prefix));
333
334        let root = Path::from("/");
335        let content_list = flatten_list_stream(&prefix, Some(&root)).await.unwrap();
336        assert_eq!(content_list, slice::from_ref(&location_prefix));
337
338        let read_data = prefix
339            .get(&location_prefix)
340            .await
341            .unwrap()
342            .bytes()
343            .await
344            .unwrap();
345        assert_eq!(&*read_data, data);
346
347        let target_prefix = Path::from("/test_written.json");
348        prefix
349            .put(&target_prefix, data.clone().into())
350            .await
351            .unwrap();
352
353        prefix.delete(&location_prefix).await.unwrap();
354
355        let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
356
357        let err = local.get(&location).await.unwrap_err();
358        assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
359
360        let location = Path::from("prefix/test_written.json");
361        let read_data = local.get(&location).await.unwrap().bytes().await.unwrap();
362        assert_eq!(&*read_data, data)
363    }
364
365    // Regression test for
366    // https://github.com/apache/arrow-rs-object-store/issues/664:
367    // head/get returned an ObjectMeta whose location still contained the
368    // prefix, so round-tripping the returned location back into the store would
369    // double-prefix it.
370    #[tokio::test]
371    async fn prefix_head_get_strip_prefix() {
372        let store = PrefixStore::new(InMemory::new(), "prefix");
373        let path = Path::from("test_file");
374        store.put(&path, "data".into()).await.unwrap();
375
376        let head = store.head(&path).await.unwrap();
377        assert_eq!(head.location, path);
378        let get = store.get(&path).await.unwrap();
379        assert_eq!(get.meta.location, path);
380
381        // The returned location must round-trip back into the same store.
382        store.get(&head.location).await.unwrap();
383    }
384
385    #[test]
386    fn strip_meta_preserves_version_and_etag() {
387        let prefix = Path::from("prefix");
388        let meta = ObjectMeta {
389            location: Path::from("prefix/foo"),
390            last_modified: chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap(),
391            size: 42,
392            e_tag: Some("etag-value".to_string()),
393            version: Some("version-value".to_string()),
394        };
395
396        let ObjectMeta {
397            location,
398            last_modified,
399            size,
400            e_tag,
401            version,
402        } = strip_meta(&prefix, meta.clone());
403
404        assert_eq!(location, Path::from("foo"));
405        assert_eq!(last_modified, meta.last_modified);
406        assert_eq!(size, meta.size);
407        assert_eq!(e_tag, meta.e_tag);
408        assert_eq!(version, meta.version);
409    }
410
411    #[tokio::test]
412    async fn prefix_multipart() {
413        let store = PrefixStore::new(InMemory::new(), "prefix");
414
415        multipart(&store, &store).await;
416        multipart_with_opts(&store, &store).await;
417        multipart_put_part_out_of_order(&store, &store).await;
418        multipart_out_of_order(&store).await;
419        multipart_race_condition(&store, true).await;
420    }
421
422    #[cfg(feature = "cloud-base")]
423    #[tokio::test]
424    async fn signer() {
425        #[derive(Debug)]
426        struct Foo;
427
428        #[async_trait::async_trait]
429        impl Signer for Foo {
430            async fn signed_url(
431                &self,
432                method: http::Method,
433                path: &Path,
434                _expires_in: std::time::Duration,
435            ) -> Result<url::Url> {
436                Ok(url::Url::parse(&format!("ex:{path}?method={method}")).unwrap())
437            }
438        }
439
440        assert_eq!(
441            PrefixStore::new(Foo, "prefix")
442                .signed_url(
443                    http::Method::GET,
444                    &"foo".into(),
445                    std::time::Duration::from_secs(1)
446                )
447                .await
448                .unwrap(),
449            url::Url::parse("ex:prefix/foo?method=GET").unwrap()
450        );
451        assert_eq!(
452            PrefixStore::new(Foo, "prefix")
453                .signed_urls(
454                    http::Method::GET,
455                    &["foo".into()],
456                    std::time::Duration::from_secs(1)
457                )
458                .await
459                .unwrap(),
460            vec![url::Url::parse("ex:prefix/foo?method=GET").unwrap()]
461        );
462    }
463}