Skip to main content

object_store/
integration.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//! Integration tests for custom object store implementations
19//!
20//! NB: These tests will delete everything present in the provided [`DynObjectStore`].
21//!
22//! These tests are not a stable part of the public API and breaking changes may be made
23//! in patch releases.
24//!
25//! They are intended solely for testing purposes.
26
27use crate::list::{PaginatedListOptions, PaginatedListStore};
28use crate::multipart::MultipartStore;
29use crate::path::Path;
30use crate::{
31    Attribute, Attributes, DynObjectStore, Error, GetOptions, GetRange, MultipartUpload,
32    ObjectStore, ObjectStoreExt, PutMode, PutMultipartOptions, PutPayload, UpdateVersion,
33    WriteMultipart,
34};
35use bytes::Bytes;
36use futures_util::stream::FuturesUnordered;
37use futures_util::{StreamExt, TryStreamExt};
38use rand::{RngExt, rng};
39use std::collections::HashSet;
40use std::slice;
41
42pub(crate) async fn flatten_list_stream(
43    storage: &DynObjectStore,
44    prefix: Option<&Path>,
45) -> crate::Result<Vec<Path>> {
46    storage
47        .list(prefix)
48        .map_ok(|meta| meta.location)
49        .try_collect::<Vec<Path>>()
50        .await
51}
52
53/// Tests basic read/write and listing operations
54pub async fn put_get_delete_list(storage: &DynObjectStore) {
55    delete_fixtures(storage).await;
56
57    let content_list = flatten_list_stream(storage, None).await.unwrap();
58    assert!(
59        content_list.is_empty(),
60        "Expected list to be empty; found: {content_list:?}"
61    );
62
63    let location = Path::from("test_dir/test_file.json");
64
65    let data = Bytes::from("arbitrary data");
66    storage.put(&location, data.clone().into()).await.unwrap();
67
68    let root = Path::from("/");
69
70    // List everything
71    let content_list = flatten_list_stream(storage, None).await.unwrap();
72    assert_eq!(content_list, slice::from_ref(&location));
73
74    // Should behave the same as no prefix
75    let content_list = flatten_list_stream(storage, Some(&root)).await.unwrap();
76    assert_eq!(content_list, slice::from_ref(&location));
77
78    // List with delimiter
79    let result = storage.list_with_delimiter(None).await.unwrap();
80    assert_eq!(&result.objects, &[]);
81    assert_eq!(result.common_prefixes.len(), 1);
82    assert_eq!(result.common_prefixes[0], Path::from("test_dir"));
83
84    // Should behave the same as no prefix
85    let result = storage.list_with_delimiter(Some(&root)).await.unwrap();
86    assert!(result.objects.is_empty());
87    assert_eq!(result.common_prefixes.len(), 1);
88    assert_eq!(result.common_prefixes[0], Path::from("test_dir"));
89
90    // Should return not found
91    let err = storage.get(&Path::from("test_dir")).await.unwrap_err();
92    assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
93
94    // Should return not found
95    let err = storage.head(&Path::from("test_dir")).await.unwrap_err();
96    assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
97
98    // List everything starting with a prefix that should return results
99    let prefix = Path::from("test_dir");
100    let content_list = flatten_list_stream(storage, Some(&prefix)).await.unwrap();
101    assert_eq!(content_list, slice::from_ref(&location));
102
103    // List everything starting with a prefix that shouldn't return results
104    let prefix = Path::from("something");
105    let content_list = flatten_list_stream(storage, Some(&prefix)).await.unwrap();
106    assert!(content_list.is_empty());
107
108    let read_data = storage.get(&location).await.unwrap().bytes().await.unwrap();
109    assert_eq!(&*read_data, data);
110
111    // Test range request
112    let range = 3..7;
113    let range_result = storage.get_range(&location, range.clone()).await;
114
115    let bytes = range_result.unwrap();
116    assert_eq!(bytes, data.slice(range.start as usize..range.end as usize));
117
118    let opts = GetOptions::new().with_range(Some(GetRange::Bounded(2..5)));
119    let result = storage.get_opts(&location, opts).await.unwrap();
120    // Data is `"arbitrary data"`, length 14 bytes
121    assert_eq!(result.meta.size, 14); // Should return full object size (#5272)
122    assert_eq!(result.range, 2..5);
123    let bytes = result.bytes().await.unwrap();
124    assert_eq!(bytes, b"bit".as_ref());
125
126    let out_of_range = 200..300;
127    let out_of_range_result = storage.get_range(&location, out_of_range).await;
128
129    // Should be a non-fatal error
130    out_of_range_result.unwrap_err();
131
132    let opts = GetOptions::new().with_range(Some(GetRange::Bounded(2..100)));
133    let result = storage.get_opts(&location, opts).await.unwrap();
134    assert_eq!(result.range, 2..14);
135    assert_eq!(result.meta.size, 14);
136    let bytes = result.bytes().await.unwrap();
137    assert_eq!(bytes, b"bitrary data".as_ref());
138
139    let opts = GetOptions::new().with_range(Some(GetRange::Suffix(2)));
140    match storage.get_opts(&location, opts).await {
141        Ok(result) => {
142            assert_eq!(result.range, 12..14);
143            assert_eq!(result.meta.size, 14);
144            let bytes = result.bytes().await.unwrap();
145            assert_eq!(bytes, b"ta".as_ref());
146        }
147        Err(Error::NotSupported { .. }) => {}
148        Err(e) => panic!("{e}"),
149    }
150
151    let opts = GetOptions::new().with_range(Some(GetRange::Suffix(100)));
152    match storage.get_opts(&location, opts).await {
153        Ok(result) => {
154            assert_eq!(result.range, 0..14);
155            assert_eq!(result.meta.size, 14);
156            let bytes = result.bytes().await.unwrap();
157            assert_eq!(bytes, b"arbitrary data".as_ref());
158        }
159        Err(Error::NotSupported { .. }) => {}
160        Err(e) => panic!("{e}"),
161    }
162
163    let opts = GetOptions::new().with_range(Some(GetRange::Offset(3)));
164    let result = storage.get_opts(&location, opts).await.unwrap();
165    assert_eq!(result.range, 3..14);
166    assert_eq!(result.meta.size, 14);
167    let bytes = result.bytes().await.unwrap();
168    assert_eq!(bytes, b"itrary data".as_ref());
169
170    let opts = GetOptions::new().with_range(Some(GetRange::Offset(100)));
171    storage.get_opts(&location, opts).await.unwrap_err();
172
173    let ranges = vec![0..1, 2..3, 0..5];
174    let bytes = storage.get_ranges(&location, &ranges).await.unwrap();
175    for (range, bytes) in ranges.iter().zip(bytes) {
176        assert_eq!(bytes, data.slice(range.start as usize..range.end as usize));
177    }
178
179    let head = storage.head(&location).await.unwrap();
180    assert_eq!(head.size, data.len() as u64);
181    assert_eq!(head.location, location);
182
183    storage.delete(&location).await.unwrap();
184
185    let content_list = flatten_list_stream(storage, None).await.unwrap();
186    assert!(content_list.is_empty());
187
188    let err = storage.get(&location).await.unwrap_err();
189    assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
190
191    let err = storage.head(&location).await.unwrap_err();
192    assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
193
194    // Test handling of paths containing an encoded delimiter
195
196    let file_with_delimiter = Path::from_iter(["a", "b/c", "foo.file"]);
197    storage
198        .put(&file_with_delimiter, "arbitrary".into())
199        .await
200        .unwrap();
201
202    let files = flatten_list_stream(storage, None).await.unwrap();
203    assert_eq!(files, vec![file_with_delimiter.clone()]);
204
205    let files = flatten_list_stream(storage, Some(&Path::from("a/b")))
206        .await
207        .unwrap();
208    assert!(files.is_empty());
209
210    let files = storage
211        .list_with_delimiter(Some(&Path::from("a/b")))
212        .await
213        .unwrap();
214    assert!(files.common_prefixes.is_empty());
215    assert!(files.objects.is_empty());
216
217    let files = storage
218        .list_with_delimiter(Some(&Path::from("a")))
219        .await
220        .unwrap();
221    assert_eq!(files.common_prefixes, vec![Path::from_iter(["a", "b/c"])]);
222    assert!(files.objects.is_empty());
223
224    let files = storage
225        .list_with_delimiter(Some(&Path::from_iter(["a", "b/c"])))
226        .await
227        .unwrap();
228    assert!(files.common_prefixes.is_empty());
229    assert_eq!(files.objects.len(), 1);
230    assert_eq!(files.objects[0].location, file_with_delimiter);
231
232    storage.delete(&file_with_delimiter).await.unwrap();
233
234    // Test handling of paths containing non-ASCII characters, e.g. emoji
235
236    let emoji_prefix = Path::from("🙀");
237    let emoji_file = Path::from("🙀/😀.parquet");
238    storage.put(&emoji_file, "arbitrary".into()).await.unwrap();
239
240    storage.head(&emoji_file).await.unwrap();
241    storage
242        .get(&emoji_file)
243        .await
244        .unwrap()
245        .bytes()
246        .await
247        .unwrap();
248
249    let files = flatten_list_stream(storage, Some(&emoji_prefix))
250        .await
251        .unwrap();
252
253    assert_eq!(files, vec![emoji_file.clone()]);
254
255    let dst = Path::from("foo.parquet");
256    storage.copy(&emoji_file, &dst).await.unwrap();
257    let mut files = flatten_list_stream(storage, None).await.unwrap();
258    files.sort_unstable();
259    assert_eq!(files, vec![emoji_file.clone(), dst.clone()]);
260
261    let dst2 = Path::from("new/nested/foo.parquet");
262    storage.copy(&emoji_file, &dst2).await.unwrap();
263    let mut files = flatten_list_stream(storage, None).await.unwrap();
264    files.sort_unstable();
265    assert_eq!(files, vec![emoji_file.clone(), dst.clone(), dst2.clone()]);
266
267    let dst3 = Path::from("new/nested2/bar.parquet");
268    storage.rename(&dst, &dst3).await.unwrap();
269    let mut files = flatten_list_stream(storage, None).await.unwrap();
270    files.sort_unstable();
271    assert_eq!(files, vec![emoji_file.clone(), dst2.clone(), dst3.clone()]);
272
273    let err = storage.head(&dst).await.unwrap_err();
274    assert!(matches!(err, Error::NotFound { .. }));
275
276    storage.delete(&emoji_file).await.unwrap();
277    storage.delete(&dst3).await.unwrap();
278    storage.delete(&dst2).await.unwrap();
279    let files = flatten_list_stream(storage, Some(&emoji_prefix))
280        .await
281        .unwrap();
282    assert!(files.is_empty());
283
284    // Test handling of paths containing percent-encoded sequences
285
286    // "HELLO" percent encoded
287    let hello_prefix = Path::parse("%48%45%4C%4C%4F").unwrap();
288    let path = hello_prefix.clone().join("foo.parquet");
289
290    storage.put(&path, vec![0, 1].into()).await.unwrap();
291    let files = flatten_list_stream(storage, Some(&hello_prefix))
292        .await
293        .unwrap();
294    assert_eq!(files, vec![path.clone()]);
295
296    // Cannot list by decoded representation
297    let files = flatten_list_stream(storage, Some(&Path::from("HELLO")))
298        .await
299        .unwrap();
300    assert!(files.is_empty());
301
302    // Cannot access by decoded representation
303    let err = storage
304        .head(&Path::from("HELLO/foo.parquet"))
305        .await
306        .unwrap_err();
307    assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
308
309    storage.delete(&path).await.unwrap();
310
311    // Test handling of unicode paths
312    let path = Path::parse("🇦🇺/$shenanigans@@~.txt").unwrap();
313    storage.put(&path, "test".into()).await.unwrap();
314
315    let r = storage.get(&path).await.unwrap();
316    assert_eq!(r.bytes().await.unwrap(), "test");
317
318    let dir = Path::parse("🇦🇺").unwrap();
319    let r = storage.list_with_delimiter(None).await.unwrap();
320    assert!(r.common_prefixes.contains(&dir));
321
322    let r = storage.list_with_delimiter(Some(&dir)).await.unwrap();
323    assert_eq!(r.objects.len(), 1);
324    assert_eq!(r.objects[0].location, path);
325
326    storage.delete(&path).await.unwrap();
327
328    // Can also write non-percent encoded sequences
329    let path = Path::parse("%Q.parquet").unwrap();
330    storage.put(&path, vec![0, 1].into()).await.unwrap();
331
332    let files = flatten_list_stream(storage, None).await.unwrap();
333    assert_eq!(files, vec![path.clone()]);
334
335    storage.delete(&path).await.unwrap();
336
337    let path = Path::parse("foo bar/I contain spaces.parquet").unwrap();
338    storage.put(&path, vec![0, 1].into()).await.unwrap();
339    storage.head(&path).await.unwrap();
340
341    let files = flatten_list_stream(storage, Some(&Path::from("foo bar")))
342        .await
343        .unwrap();
344    assert_eq!(files, vec![path.clone()]);
345
346    storage.delete(&path).await.unwrap();
347
348    let files = flatten_list_stream(storage, None).await.unwrap();
349    assert!(files.is_empty(), "{files:?}");
350
351    // Test list order
352    let files = vec![
353        Path::from("a a/b.file"),
354        Path::parse("a%2Fa.file").unwrap(),
355        Path::from("a/😀.file"),
356        Path::from("a/a file"),
357        Path::parse("a/a%2F.file").unwrap(),
358        Path::from("a/a.file"),
359        Path::from("a/a/b.file"),
360        Path::from("a/b.file"),
361        Path::from("aa/a.file"),
362        Path::from("ab/a.file"),
363    ];
364
365    for file in &files {
366        storage.put(file, "foo".into()).await.unwrap();
367    }
368
369    let cases = [
370        (None, Path::from("a")),
371        (None, Path::from("a/a file")),
372        (None, Path::from("a/a/b.file")),
373        (None, Path::from("ab/a.file")),
374        (None, Path::from("a%2Fa.file")),
375        (None, Path::from("a/😀.file")),
376        (Some(Path::from("a")), Path::from("")),
377        (Some(Path::from("a")), Path::from("a")),
378        (Some(Path::from("a")), Path::from("a/😀")),
379        (Some(Path::from("a")), Path::from("a/😀.file")),
380        (Some(Path::from("a")), Path::from("a/b")),
381        (Some(Path::from("a")), Path::from("a/a/b.file")),
382    ];
383
384    for (prefix, offset) in cases {
385        let s = storage.list_with_offset(prefix.as_ref(), &offset);
386        let mut actual: Vec<_> = s.map_ok(|x| x.location).try_collect().await.unwrap();
387
388        actual.sort_unstable();
389
390        let expected: Vec<_> = files
391            .iter()
392            .filter(|x| {
393                let prefix_match = prefix.as_ref().map(|p| x.prefix_matches(p)).unwrap_or(true);
394                prefix_match && *x > &offset
395            })
396            .cloned()
397            .collect();
398
399        assert_eq!(actual, expected, "{prefix:?} - {offset:?}");
400    }
401
402    // Test bulk delete
403    let paths = vec![
404        Path::from("a/a.file"),
405        Path::from("a/a/b.file"),
406        Path::from("aa/a.file"),
407        Path::from("does_not_exist"),
408        Path::from("I'm a < & weird path"),
409        Path::from("ab/a.file"),
410        Path::from("a/😀.file"),
411    ];
412
413    storage.put(&paths[4], "foo".into()).await.unwrap();
414
415    let out_paths = storage
416        .delete_stream(futures_util::stream::iter(paths.clone()).map(Ok).boxed())
417        .collect::<Vec<_>>()
418        .await;
419
420    assert_eq!(out_paths.len(), paths.len());
421
422    let expect_errors = [3];
423
424    for (i, input_path) in paths.iter().enumerate() {
425        let err = storage.head(input_path).await.unwrap_err();
426        assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
427
428        if expect_errors.contains(&i) {
429            // Some object stores will report NotFound, but others (such as S3) will
430            // report success regardless.
431            match &out_paths[i] {
432                Err(Error::NotFound { path: out_path, .. }) => {
433                    assert!(out_path.ends_with(&input_path.to_string()));
434                }
435                Ok(out_path) => {
436                    assert_eq!(out_path, input_path);
437                }
438                _ => panic!("unexpected error"),
439            }
440        } else {
441            assert_eq!(out_paths[i].as_ref().unwrap(), input_path);
442        }
443    }
444
445    delete_fixtures(storage).await;
446
447    let path = Path::from("empty");
448    storage.put(&path, PutPayload::default()).await.unwrap();
449    let meta = storage.head(&path).await.unwrap();
450    assert_eq!(meta.size, 0);
451    let data = storage.get(&path).await.unwrap().bytes().await.unwrap();
452    assert_eq!(data.len(), 0);
453
454    storage.delete(&path).await.unwrap();
455}
456
457/// Tests the ability to read and write [`Attributes`]
458pub async fn put_get_attributes(integration: &dyn ObjectStore) {
459    // Test handling of attributes
460    let attributes = Attributes::from_iter([
461        (Attribute::CacheControl, "max-age=604800"),
462        (
463            Attribute::ContentDisposition,
464            r#"attachment; filename="test.html""#,
465        ),
466        (Attribute::ContentEncoding, "gzip"),
467        (Attribute::ContentLanguage, "en-US"),
468        (Attribute::ContentType, "text/html; charset=utf-8"),
469        (Attribute::Metadata("test_key".into()), "test_value"),
470        (
471            Attribute::Metadata("key_with_spaces".into()),
472            "hello   world",
473        ),
474    ]);
475
476    let path = Path::from("attributes");
477    let opts = attributes.clone().into();
478    match integration.put_opts(&path, "foo".into(), opts).await {
479        Ok(_) => {
480            let r = integration.get(&path).await.unwrap();
481            assert_eq!(r.attributes, attributes);
482        }
483        Err(Error::NotImplemented { .. }) => {}
484        Err(e) => panic!("{e}"),
485    }
486
487    let opts = attributes.clone().into();
488    match integration.put_multipart_opts(&path, opts).await {
489        Ok(mut w) => {
490            w.put_part("foo".into()).await.unwrap();
491            w.complete().await.unwrap();
492
493            let r = integration.get(&path).await.unwrap();
494            assert_eq!(r.attributes, attributes);
495        }
496        Err(Error::NotImplemented { .. }) => {}
497        Err(e) => panic!("{e}"),
498    }
499}
500
501/// Tests conditional read requests
502pub async fn get_opts(storage: &dyn ObjectStore) {
503    let path = Path::from("test");
504    storage.put(&path, "foo".into()).await.unwrap();
505    let meta = storage.head(&path).await.unwrap();
506
507    let options = GetOptions::new().with_if_unmodified_since(Some(meta.last_modified));
508    match storage.get_opts(&path, options).await {
509        Ok(_) | Err(Error::NotSupported { .. }) => {}
510        Err(e) => panic!("{e}"),
511    }
512
513    let options = GetOptions::new().with_if_unmodified_since(Some(
514        meta.last_modified + chrono::Duration::try_hours(10).unwrap(),
515    ));
516    match storage.get_opts(&path, options).await {
517        Ok(_) | Err(Error::NotSupported { .. }) => {}
518        Err(e) => panic!("{e}"),
519    }
520
521    let options = GetOptions::new().with_if_unmodified_since(Some(
522        meta.last_modified - chrono::Duration::try_hours(10).unwrap(),
523    ));
524    match storage.get_opts(&path, options).await {
525        Err(Error::Precondition { .. } | Error::NotSupported { .. }) => {}
526        d => panic!("{d:?}"),
527    }
528
529    let options = GetOptions::new().with_if_modified_since(Some(meta.last_modified));
530    match storage.get_opts(&path, options).await {
531        Err(Error::NotModified { .. } | Error::NotSupported { .. }) => {}
532        d => panic!("{d:?}"),
533    }
534
535    let options = GetOptions::new().with_if_modified_since(Some(
536        meta.last_modified - chrono::Duration::try_hours(10).unwrap(),
537    ));
538    match storage.get_opts(&path, options).await {
539        Ok(_) | Err(Error::NotSupported { .. }) => {}
540        Err(e) => panic!("{e}"),
541    }
542
543    let tag = meta.e_tag.unwrap();
544    let options = GetOptions::new().with_if_match(Some(tag.clone()));
545    storage.get_opts(&path, options).await.unwrap();
546
547    let options = GetOptions::new().with_if_match(Some("invalid".to_string()));
548    let err = storage.get_opts(&path, options).await.unwrap_err();
549    assert!(matches!(err, Error::Precondition { .. }), "{err}");
550
551    let options = GetOptions::new().with_if_none_match(Some(tag.clone()));
552    let err = storage.get_opts(&path, options).await.unwrap_err();
553    assert!(matches!(err, Error::NotModified { .. }), "{err}");
554
555    let options = GetOptions::new().with_if_none_match(Some("invalid".to_string()));
556    storage.get_opts(&path, options).await.unwrap();
557
558    let result = storage.put(&path, "test".into()).await.unwrap();
559    let new_tag = result.e_tag.unwrap();
560    assert_ne!(tag, new_tag);
561
562    let meta = storage.head(&path).await.unwrap();
563    assert_eq!(meta.e_tag.unwrap(), new_tag);
564
565    let options = GetOptions::new().with_if_match(Some(new_tag.clone()));
566    storage.get_opts(&path, options).await.unwrap();
567
568    let options = GetOptions::new().with_if_match(Some(tag));
569    let err = storage.get_opts(&path, options).await.unwrap_err();
570    assert!(matches!(err, Error::Precondition { .. }), "{err}");
571
572    if let Some(version) = meta.version {
573        storage.put(&path, "bar".into()).await.unwrap();
574
575        let options = GetOptions::new().with_version(Some(version));
576
577        // Can retrieve previous version
578        let get_opts = storage.get_opts(&path, options).await.unwrap();
579        let old = get_opts.bytes().await.unwrap();
580        assert_eq!(old, b"test".as_slice());
581
582        // Current version contains the updated data
583        let current = storage.get(&path).await.unwrap().bytes().await.unwrap();
584        assert_eq!(&current, b"bar".as_slice());
585    }
586}
587
588/// Tests conditional writes
589pub async fn put_opts(storage: &dyn ObjectStore, supports_update: bool) {
590    delete_fixtures(storage).await;
591    let path = Path::from("put_opts");
592    let v1 = storage
593        .put_opts(&path, "a".into(), PutMode::Create.into())
594        .await
595        .unwrap();
596
597    let err = storage
598        .put_opts(&path, "b".into(), PutMode::Create.into())
599        .await
600        .unwrap_err();
601    assert!(matches!(err, Error::AlreadyExists { .. }), "{err}");
602
603    let b = storage.get(&path).await.unwrap().bytes().await.unwrap();
604    assert_eq!(b.as_ref(), b"a");
605
606    if !supports_update {
607        let err = storage
608            .put_opts(&path, "c".into(), PutMode::Update(v1.clone().into()).into())
609            .await
610            .unwrap_err();
611        assert!(matches!(err, Error::NotImplemented { .. }), "{err}");
612
613        return;
614    }
615
616    let v2 = storage
617        .put_opts(&path, "c".into(), PutMode::Update(v1.clone().into()).into())
618        .await
619        .unwrap();
620
621    let b = storage.get(&path).await.unwrap().bytes().await.unwrap();
622    assert_eq!(b.as_ref(), b"c");
623
624    let err = storage
625        .put_opts(&path, "d".into(), PutMode::Update(v1.into()).into())
626        .await
627        .unwrap_err();
628    assert!(matches!(err, Error::Precondition { .. }), "{err}");
629
630    storage
631        .put_opts(&path, "e".into(), PutMode::Update(v2.clone().into()).into())
632        .await
633        .unwrap();
634
635    let b = storage.get(&path).await.unwrap().bytes().await.unwrap();
636    assert_eq!(b.as_ref(), b"e");
637
638    // Update not exists
639    let path = Path::from("I don't exist");
640    let err = storage
641        .put_opts(&path, "e".into(), PutMode::Update(v2.into()).into())
642        .await
643        .unwrap_err();
644    assert!(matches!(err, Error::Precondition { .. }), "{err}");
645
646    const NUM_WORKERS: usize = 5;
647    const NUM_INCREMENTS: usize = 10;
648
649    let path = Path::from("RACE");
650    let mut futures: FuturesUnordered<_> = (0..NUM_WORKERS)
651        .map(|_| async {
652            for _ in 0..NUM_INCREMENTS {
653                loop {
654                    match storage.get(&path).await {
655                        Ok(r) => {
656                            let mode = PutMode::Update(UpdateVersion {
657                                e_tag: r.meta.e_tag.clone(),
658                                version: r.meta.version.clone(),
659                            });
660
661                            let b = r.bytes().await.unwrap();
662                            let v: usize = std::str::from_utf8(&b).unwrap().parse().unwrap();
663                            let new = (v + 1).to_string();
664
665                            match storage.put_opts(&path, new.into(), mode.into()).await {
666                                Ok(_) => break,
667                                Err(Error::Precondition { .. }) => continue,
668                                Err(e) => return Err(e),
669                            }
670                        }
671                        Err(Error::NotFound { .. }) => {
672                            let mode = PutMode::Create;
673                            match storage.put_opts(&path, "1".into(), mode.into()).await {
674                                Ok(_) => break,
675                                Err(Error::AlreadyExists { .. }) => continue,
676                                Err(e) => return Err(e),
677                            }
678                        }
679                        Err(e) => return Err(e),
680                    }
681                }
682            }
683            Ok(())
684        })
685        .collect();
686
687    while futures.next().await.transpose().unwrap().is_some() {}
688    let b = storage.get(&path).await.unwrap().bytes().await.unwrap();
689    let v = std::str::from_utf8(&b).unwrap().parse::<usize>().unwrap();
690    assert_eq!(v, NUM_WORKERS * NUM_INCREMENTS);
691}
692
693/// Returns a chunk of length `chunk_length`
694fn get_chunk(chunk_length: usize) -> Bytes {
695    let mut data = vec![0_u8; chunk_length];
696    let mut rng = rng();
697    // Set a random selection of bytes
698    for _ in 0..1000 {
699        data[rng.random_range(0..chunk_length)] = rng.random();
700    }
701    data.into()
702}
703
704/// Returns `num_chunks` of length `chunks`
705fn get_chunks(chunk_length: usize, num_chunks: usize) -> Vec<Bytes> {
706    (0..num_chunks).map(|_| get_chunk(chunk_length)).collect()
707}
708
709/// Tests the ability to perform multipart writes
710pub async fn stream_get(storage: &DynObjectStore) {
711    let location = Path::from("test_dir/test_upload_file.txt");
712
713    // Can write to storage
714    let data = get_chunks(5 * 1024 * 1024, 3);
715    let bytes_expected = data.concat();
716    let mut upload = storage.put_multipart(&location).await.unwrap();
717    let uploads = data.into_iter().map(|x| upload.put_part(x.into()));
718    futures_util::future::try_join_all(uploads).await.unwrap();
719
720    // Object should not yet exist in store
721    let meta_res = storage.head(&location).await;
722    assert!(meta_res.is_err());
723    assert!(matches!(
724        meta_res.unwrap_err(),
725        crate::Error::NotFound { .. }
726    ));
727
728    let files = flatten_list_stream(storage, None).await.unwrap();
729    assert_eq!(&files, &[]);
730
731    let result = storage.list_with_delimiter(None).await.unwrap();
732    assert_eq!(&result.objects, &[]);
733
734    upload.complete().await.unwrap();
735
736    let bytes_written = storage.get(&location).await.unwrap().bytes().await.unwrap();
737    assert_eq!(bytes_expected, bytes_written);
738
739    // Can overwrite some storage
740    // Sizes chosen to ensure we write three parts
741    let data = get_chunks(3_200_000, 7);
742    let bytes_expected = data.concat();
743    let upload = storage.put_multipart(&location).await.unwrap();
744    let mut writer = WriteMultipart::new(upload);
745    for chunk in &data {
746        writer.write(chunk)
747    }
748    writer.finish().await.unwrap();
749    let bytes_written = storage.get(&location).await.unwrap().bytes().await.unwrap();
750    assert_eq!(bytes_expected, bytes_written);
751
752    let location = Path::from("test_dir/test_put_part.txt");
753    let upload = storage.put_multipart(&location).await.unwrap();
754    let mut write = WriteMultipart::new(upload);
755    write.put(vec![0; 2].into());
756    write.put(vec![3; 4].into());
757    write.finish().await.unwrap();
758
759    let meta = storage.head(&location).await.unwrap();
760    assert_eq!(meta.size, 6);
761
762    let location = Path::from("test_dir/test_put_part_mixed.txt");
763    let upload = storage.put_multipart(&location).await.unwrap();
764    let mut write = WriteMultipart::new(upload);
765    write.put(vec![0; 2].into());
766    write.write(&[1, 2, 3]);
767    write.put(vec![4, 5, 6, 7].into());
768    write.finish().await.unwrap();
769
770    let r = storage.get(&location).await.unwrap();
771    let r = r.bytes().await.unwrap();
772    assert_eq!(r.as_ref(), &[0, 0, 1, 2, 3, 4, 5, 6, 7]);
773
774    // We can abort an empty write
775    let location = Path::from("test_dir/test_abort_upload.txt");
776    let mut upload = storage.put_multipart(&location).await.unwrap();
777    upload.abort().await.unwrap();
778    let get_res = storage.get(&location).await;
779    assert!(get_res.is_err());
780    assert!(matches!(
781        get_res.unwrap_err(),
782        crate::Error::NotFound { .. }
783    ));
784
785    // We can abort an in-progress write
786    let mut upload = storage.put_multipart(&location).await.unwrap();
787    upload
788        .put_part(data.first().unwrap().clone().into())
789        .await
790        .unwrap();
791
792    upload.abort().await.unwrap();
793    let get_res = storage.get(&location).await;
794    assert!(get_res.is_err());
795    assert!(matches!(get_res.unwrap_err(), Error::NotFound { .. }));
796}
797
798/// Tests that directories are transparent
799pub async fn list_uses_directories_correctly(storage: &DynObjectStore) {
800    delete_fixtures(storage).await;
801
802    let content_list = flatten_list_stream(storage, None).await.unwrap();
803    assert!(
804        content_list.is_empty(),
805        "Expected list to be empty; found: {content_list:?}"
806    );
807
808    let location1 = Path::from("foo/x.json");
809    let location2 = Path::from("foo.bar/y.json");
810
811    let data = PutPayload::from("arbitrary data");
812    storage.put(&location1, data.clone()).await.unwrap();
813    storage.put(&location2, data).await.unwrap();
814
815    let prefix = Path::from("foo");
816    let content_list = flatten_list_stream(storage, Some(&prefix)).await.unwrap();
817    assert_eq!(content_list, slice::from_ref(&location1));
818
819    let result = storage.list_with_delimiter(Some(&prefix)).await.unwrap();
820    assert_eq!(result.objects.len(), 1);
821    assert_eq!(result.objects[0].location, location1);
822    assert_eq!(result.common_prefixes, &[]);
823
824    // Listing an existing path (file) should return an empty list:
825    // https://github.com/apache/arrow-rs/issues/3712
826    let content_list = flatten_list_stream(storage, Some(&location1))
827        .await
828        .unwrap();
829    assert_eq!(content_list, &[]);
830
831    let list = storage.list_with_delimiter(Some(&location1)).await.unwrap();
832    assert_eq!(list.objects, &[]);
833    assert_eq!(list.common_prefixes, &[]);
834
835    let prefix = Path::from("foo/x");
836    let content_list = flatten_list_stream(storage, Some(&prefix)).await.unwrap();
837    assert_eq!(content_list, &[]);
838
839    let list = storage.list_with_delimiter(Some(&prefix)).await.unwrap();
840    assert_eq!(list.objects, &[]);
841    assert_eq!(list.common_prefixes, &[]);
842}
843
844/// Tests listing with delimiter
845pub async fn list_with_delimiter(storage: &DynObjectStore) {
846    delete_fixtures(storage).await;
847
848    // ==================== check: store is empty ====================
849    let content_list = flatten_list_stream(storage, None).await.unwrap();
850    assert!(content_list.is_empty());
851
852    // ==================== do: create files ====================
853    let data = Bytes::from("arbitrary data");
854
855    let files: Vec<_> = [
856        "test_file",
857        "mydb/wb/000/000/000.segment",
858        "mydb/wb/000/000/001.segment",
859        "mydb/wb/000/000/002.segment",
860        "mydb/wb/001/001/000.segment",
861        "mydb/wb/foo.json",
862        "mydb/wbwbwb/111/222/333.segment",
863        "mydb/data/whatevs",
864    ]
865    .iter()
866    .map(|&s| Path::from(s))
867    .collect();
868
869    for f in &files {
870        storage.put(f, data.clone().into()).await.unwrap();
871    }
872
873    // ==================== check: prefix-list `mydb/wb` (directory) ====================
874    let prefix = Path::from("mydb/wb");
875
876    let expected_000 = Path::from("mydb/wb/000");
877    let expected_001 = Path::from("mydb/wb/001");
878    let expected_location = Path::from("mydb/wb/foo.json");
879
880    let result = storage.list_with_delimiter(Some(&prefix)).await.unwrap();
881
882    assert_eq!(result.common_prefixes, vec![expected_000, expected_001]);
883    assert_eq!(result.objects.len(), 1);
884
885    let object = &result.objects[0];
886
887    assert_eq!(object.location, expected_location);
888    assert_eq!(object.size, data.len() as u64);
889
890    // ==================== check: prefix-list `mydb/wb/000/000/001` (partial filename doesn't match) ====================
891    let prefix = Path::from("mydb/wb/000/000/001");
892
893    let result = storage.list_with_delimiter(Some(&prefix)).await.unwrap();
894    assert!(result.common_prefixes.is_empty());
895    assert_eq!(result.objects.len(), 0);
896
897    // ==================== check: prefix-list `not_there` (non-existing prefix) ====================
898    let prefix = Path::from("not_there");
899
900    let result = storage.list_with_delimiter(Some(&prefix)).await.unwrap();
901    assert!(result.common_prefixes.is_empty());
902    assert!(result.objects.is_empty());
903
904    // ==================== do: remove all files ====================
905    for f in &files {
906        storage.delete(f).await.unwrap();
907    }
908
909    // ==================== check: store is empty ====================
910    let content_list = flatten_list_stream(storage, None).await.unwrap();
911    assert!(content_list.is_empty());
912}
913
914/// Tests fetching a non-existent object returns a not found error
915pub async fn get_nonexistent_object(
916    storage: &DynObjectStore,
917    location: Option<Path>,
918) -> crate::Result<Bytes> {
919    let location = location.unwrap_or_else(|| Path::from("this_file_should_not_exist"));
920
921    let err = storage.head(&location).await.unwrap_err();
922    assert!(matches!(err, Error::NotFound { .. }));
923
924    storage.get(&location).await?.bytes().await
925}
926
927/// Tests copying
928pub async fn rename_and_copy(storage: &DynObjectStore) {
929    // Create two objects
930    let path1 = Path::from("test1");
931    let path2 = Path::from("test2");
932    let contents1 = Bytes::from("cats");
933    let contents2 = Bytes::from("dogs");
934
935    // copy() make both objects identical
936    storage.put(&path1, contents1.clone().into()).await.unwrap();
937    storage.put(&path2, contents2.clone().into()).await.unwrap();
938    storage.copy(&path1, &path2).await.unwrap();
939    let new_contents = storage.get(&path2).await.unwrap().bytes().await.unwrap();
940    assert_eq!(&new_contents, &contents1);
941
942    // rename() copies contents and deletes original
943    storage.put(&path1, contents1.clone().into()).await.unwrap();
944    storage.put(&path2, contents2.clone().into()).await.unwrap();
945    storage.rename(&path1, &path2).await.unwrap();
946    let new_contents = storage.get(&path2).await.unwrap().bytes().await.unwrap();
947    assert_eq!(&new_contents, &contents1);
948    let result = storage.get(&path1).await;
949    assert!(result.is_err());
950    assert!(matches!(result.unwrap_err(), Error::NotFound { .. }));
951
952    // Clean up
953    storage.delete(&path2).await.unwrap();
954}
955
956/// Tests copy if not exists
957pub async fn copy_if_not_exists(storage: &DynObjectStore) {
958    // Create two objects
959    let path1 = Path::from("test1");
960    let path2 = Path::from("not_exists_nested/test2");
961    let contents1 = Bytes::from("cats");
962    let contents2 = Bytes::from("dogs");
963
964    // copy_if_not_exists() errors if destination already exists
965    storage.put(&path1, contents1.clone().into()).await.unwrap();
966    storage.put(&path2, contents2.clone().into()).await.unwrap();
967    let result = storage.copy_if_not_exists(&path1, &path2).await;
968    assert!(result.is_err());
969    assert!(matches!(
970        result.unwrap_err(),
971        crate::Error::AlreadyExists { .. }
972    ));
973
974    // copy_if_not_exists() copies contents and allows deleting original
975    storage.delete(&path2).await.unwrap();
976    storage.copy_if_not_exists(&path1, &path2).await.unwrap();
977    storage.delete(&path1).await.unwrap();
978    let new_contents = storage.get(&path2).await.unwrap().bytes().await.unwrap();
979    assert_eq!(&new_contents, &contents1);
980    let result = storage.get(&path1).await;
981    assert!(result.is_err());
982    assert!(matches!(result.unwrap_err(), crate::Error::NotFound { .. }));
983
984    // Clean up
985    storage.delete(&path2).await.unwrap();
986}
987
988/// Tests copy and renaming behaviour of non-existent objects
989pub async fn copy_rename_nonexistent_object(storage: &DynObjectStore) {
990    // Create empty source object
991    let path1 = Path::from("test1");
992
993    // Create destination object
994    let path2 = Path::from("test2");
995    storage.put(&path2, "hello".into()).await.unwrap();
996
997    // copy() errors if source does not exist
998    let result = storage.copy(&path1, &path2).await;
999    assert!(result.is_err());
1000    assert!(matches!(result.unwrap_err(), crate::Error::NotFound { .. }));
1001
1002    // rename() errors if source does not exist
1003    let result = storage.rename(&path1, &path2).await;
1004    assert!(result.is_err());
1005    assert!(matches!(result.unwrap_err(), crate::Error::NotFound { .. }));
1006
1007    // copy_if_not_exists() errors if source does not exist
1008    let result = storage.copy_if_not_exists(&path1, &path2).await;
1009    assert!(result.is_err());
1010    assert!(matches!(result.unwrap_err(), crate::Error::NotFound { .. }));
1011
1012    // Clean up
1013    storage.delete(&path2).await.unwrap();
1014}
1015
1016/// Tests [`MultipartStore`]
1017pub async fn multipart(storage: &dyn ObjectStore, multipart: &dyn MultipartStore) {
1018    let path = Path::from("test_multipart");
1019    let chunk_size = 5 * 1024 * 1024;
1020
1021    let chunks = get_chunks(chunk_size, 2);
1022
1023    let id = multipart.create_multipart(&path).await.unwrap();
1024
1025    let parts: Vec<_> = futures_util::stream::iter(chunks)
1026        .enumerate()
1027        .map(|(idx, b)| multipart.put_part(&path, &id, idx, b.into()))
1028        .buffered(2)
1029        .try_collect()
1030        .await
1031        .unwrap();
1032
1033    multipart
1034        .complete_multipart(&path, &id, parts)
1035        .await
1036        .unwrap();
1037
1038    let meta = storage.head(&path).await.unwrap();
1039    assert_eq!(meta.size, chunk_size as u64 * 2);
1040
1041    // Empty case
1042    let path = Path::from("test_empty_multipart");
1043
1044    let id = multipart.create_multipart(&path).await.unwrap();
1045
1046    let parts = vec![];
1047
1048    multipart
1049        .complete_multipart(&path, &id, parts)
1050        .await
1051        .unwrap();
1052
1053    let meta = storage.head(&path).await.unwrap();
1054    assert_eq!(meta.size, 0);
1055}
1056
1057/// Tests [`MultipartStore::create_multipart_opts`]
1058pub async fn multipart_with_opts(storage: &dyn ObjectStore, multipart: &dyn MultipartStore) {
1059    let path = Path::from("test_multipart_with_opts");
1060    let chunk_size = 5 * 1024 * 1024;
1061    let chunks = get_chunks(chunk_size, 2);
1062    let attributes =
1063        Attributes::from_iter([(Attribute::Metadata("test_key".into()), "test_value")]);
1064    let opts = PutMultipartOptions {
1065        attributes: attributes.clone(),
1066        ..Default::default()
1067    };
1068
1069    let id = multipart.create_multipart_opts(&path, opts).await.unwrap();
1070
1071    let parts: Vec<_> = futures_util::stream::iter(chunks)
1072        .enumerate()
1073        .map(|(idx, b)| multipart.put_part(&path, &id, idx, b.into()))
1074        .buffered(2)
1075        .try_collect()
1076        .await
1077        .unwrap();
1078
1079    multipart
1080        .complete_multipart(&path, &id, parts)
1081        .await
1082        .unwrap();
1083
1084    let result = storage.get(&path).await.unwrap();
1085    assert_eq!(result.meta.size, chunk_size as u64 * 2);
1086    // Some providers add default attributes, e.g. S3 may report a default Content-Type.
1087    for (attribute, value) in &attributes {
1088        assert_eq!(result.attributes.get(attribute), Some(value));
1089    }
1090}
1091
1092/// Tests that [`MultipartStore::put_part`] may be invoked with non-sequential part indices.
1093pub async fn multipart_put_part_out_of_order(
1094    storage: &dyn ObjectStore,
1095    multipart: &dyn MultipartStore,
1096) {
1097    let path = Path::from("test_multipart_put_part_out_of_order");
1098
1099    // S3: each part except the last must be ≥ 5 MiB.
1100    const MIN_MULTIPART_PART: usize = 5 * 1024 * 1024;
1101    let part0_data = Bytes::from(vec![0xAA_u8; MIN_MULTIPART_PART]);
1102    let part1_data = Bytes::from(vec![0xBB_u8; MIN_MULTIPART_PART]);
1103    let part2_data = Bytes::from_static(b"tail");
1104
1105    let upload_id = multipart.create_multipart(&path).await.unwrap();
1106
1107    let part2 = multipart
1108        .put_part(&path, &upload_id, 2, PutPayload::from(part2_data.clone()))
1109        .await
1110        .unwrap();
1111
1112    let part0 = multipart
1113        .put_part(&path, &upload_id, 0, PutPayload::from(part0_data.clone()))
1114        .await
1115        .unwrap();
1116
1117    let part1 = multipart
1118        .put_part(&path, &upload_id, 1, PutPayload::from(part1_data.clone()))
1119        .await
1120        .unwrap();
1121
1122    let result = multipart
1123        .complete_multipart(&path, &upload_id, vec![part0, part1, part2])
1124        .await
1125        .unwrap();
1126    assert!(result.e_tag.is_some(), "Expected e_tag in PutResult");
1127
1128    let data = storage.get(&path).await.unwrap().bytes().await.unwrap();
1129    assert_eq!(
1130        data.len(),
1131        part0_data.len() + part1_data.len() + part2_data.len()
1132    );
1133    assert!(data[..MIN_MULTIPART_PART].iter().all(|&b| b == 0xAA));
1134    assert!(
1135        data[MIN_MULTIPART_PART..2 * MIN_MULTIPART_PART]
1136            .iter()
1137            .all(|&b| b == 0xBB)
1138    );
1139    assert_eq!(&data[2 * MIN_MULTIPART_PART..], part2_data.as_ref());
1140}
1141
1142async fn delete_fixtures(storage: &DynObjectStore) {
1143    let paths = storage.list(None).map_ok(|meta| meta.location).boxed();
1144    storage
1145        .delete_stream(paths)
1146        .try_collect::<Vec<_>>()
1147        .await
1148        .unwrap();
1149}
1150
1151/// Tests a race condition where 2 threads are performing multipart writes to the same path
1152pub async fn multipart_race_condition(storage: &dyn ObjectStore, last_writer_wins: bool) {
1153    let path = Path::from("test_multipart_race_condition");
1154
1155    let mut multipart_upload_1 = storage.put_multipart(&path).await.unwrap();
1156    let mut multipart_upload_2 = storage.put_multipart(&path).await.unwrap();
1157
1158    /// Create a string like `"1:"`  followed by `part` padded to 5,300,000 places
1159    ///
1160    /// equivalent of format!("{prefix}:{part:05300000}"), which is no longer supported
1161    ///
1162    /// See: <https://github.com/apache/arrow-rs-object-store/issues/343>
1163    fn make_payload(prefix: u8, part: u8) -> Vec<u8> {
1164        // prefix = 1 byte
1165        // ':' = 1 byte
1166        let mut payload = vec![b'0'; 5_300_002];
1167        payload[0] = prefix;
1168        payload[1] = b':';
1169        payload[2] = part;
1170        payload
1171    }
1172
1173    // Upload parts interleaved
1174    multipart_upload_1
1175        .put_part(Bytes::from(make_payload(b'1', 0)).into())
1176        .await
1177        .unwrap();
1178    multipart_upload_2
1179        .put_part(Bytes::from(make_payload(b'2', 0)).into())
1180        .await
1181        .unwrap();
1182
1183    multipart_upload_2
1184        .put_part(Bytes::from(make_payload(b'2', 1)).into())
1185        .await
1186        .unwrap();
1187    multipart_upload_1
1188        .put_part(Bytes::from(make_payload(b'1', 1)).into())
1189        .await
1190        .unwrap();
1191
1192    multipart_upload_1
1193        .put_part(Bytes::from(make_payload(b'1', 2)).into())
1194        .await
1195        .unwrap();
1196    multipart_upload_2
1197        .put_part(Bytes::from(make_payload(b'2', 2)).into())
1198        .await
1199        .unwrap();
1200
1201    multipart_upload_2
1202        .put_part(Bytes::from(make_payload(b'2', 3)).into())
1203        .await
1204        .unwrap();
1205    multipart_upload_1
1206        .put_part(Bytes::from(make_payload(b'1', 3)).into())
1207        .await
1208        .unwrap();
1209
1210    multipart_upload_1
1211        .put_part(Bytes::from(make_payload(b'1', 4)).into())
1212        .await
1213        .unwrap();
1214    multipart_upload_2
1215        .put_part(Bytes::from(make_payload(b'2', 4)).into())
1216        .await
1217        .unwrap();
1218
1219    multipart_upload_1.complete().await.unwrap();
1220
1221    if last_writer_wins {
1222        multipart_upload_2.complete().await.unwrap();
1223    } else {
1224        let err = multipart_upload_2.complete().await.unwrap_err();
1225
1226        assert!(matches!(err, crate::Error::Generic { .. }), "{err}");
1227    }
1228
1229    let get_result = storage.get(&path).await.unwrap();
1230    let result_bytes = get_result.bytes().await.unwrap();
1231
1232    let expected_writer_prefix = if last_writer_wins { b'2' } else { b'1' };
1233    let mut expected_writer_contents = vec![];
1234    for part in 0..5 {
1235        expected_writer_contents.append(&mut make_payload(expected_writer_prefix, part));
1236    }
1237
1238    assert!(result_bytes.starts_with(&expected_writer_contents));
1239}
1240
1241/// Tests performing out of order multipart uploads
1242pub async fn multipart_out_of_order(storage: &dyn ObjectStore) {
1243    let path = Path::from("test_multipart_out_of_order");
1244    let mut multipart_upload = storage.put_multipart(&path).await.unwrap();
1245
1246    let part1 = std::iter::repeat_n(b'1', 5 * 1024 * 1024).collect::<Bytes>();
1247    let part2 = std::iter::repeat_n(b'2', 5 * 1024 * 1024).collect::<Bytes>();
1248    let part3 = std::iter::repeat_n(b'3', 5 * 1024 * 1024).collect::<Bytes>();
1249    let full = [part1.as_ref(), part2.as_ref(), part3.as_ref()].concat();
1250
1251    let fut1 = multipart_upload.put_part(part1.into());
1252    let fut2 = multipart_upload.put_part(part2.into());
1253    let fut3 = multipart_upload.put_part(part3.into());
1254    // note order is 2,3,1 , different than the parts were created in
1255    fut2.await.unwrap();
1256    fut3.await.unwrap();
1257    fut1.await.unwrap();
1258
1259    multipart_upload.complete().await.unwrap();
1260
1261    let result = storage.get(&path).await.unwrap();
1262    let bytes = result.bytes().await.unwrap();
1263    assert_eq!(bytes, full);
1264}
1265
1266/// Tests [`PaginatedListStore`]
1267pub async fn list_paginated(storage: &dyn ObjectStore, list: &dyn PaginatedListStore) {
1268    delete_fixtures(storage).await;
1269
1270    let r = list.list_paginated(None, Default::default()).await.unwrap();
1271    assert_eq!(r.page_token, None);
1272    assert_eq!(r.result.objects, vec![]);
1273    assert_eq!(r.result.common_prefixes, vec![]);
1274
1275    let p1 = Path::from("foo/bar");
1276    let p2 = Path::from("foo/bax");
1277    let p3 = Path::from("foo/baz/bar");
1278    let p4 = Path::from("foo/baz/banana");
1279    let p5 = Path::from("fob/banana");
1280    let p6 = Path::from("fongle/banana");
1281
1282    let paths = HashSet::from_iter([&p1, &p2, &p3, &p4, &p5, &p6]);
1283
1284    for path in &paths {
1285        storage.put(path, vec![1].into()).await.unwrap();
1286    }
1287
1288    // Test basic listing
1289
1290    let mut listed = HashSet::new();
1291    let mut opts = PaginatedListOptions {
1292        max_keys: Some(5),
1293        ..Default::default()
1294    };
1295    let ret = list.list_paginated(None, opts.clone()).await.unwrap();
1296    assert_eq!(ret.result.objects.len(), 5);
1297    listed.extend(ret.result.objects.iter().map(|x| &x.location));
1298
1299    opts.page_token = Some(ret.page_token.unwrap());
1300    let ret = list.list_paginated(None, opts.clone()).await.unwrap();
1301    assert_eq!(ret.result.objects.len(), 1);
1302    listed.extend(ret.result.objects.iter().map(|x| &x.location));
1303
1304    assert_eq!(listed, paths);
1305
1306    // List with prefix
1307    let prefix = Some("foo/");
1308    opts.page_token = None;
1309    let ret = list.list_paginated(prefix, opts.clone()).await.unwrap();
1310    assert_eq!(ret.result.objects.len(), 4);
1311    assert!(ret.page_token.is_none());
1312
1313    let actual = HashSet::from_iter(ret.result.objects.iter().map(|x| &x.location));
1314    assert_eq!(actual, HashSet::<&Path>::from_iter([&p1, &p2, &p3, &p4]));
1315
1316    // List with partial prefix
1317    let prefix = Some("fo");
1318    opts.page_token = None;
1319    let ret = list.list_paginated(prefix, opts.clone()).await.unwrap();
1320    assert_eq!(ret.result.objects.len(), 5);
1321    listed.extend(ret.result.objects.iter().map(|x| &x.location));
1322
1323    opts.page_token = Some(ret.page_token.unwrap());
1324    let ret = list.list_paginated(prefix, opts.clone()).await.unwrap();
1325    assert_eq!(ret.result.objects.len(), 1);
1326    listed.extend(ret.result.objects.iter().map(|x| &x.location));
1327
1328    assert_eq!(listed, paths);
1329
1330    // List with prefix and delimiter
1331    let prefix = Some("foo/");
1332    opts.page_token = None;
1333    opts.delimiter = Some("/".into());
1334    let ret = list.list_paginated(prefix, opts.clone()).await.unwrap();
1335    assert_eq!(ret.result.objects.len(), 2);
1336    assert_eq!(ret.result.common_prefixes, vec![Path::from("foo/baz")]);
1337    assert!(ret.page_token.is_none());
1338
1339    let actual = HashSet::from_iter(ret.result.objects.iter().map(|x| &x.location));
1340    assert_eq!(actual, HashSet::<&Path>::from_iter([&p1, &p2]));
1341
1342    // List with partial prefix and delimiter
1343    let prefix = Some("fo");
1344    opts.page_token = None;
1345    opts.delimiter = Some("/".into());
1346    let ret = list.list_paginated(prefix, opts.clone()).await.unwrap();
1347    assert_eq!(ret.result.objects.len(), 0);
1348    assert_eq!(
1349        HashSet::<Path>::from_iter(ret.result.common_prefixes),
1350        HashSet::from_iter([Path::from("foo"), Path::from("fob"), Path::from("fongle")])
1351    );
1352    assert!(ret.page_token.is_none());
1353}
1354
1355/// Tests that [`ObjectStore::list_with_offset`] returns an exclusive list
1356/// that does not include the offset value itself.
1357/// This is needed because some object stores (i.e. Azure) return inclusive results,
1358/// while AWS S3 and GCP return exclusive results.
1359pub async fn list_with_offset_exclusivity(storage: &DynObjectStore) {
1360    delete_fixtures(storage).await;
1361
1362    // Create a set of test files with predictable ordering
1363    let files = vec![
1364        Path::from("test_offset/file_a.txt"),
1365        Path::from("test_offset/file_b.txt"),
1366        Path::from("test_offset/file_c.txt"),
1367        Path::from("test_offset/file_d.txt"),
1368        Path::from("test_offset/file_e.txt"),
1369    ];
1370
1371    let data = Bytes::from("test data");
1372    for file in &files {
1373        storage.put(file, data.clone().into()).await.unwrap();
1374    }
1375
1376    // Test 1: Offset at file_b, should return c, d, e but NOT b
1377    let offset = Path::from("test_offset/file_b.txt");
1378    let result: Vec<Path> = storage
1379        .list_with_offset(Some(&Path::from("test_offset")), &offset)
1380        .map_ok(|meta| meta.location)
1381        .try_collect()
1382        .await
1383        .unwrap();
1384
1385    // Verify that the offset file (file_b) is NOT in the results
1386    assert!(
1387        !result.contains(&offset),
1388        "list_with_offset should not include the offset value itself. Found {:?} in results: {:?}",
1389        offset,
1390        result
1391    );
1392
1393    // Verify that only files after the offset are returned
1394    let expected = vec![
1395        Path::from("test_offset/file_c.txt"),
1396        Path::from("test_offset/file_d.txt"),
1397        Path::from("test_offset/file_e.txt"),
1398    ];
1399    assert_eq!(
1400        result.len(),
1401        expected.len(),
1402        "Expected {} files after offset, got {}",
1403        expected.len(),
1404        result.len()
1405    );
1406    for file in &expected {
1407        assert!(
1408            result.contains(file),
1409            "Expected file {:?} to be in results",
1410            file
1411        );
1412    }
1413
1414    // Test 2: Offset at file_a, should return b, c, d, e but NOT a
1415    let offset = Path::from("test_offset/file_a.txt");
1416    let result: Vec<Path> = storage
1417        .list_with_offset(Some(&Path::from("test_offset")), &offset)
1418        .map_ok(|meta| meta.location)
1419        .try_collect()
1420        .await
1421        .unwrap();
1422
1423    assert!(
1424        !result.contains(&offset),
1425        "list_with_offset should not include the offset value itself. Found {:?} in results: {:?}",
1426        offset,
1427        result
1428    );
1429    assert_eq!(result.len(), 4, "Expected 4 files after first file");
1430
1431    // Test 3: Offset at file_e (last file), should return empty list
1432    let offset = Path::from("test_offset/file_e.txt");
1433    let result: Vec<Path> = storage
1434        .list_with_offset(Some(&Path::from("test_offset")), &offset)
1435        .map_ok(|meta| meta.location)
1436        .try_collect()
1437        .await
1438        .unwrap();
1439
1440    assert!(
1441        result.is_empty(),
1442        "list_with_offset with offset at last file should return empty list, got: {:?}",
1443        result
1444    );
1445
1446    // Test 4: Test without prefix - offset should still be exclusive
1447    let offset = Path::from("test_offset/file_c.txt");
1448    let result: Vec<Path> = storage
1449        .list_with_offset(None, &offset)
1450        .map_ok(|meta| meta.location)
1451        .try_collect()
1452        .await
1453        .unwrap();
1454
1455    assert!(
1456        !result.contains(&offset),
1457        "list_with_offset without prefix should not include the offset value itself. Found {:?} in results: {:?}",
1458        offset,
1459        result
1460    );
1461
1462    // Verify only files after file_c are included
1463    let expected = vec![
1464        Path::from("test_offset/file_d.txt"),
1465        Path::from("test_offset/file_e.txt"),
1466    ];
1467    assert_eq!(result.len(), expected.len());
1468    for file in &expected {
1469        assert!(result.contains(file), "Expected file {:?} in results", file);
1470    }
1471
1472    // Clean up
1473    delete_fixtures(storage).await;
1474}
1475
1476#[cfg(all(
1477    feature = "reqwest",
1478    not(all(target_arch = "wasm32", target_os = "wasi"))
1479))]
1480mod marker {
1481    use super::*;
1482    use crate::ClientOptions;
1483    use crate::client::{
1484        HttpClient, HttpConnector, HttpError, HttpRequest, HttpResponse, HttpService,
1485        ReqwestConnector,
1486    };
1487    use async_trait::async_trait;
1488
1489    /// A marker inserted into response extensions by [`MarkerHttpConnector`]
1490    #[derive(Clone, Debug, PartialEq)]
1491    struct Marker;
1492
1493    /// [`HttpService`] middleware that tags every response with a [`Marker`]
1494    #[derive(Debug)]
1495    struct MarkerService(HttpClient);
1496
1497    #[async_trait]
1498    impl HttpService for MarkerService {
1499        async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError> {
1500            let mut response = self.0.execute(req).await?;
1501            response.extensions_mut().insert(Marker);
1502            Ok(response)
1503        }
1504    }
1505
1506    /// An [`HttpConnector`] that tags the extensions of every HTTP response with a
1507    /// marker, allowing [`response_extensions`] to verify their propagation
1508    #[derive(Debug, Default)]
1509    pub struct MarkerHttpConnector(ReqwestConnector);
1510
1511    impl HttpConnector for MarkerHttpConnector {
1512        fn connect(&self, options: &ClientOptions) -> crate::Result<HttpClient> {
1513            let client = self.0.connect(options)?;
1514            Ok(HttpClient::new(MarkerService(client)))
1515        }
1516    }
1517
1518    /// Tests that the extensions of HTTP responses are propagated to returned results
1519    ///
1520    /// The provided store must have been built with [`MarkerHttpConnector`]
1521    pub async fn response_extensions(storage: &DynObjectStore, test_multipart: bool) {
1522        delete_fixtures(storage).await;
1523
1524        let location = Path::from("test_response_extensions/file.txt");
1525        let data = Bytes::from("arbitrary data");
1526
1527        let put = storage.put(&location, data.clone().into()).await.unwrap();
1528        assert!(
1529            put.extensions.get::<Marker>().is_some(),
1530            "PutResult should contain the response extensions"
1531        );
1532
1533        let get = storage.get(&location).await.unwrap();
1534        assert!(
1535            get.extensions.get::<Marker>().is_some(),
1536            "GetResult should contain the response extensions"
1537        );
1538
1539        let list = storage.list_with_delimiter(None).await.unwrap();
1540        assert!(
1541            list.extensions.get::<Marker>().is_some(),
1542            "ListResult should contain the response extensions"
1543        );
1544
1545        if test_multipart {
1546            let mut upload = storage.put_multipart(&location).await.unwrap();
1547            upload.put_part(data.into()).await.unwrap();
1548            let complete = upload.complete().await.unwrap();
1549            assert!(
1550                complete.extensions.get::<Marker>().is_some(),
1551                "PutResult of a completed multipart upload should contain the response extensions"
1552            );
1553        }
1554
1555        delete_fixtures(storage).await;
1556    }
1557}
1558
1559#[cfg(all(
1560    feature = "reqwest",
1561    not(all(target_arch = "wasm32", target_os = "wasi"))
1562))]
1563pub use marker::{MarkerHttpConnector, response_extensions};