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_urls(
264        &self,
265        method: http::Method,
266        paths: &[Path],
267        expires_in: std::time::Duration,
268    ) -> Result<Vec<url::Url>> {
269        self.inner
270            .signed_urls(
271                method,
272                &paths.iter().map(|p| self.full_path(p)).collect::<Vec<_>>(),
273                expires_in,
274            )
275            .await
276    }
277}
278
279#[cfg(not(target_arch = "wasm32"))]
280#[cfg(test)]
281mod tests {
282    use std::slice;
283
284    use super::*;
285    use crate::local::LocalFileSystem;
286    use crate::memory::InMemory;
287    use crate::{ObjectStoreExt, integration::*};
288
289    use tempfile::TempDir;
290
291    #[tokio::test]
292    async fn prefix_test() {
293        let root = TempDir::new().unwrap();
294        let inner = LocalFileSystem::new_with_prefix(root.path()).unwrap();
295        let integration = PrefixStore::new(inner, "prefix");
296
297        put_get_delete_list(&integration).await;
298        get_opts(&integration).await;
299        list_uses_directories_correctly(&integration).await;
300        list_with_delimiter(&integration).await;
301        rename_and_copy(&integration).await;
302        copy_if_not_exists(&integration).await;
303        stream_get(&integration).await;
304    }
305
306    #[tokio::test]
307    async fn prefix_test_applies_prefix() {
308        let tmpdir = TempDir::new().unwrap();
309        let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
310
311        let location = Path::from("prefix/test_file.json");
312        let data = Bytes::from("arbitrary data");
313
314        local.put(&location, data.clone().into()).await.unwrap();
315
316        let prefix = PrefixStore::new(local, "prefix");
317        let location_prefix = Path::from("test_file.json");
318
319        let content_list = flatten_list_stream(&prefix, None).await.unwrap();
320        assert_eq!(content_list, slice::from_ref(&location_prefix));
321
322        let root = Path::from("/");
323        let content_list = flatten_list_stream(&prefix, Some(&root)).await.unwrap();
324        assert_eq!(content_list, slice::from_ref(&location_prefix));
325
326        let read_data = prefix
327            .get(&location_prefix)
328            .await
329            .unwrap()
330            .bytes()
331            .await
332            .unwrap();
333        assert_eq!(&*read_data, data);
334
335        let target_prefix = Path::from("/test_written.json");
336        prefix
337            .put(&target_prefix, data.clone().into())
338            .await
339            .unwrap();
340
341        prefix.delete(&location_prefix).await.unwrap();
342
343        let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
344
345        let err = local.get(&location).await.unwrap_err();
346        assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
347
348        let location = Path::from("prefix/test_written.json");
349        let read_data = local.get(&location).await.unwrap().bytes().await.unwrap();
350        assert_eq!(&*read_data, data)
351    }
352
353    // Regression test for
354    // https://github.com/apache/arrow-rs-object-store/issues/664:
355    // head/get returned an ObjectMeta whose location still contained the
356    // prefix, so round-tripping the returned location back into the store would
357    // double-prefix it.
358    #[tokio::test]
359    async fn prefix_head_get_strip_prefix() {
360        let store = PrefixStore::new(InMemory::new(), "prefix");
361        let path = Path::from("test_file");
362        store.put(&path, "data".into()).await.unwrap();
363
364        let head = store.head(&path).await.unwrap();
365        assert_eq!(head.location, path);
366        let get = store.get(&path).await.unwrap();
367        assert_eq!(get.meta.location, path);
368
369        // The returned location must round-trip back into the same store.
370        store.get(&head.location).await.unwrap();
371    }
372
373    #[test]
374    fn strip_meta_preserves_version_and_etag() {
375        let prefix = Path::from("prefix");
376        let meta = ObjectMeta {
377            location: Path::from("prefix/foo"),
378            last_modified: chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap(),
379            size: 42,
380            e_tag: Some("etag-value".to_string()),
381            version: Some("version-value".to_string()),
382        };
383
384        let ObjectMeta {
385            location,
386            last_modified,
387            size,
388            e_tag,
389            version,
390        } = strip_meta(&prefix, meta.clone());
391
392        assert_eq!(location, Path::from("foo"));
393        assert_eq!(last_modified, meta.last_modified);
394        assert_eq!(size, meta.size);
395        assert_eq!(e_tag, meta.e_tag);
396        assert_eq!(version, meta.version);
397    }
398
399    #[tokio::test]
400    async fn prefix_multipart() {
401        let store = PrefixStore::new(InMemory::new(), "prefix");
402
403        multipart(&store, &store).await;
404        multipart_with_opts(&store, &store).await;
405        multipart_put_part_out_of_order(&store, &store).await;
406        multipart_out_of_order(&store).await;
407        multipart_race_condition(&store, true).await;
408    }
409
410    #[cfg(feature = "cloud-base")]
411    #[tokio::test]
412    async fn signer() {
413        #[derive(Debug)]
414        struct Foo;
415
416        #[async_trait::async_trait]
417        impl Signer for Foo {
418            async fn signed_url(
419                &self,
420                method: http::Method,
421                path: &Path,
422                _expires_in: std::time::Duration,
423            ) -> Result<url::Url> {
424                Ok(url::Url::parse(&format!("ex:{path}?method={method}")).unwrap())
425            }
426        }
427
428        assert_eq!(
429            PrefixStore::new(Foo, "prefix")
430                .signed_url(
431                    http::Method::GET,
432                    &"foo".into(),
433                    std::time::Duration::from_secs(1)
434                )
435                .await
436                .unwrap(),
437            url::Url::parse("ex:prefix/foo?method=GET").unwrap()
438        );
439        assert_eq!(
440            PrefixStore::new(Foo, "prefix")
441                .signed_urls(
442                    http::Method::GET,
443                    &["foo".into()],
444                    std::time::Duration::from_secs(1)
445                )
446                .await
447                .unwrap(),
448            vec![url::Url::parse("ex:prefix/foo?method=GET").unwrap()]
449        );
450    }
451}