Skip to main content

object_store/
compat.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::fmt::{self, Debug, Display, Formatter};
17use std::future::IntoFuture;
18use std::io;
19use std::ops::Range;
20use std::sync::Arc;
21
22use async_trait::async_trait;
23use bytes::Bytes;
24use datafusion_object_store::path::Path;
25use datafusion_object_store::{
26    Attribute, Attributes, CopyMode, CopyOptions, GetOptions, GetRange, GetResult,
27    GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore as ArrowObjectStore,
28    PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult, UploadPart,
29};
30use futures::stream::BoxStream;
31use futures::{FutureExt, StreamExt, TryStreamExt};
32use opendal::options::CopyOptions as OpendalCopyOptions;
33use opendal::raw::percent_decode_path;
34use opendal::{Buffer, Operator, OperatorInfo, Writer};
35use tokio::sync::{Mutex, oneshot};
36
37/// OpendalStore implements ObjectStore trait by using opendal.
38///
39/// This allows users to use opendal as an object store without extra cost.
40///
41/// Visit [`opendal::services`] for more information about supported services.
42///
43/// ```no_run
44/// use std::sync::Arc;
45///
46/// use bytes::Bytes;
47/// use object_store::path::Path;
48/// use object_store::ObjectStore;
49/// use object_store_opendal::OpendalStore;
50/// use opendal::services::S3;
51/// use opendal::{Builder, Operator};
52///
53/// #[tokio::main]
54/// async fn main() {
55///    let builder = S3::default()
56///     .access_key_id("my_access_key")
57///     .secret_access_key("my_secret_key")
58///     .endpoint("my_endpoint")
59///     .region("my_region");
60///
61///     // Create a new operator
62///     let operator = Operator::new(builder).unwrap().finish();
63///
64///     // Create a new object store
65///     let object_store = Arc::new(OpendalStore::new(operator));
66///
67///     let path = Path::from("data/nested/test.txt");
68///     let bytes = Bytes::from_static(b"hello, world! I am nested.");
69///
70///     object_store.put(&path, bytes.clone().into()).await.unwrap();
71///
72///     let content = object_store
73///         .get(&path)
74///         .await
75///         .unwrap()
76///         .bytes()
77///         .await
78///         .unwrap();
79///
80///     assert_eq!(content, bytes);
81/// }
82/// ```
83#[derive(Clone)]
84pub struct OpendalStore {
85    info: Arc<OperatorInfo>,
86    inner: Operator,
87}
88
89impl OpendalStore {
90    /// Create OpendalStore by given Operator.
91    pub fn new(op: Operator) -> Self {
92        Self {
93            info: op.info().into(),
94            inner: op,
95        }
96    }
97
98    /// Get the Operator info.
99    pub fn info(&self) -> &OperatorInfo {
100        self.info.as_ref()
101    }
102
103    /// Copy a file from one location to another.
104    async fn copy_request(
105        &self,
106        from: &Path,
107        to: &Path,
108        if_not_exists: bool,
109    ) -> datafusion_object_store::Result<()> {
110        let mut copy_options = OpendalCopyOptions::default();
111        if if_not_exists {
112            copy_options.if_not_exists = true;
113        }
114
115        // Perform the copy operation
116        self.inner
117            .copy_options(
118                &percent_decode_path(from.as_ref()),
119                &percent_decode_path(to.as_ref()),
120                copy_options,
121            )
122            .await
123            .map_err(|err| {
124                if if_not_exists && err.kind() == opendal::ErrorKind::AlreadyExists {
125                    datafusion_object_store::Error::AlreadyExists {
126                        path: to.to_string(),
127                        source: Box::new(err),
128                    }
129                } else {
130                    format_object_store_error(err, from.as_ref())
131                }
132            })?;
133
134        Ok(())
135    }
136}
137
138impl Debug for OpendalStore {
139    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
140        f.debug_struct("OpendalStore")
141            .field("scheme", &self.info.scheme())
142            .field("name", &self.info.name())
143            .field("root", &self.info.root())
144            .field("capability", &self.info.full_capability())
145            .finish()
146    }
147}
148
149impl Display for OpendalStore {
150    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
151        let info = self.inner.info();
152        write!(
153            f,
154            "Opendal({}, bucket={}, root={})",
155            info.scheme(),
156            info.name(),
157            info.root()
158        )
159    }
160}
161
162impl From<Operator> for OpendalStore {
163    fn from(value: Operator) -> Self {
164        Self::new(value)
165    }
166}
167
168#[async_trait]
169impl ArrowObjectStore for OpendalStore {
170    async fn put_opts(
171        &self,
172        location: &Path,
173        bytes: PutPayload,
174        opts: PutOptions,
175    ) -> datafusion_object_store::Result<PutResult> {
176        let decoded_location = percent_decode_path(location.as_ref());
177        let mut future_write = self
178            .inner
179            .write_with(&decoded_location, Buffer::from_iter(bytes));
180        let opts_mode = opts.mode.clone();
181        match opts.mode {
182            PutMode::Overwrite => {}
183            PutMode::Create => {
184                future_write = future_write.if_not_exists(true);
185            }
186            PutMode::Update(update_version) => {
187                let Some(etag) = update_version.e_tag else {
188                    return Err(datafusion_object_store::Error::NotSupported {
189                        source: Box::new(opendal::Error::new(
190                            opendal::ErrorKind::Unsupported,
191                            "etag is required for conditional put",
192                        )),
193                    });
194                };
195                future_write = future_write.if_match(etag.as_str());
196            }
197        }
198        let rp = future_write.await.map_err(|err| {
199            match format_object_store_error(err, location.as_ref()) {
200                datafusion_object_store::Error::Precondition { path, source }
201                    if opts_mode == PutMode::Create =>
202                {
203                    datafusion_object_store::Error::AlreadyExists { path, source }
204                }
205                e => e,
206            }
207        })?;
208
209        let e_tag = rp.etag().map(|s| s.to_string());
210        let version = rp.version().map(|s| s.to_string());
211
212        Ok(PutResult { e_tag, version })
213    }
214
215    async fn put_multipart_opts(
216        &self,
217        location: &Path,
218        opts: PutMultipartOptions,
219    ) -> datafusion_object_store::Result<Box<dyn MultipartUpload>> {
220        const DEFAULT_CONCURRENT: usize = 8;
221
222        let mut options = opendal::options::WriteOptions {
223            concurrent: DEFAULT_CONCURRENT,
224            ..Default::default()
225        };
226
227        let mut user_metadata = HashMap::new();
228
229        for (key, value) in opts.attributes.iter() {
230            match key {
231                Attribute::CacheControl => {
232                    options.cache_control = Some(value.to_string());
233                }
234                Attribute::ContentDisposition => {
235                    options.content_disposition = Some(value.to_string());
236                }
237                Attribute::ContentEncoding => {
238                    options.content_encoding = Some(value.to_string());
239                }
240                Attribute::ContentLanguage => continue,
241                Attribute::ContentType => {
242                    options.content_type = Some(value.to_string());
243                }
244                Attribute::Metadata(k) => {
245                    user_metadata.insert(k.to_string(), value.to_string());
246                }
247                _ => {}
248            }
249        }
250
251        if !user_metadata.is_empty() {
252            options.user_metadata = Some(user_metadata);
253        }
254
255        let decoded_location = percent_decode_path(location.as_ref());
256        let writer = self
257            .inner
258            .writer_options(&decoded_location, options)
259            .await
260            .map_err(|err| format_object_store_error(err, location.as_ref()))?;
261        let upload = OpendalMultipartUpload::new(writer, location.clone());
262
263        Ok(Box::new(upload))
264    }
265
266    async fn get_opts(
267        &self,
268        location: &Path,
269        options: GetOptions,
270    ) -> datafusion_object_store::Result<GetResult> {
271        let raw_location = percent_decode_path(location.as_ref());
272        let meta = {
273            let mut s = self.inner.stat_with(&raw_location);
274            if let Some(version) = &options.version {
275                s = s.version(version.as_str())
276            }
277            if let Some(if_match) = &options.if_match {
278                s = s.if_match(if_match.as_str());
279            }
280            if let Some(if_none_match) = &options.if_none_match {
281                s = s.if_none_match(if_none_match.as_str());
282            }
283            if let Some(if_modified_since) =
284                options.if_modified_since.and_then(datetime_to_timestamp)
285            {
286                s = s.if_modified_since(if_modified_since);
287            }
288            if let Some(if_unmodified_since) =
289                options.if_unmodified_since.and_then(datetime_to_timestamp)
290            {
291                s = s.if_unmodified_since(if_unmodified_since);
292            }
293            s.await
294                .map_err(|err| format_object_store_error(err, location.as_ref()))?
295        };
296
297        let mut attributes = Attributes::new();
298        if let Some(user_meta) = meta.user_metadata() {
299            for (key, value) in user_meta {
300                attributes.insert(
301                    Attribute::Metadata(key.clone().into()),
302                    value.clone().into(),
303                );
304            }
305        }
306
307        let meta = ObjectMeta {
308            location: location.clone(),
309            last_modified: meta
310                .last_modified()
311                .and_then(timestamp_to_datetime)
312                .unwrap_or_default(),
313            size: meta.content_length(),
314            e_tag: meta.etag().map(|x| x.to_string()),
315            version: meta.version().map(|x| x.to_string()),
316        };
317
318        if options.head {
319            return Ok(GetResult {
320                payload: GetResultPayload::Stream(Box::pin(futures::stream::empty())),
321                range: 0..0,
322                meta,
323                attributes,
324            });
325        }
326
327        let reader = {
328            let mut r = self.inner.reader_with(raw_location.as_ref());
329            if let Some(version) = options.version {
330                r = r.version(version.as_str());
331            }
332            if let Some(if_match) = options.if_match {
333                r = r.if_match(if_match.as_str());
334            }
335            if let Some(if_none_match) = options.if_none_match {
336                r = r.if_none_match(if_none_match.as_str());
337            }
338            if let Some(if_modified_since) =
339                options.if_modified_since.and_then(datetime_to_timestamp)
340            {
341                r = r.if_modified_since(if_modified_since);
342            }
343            if let Some(if_unmodified_since) =
344                options.if_unmodified_since.and_then(datetime_to_timestamp)
345            {
346                r = r.if_unmodified_since(if_unmodified_since);
347            }
348            r.await
349                .map_err(|err| format_object_store_error(err, location.as_ref()))?
350        };
351
352        let read_range = match options.range {
353            Some(GetRange::Bounded(r)) => {
354                if r.start >= r.end || r.start >= meta.size {
355                    0..0
356                } else {
357                    let end = r.end.min(meta.size);
358                    r.start..end
359                }
360            }
361            Some(GetRange::Offset(r)) => {
362                if r < meta.size {
363                    r..meta.size
364                } else {
365                    0..0
366                }
367            }
368            Some(GetRange::Suffix(r)) if r < meta.size => (meta.size - r)..meta.size,
369            _ => 0..meta.size,
370        };
371
372        let stream = reader
373            .into_bytes_stream(read_range.start..read_range.end)
374            .await
375            .map_err(|err| format_object_store_error(err, location.as_ref()))?
376            .map_ok(|buf| buf)
377            .map_err(|err: io::Error| datafusion_object_store::Error::Generic {
378                store: "IoError",
379                source: Box::new(err),
380            });
381
382        Ok(GetResult {
383            payload: GetResultPayload::Stream(Box::pin(stream)),
384            range: read_range.start..read_range.end,
385            meta,
386            attributes,
387        })
388    }
389
390    async fn get_ranges(
391        &self,
392        location: &Path,
393        ranges: &[Range<u64>],
394    ) -> datafusion_object_store::Result<Vec<Bytes>> {
395        if ranges.is_empty() {
396            return Ok(Vec::new());
397        }
398
399        let raw_location = percent_decode_path(location.as_ref());
400        let reader = self
401            .inner
402            .reader_with(raw_location.as_ref())
403            .await
404            .map_err(|err| format_object_store_error(err, location.as_ref()))?;
405        let buffers = reader
406            .fetch(ranges.to_vec())
407            .await
408            .map_err(|err| format_object_store_error(err, location.as_ref()))?;
409
410        Ok(buffers.into_iter().map(|buf| buf.to_bytes()).collect())
411    }
412
413    fn delete_stream(
414        &self,
415        locations: BoxStream<'static, datafusion_object_store::Result<Path>>,
416    ) -> BoxStream<'static, datafusion_object_store::Result<Path>> {
417        let this = self.clone();
418        locations
419            .map(move |location| {
420                let this = this.clone();
421                async move {
422                    let location = location?;
423                    let decoded_location = percent_decode_path(location.as_ref());
424                    this.inner
425                        .delete(&decoded_location)
426                        .await
427                        .map_err(|err| format_object_store_error(err, location.as_ref()))?;
428                    Ok(location)
429                }
430            })
431            .buffered(10)
432            .boxed()
433    }
434
435    fn list(
436        &self,
437        prefix: Option<&Path>,
438    ) -> BoxStream<'static, datafusion_object_store::Result<ObjectMeta>> {
439        // object_store `Path` always removes trailing slash
440        // need to add it back
441        let path = prefix.map_or("".into(), |x| {
442            format!("{}/", percent_decode_path(x.as_ref()))
443        });
444
445        let this = self.clone();
446        let fut = async move {
447            let stream = this
448                .inner
449                .lister_with(&path)
450                .recursive(true)
451                .await
452                .map_err(|err| format_object_store_error(err, &path))?;
453
454            let stream = stream.then(|res| async {
455                let entry = res.map_err(|err| format_object_store_error(err, ""))?;
456                let meta = entry.metadata();
457
458                Ok(format_object_meta(entry.path(), meta))
459            });
460            Ok::<_, datafusion_object_store::Error>(stream)
461        };
462
463        fut.into_stream().try_flatten().boxed()
464    }
465
466    fn list_with_offset(
467        &self,
468        prefix: Option<&Path>,
469        offset: &Path,
470    ) -> BoxStream<'static, datafusion_object_store::Result<ObjectMeta>> {
471        let path = prefix.map_or("".into(), |x| {
472            format!("{}/", percent_decode_path(x.as_ref()))
473        });
474        let offset = offset.clone();
475
476        // clone self for 'static lifetime
477        // clone self is cheap
478        let this = self.clone();
479
480        let fut = async move {
481            let list_with_start_after = this.inner.info().full_capability().list_with_start_after;
482            let mut fut = this.inner.lister_with(&path).recursive(true);
483
484            // Use native start_after support if possible.
485            if list_with_start_after {
486                fut = fut.start_after(offset.as_ref());
487            }
488
489            let lister = fut
490                .await
491                .map_err(|err| format_object_store_error(err, &path))?
492                .then(move |entry| {
493                    let path = path.clone();
494                    let this = this.clone();
495                    async move {
496                        let entry = entry.map_err(|err| format_object_store_error(err, &path))?;
497                        let (path, metadata) = entry.into_parts();
498
499                        // If it's a dir or last_modified is present, we can use it directly.
500                        if metadata.is_dir() || metadata.last_modified().is_some() {
501                            let object_meta = format_object_meta(&path, &metadata);
502                            return Ok(object_meta);
503                        }
504
505                        let metadata = this
506                            .inner
507                            .stat(&path)
508                            .await
509                            .map_err(|err| format_object_store_error(err, &path))?;
510                        let object_meta = format_object_meta(&path, &metadata);
511                        Ok::<_, datafusion_object_store::Error>(object_meta)
512                    }
513                })
514                .boxed();
515
516            let stream = if list_with_start_after {
517                lister
518            } else {
519                lister
520                    .try_filter(move |entry| futures::future::ready(entry.location > offset))
521                    .boxed()
522            };
523
524            Ok::<_, datafusion_object_store::Error>(stream)
525        };
526
527        fut.into_stream().try_flatten().boxed()
528    }
529
530    async fn list_with_delimiter(
531        &self,
532        prefix: Option<&Path>,
533    ) -> datafusion_object_store::Result<ListResult> {
534        let path = prefix.map_or("".into(), |x| {
535            format!("{}/", percent_decode_path(x.as_ref()))
536        });
537        let mut stream = self
538            .inner
539            .lister_with(&path)
540            .into_future()
541            .await
542            .map_err(|err| format_object_store_error(err, &path))?;
543
544        let mut common_prefixes = Vec::new();
545        let mut objects = Vec::new();
546
547        while let Some(res) = stream.next().await {
548            let entry = res.map_err(|err| format_object_store_error(err, ""))?;
549            let meta = entry.metadata();
550
551            if meta.is_dir() {
552                common_prefixes.push(entry.path().into());
553            } else if meta.last_modified().is_some() {
554                objects.push(format_object_meta(entry.path(), meta));
555            } else {
556                let meta = self
557                    .inner
558                    .stat(entry.path())
559                    .await
560                    .map_err(|err| format_object_store_error(err, entry.path()))?;
561                objects.push(format_object_meta(entry.path(), &meta));
562            }
563        }
564
565        Ok(ListResult {
566            common_prefixes,
567            objects,
568        })
569    }
570
571    async fn copy_opts(
572        &self,
573        from: &Path,
574        to: &Path,
575        options: CopyOptions,
576    ) -> datafusion_object_store::Result<()> {
577        let if_not_exists = options.mode == CopyMode::Create;
578        self.copy_request(from, to, if_not_exists).await
579    }
580}
581
582/// `MultipartUpload` implementation based on `Writer` in opendal.
583///
584/// # Notes
585///
586/// OpenDAL writer can handle concurrent internally we don't generate real `UploadPart` like existing
587/// implementation do. Instead, we just write the part and notify the next task to be written.
588///
589/// The lock here doesn't really involve the write process, it's just for the notify mechanism.
590struct OpendalMultipartUpload {
591    writer: Arc<Mutex<Writer>>,
592    location: Path,
593    next_notify: oneshot::Receiver<()>,
594}
595
596impl OpendalMultipartUpload {
597    fn new(writer: Writer, location: Path) -> Self {
598        // an immediately dropped sender for the first part to write without waiting
599        let (_, rx) = oneshot::channel();
600
601        Self {
602            writer: Arc::new(Mutex::new(writer)),
603            location,
604            next_notify: rx,
605        }
606    }
607}
608
609#[async_trait]
610impl MultipartUpload for OpendalMultipartUpload {
611    fn put_part(&mut self, data: PutPayload) -> UploadPart {
612        let writer = self.writer.clone();
613        let location = self.location.clone();
614
615        // Generate next notify which will be notified after the current part is written.
616        let (tx, rx) = oneshot::channel();
617        // Fetch the notify for current part to wait for it to be written.
618        let last_rx = std::mem::replace(&mut self.next_notify, rx);
619
620        async move {
621            // Wait for the previous part to be written
622            let _ = last_rx.await;
623
624            let mut writer = writer.lock().await;
625            let result = writer
626                .write(Buffer::from_iter(data))
627                .await
628                .map_err(|err| format_object_store_error(err, location.as_ref()));
629
630            // Notify the next part to be written
631            drop(tx);
632
633            result
634        }
635        .boxed()
636    }
637
638    async fn complete(&mut self) -> datafusion_object_store::Result<PutResult> {
639        let mut writer = self.writer.lock().await;
640        let metadata = writer
641            .close()
642            .await
643            .map_err(|err| format_object_store_error(err, self.location.as_ref()))?;
644
645        let e_tag = metadata.etag().map(|s| s.to_string());
646        let version = metadata.version().map(|s| s.to_string());
647
648        Ok(PutResult { e_tag, version })
649    }
650
651    async fn abort(&mut self) -> datafusion_object_store::Result<()> {
652        let mut writer = self.writer.lock().await;
653        writer
654            .abort()
655            .await
656            .map_err(|err| format_object_store_error(err, self.location.as_ref()))
657    }
658}
659
660impl Debug for OpendalMultipartUpload {
661    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
662        f.debug_struct("OpendalMultipartUpload")
663            .field("location", &self.location)
664            .finish()
665    }
666}
667
668fn format_object_store_error(err: opendal::Error, path: &str) -> datafusion_object_store::Error {
669    match err.kind() {
670        opendal::ErrorKind::NotFound => datafusion_object_store::Error::NotFound {
671            path: path.to_string(),
672            source: Box::new(err),
673        },
674        opendal::ErrorKind::Unsupported => datafusion_object_store::Error::NotSupported {
675            source: Box::new(err),
676        },
677        opendal::ErrorKind::AlreadyExists => datafusion_object_store::Error::AlreadyExists {
678            path: path.to_string(),
679            source: Box::new(err),
680        },
681        opendal::ErrorKind::ConditionNotMatch => datafusion_object_store::Error::Precondition {
682            path: path.to_string(),
683            source: Box::new(err),
684        },
685        kind => datafusion_object_store::Error::Generic {
686            store: kind.into_static(),
687            source: Box::new(err),
688        },
689    }
690}
691
692fn format_object_meta(path: &str, meta: &opendal::Metadata) -> ObjectMeta {
693    ObjectMeta {
694        location: path.into(),
695        last_modified: meta
696            .last_modified()
697            .and_then(timestamp_to_datetime)
698            .unwrap_or_default(),
699        size: meta.content_length(),
700        e_tag: meta.etag().map(|x| x.to_string()),
701        version: meta.version().map(|x| x.to_string()),
702    }
703}
704
705fn timestamp_to_datetime(ts: opendal::raw::Timestamp) -> Option<chrono::DateTime<chrono::Utc>> {
706    let ts = ts.into_inner();
707    chrono::DateTime::<chrono::Utc>::from_timestamp(ts.as_second(), ts.subsec_nanosecond() as u32)
708}
709
710fn datetime_to_timestamp(dt: chrono::DateTime<chrono::Utc>) -> Option<opendal::raw::Timestamp> {
711    opendal::raw::Timestamp::new(dt.timestamp(), dt.timestamp_subsec_nanos() as i32).ok()
712}
713
714#[cfg(test)]
715mod tests {
716    use std::sync::Arc;
717
718    use bytes::Bytes;
719    use datafusion_object_store::path::Path;
720    use datafusion_object_store::{
721        ObjectStore as ArrowObjectStore, ObjectStoreExt, WriteMultipart,
722    };
723    use opendal::{Operator, services};
724    use rand::{Rng, RngCore};
725
726    use super::*;
727
728    async fn create_test_object_store() -> Arc<dyn ArrowObjectStore> {
729        let op = Operator::new(services::Memory::default()).unwrap().finish();
730        let object_store = Arc::new(OpendalStore::new(op));
731
732        let path: Path = "data/test.txt".into();
733        let bytes = Bytes::from_static(b"hello, world!");
734        object_store.put(&path, bytes.into()).await.unwrap();
735
736        let path: Path = "data/nested/test.txt".into();
737        let bytes = Bytes::from_static(b"hello, world! I am nested.");
738        object_store.put(&path, bytes.into()).await.unwrap();
739
740        object_store
741    }
742
743    #[tokio::test]
744    async fn test_basic() {
745        let op = Operator::new(services::Memory::default()).unwrap().finish();
746        let object_store: Arc<dyn ArrowObjectStore> = Arc::new(OpendalStore::new(op));
747
748        // Retrieve a specific file
749        let path: Path = "data/test.txt".into();
750
751        let bytes = Bytes::from_static(b"hello, world!");
752        object_store.put(&path, bytes.clone().into()).await.unwrap();
753
754        let meta = object_store.head(&path).await.unwrap();
755
756        assert_eq!(meta.size, 13);
757
758        assert_eq!(
759            object_store
760                .get(&path)
761                .await
762                .unwrap()
763                .bytes()
764                .await
765                .unwrap(),
766            bytes
767        );
768    }
769
770    #[tokio::test]
771    async fn test_put_multipart() {
772        let op = Operator::new(services::Memory::default()).unwrap().finish();
773        let object_store: Arc<dyn ArrowObjectStore> = Arc::new(OpendalStore::new(op));
774
775        let mut rng = rand::rng();
776
777        // Case complete
778        let path: Path = "data/test_complete.txt".into();
779        let upload = object_store.put_multipart(&path).await.unwrap();
780
781        let mut write = WriteMultipart::new(upload);
782
783        let mut all_bytes = vec![];
784        let round = rng.random_range(1..=1024);
785        for _ in 0..round {
786            let size = rng.random_range(1..=1024);
787            let mut bytes = vec![0; size];
788            rng.fill_bytes(&mut bytes);
789
790            all_bytes.extend_from_slice(&bytes);
791            write.put(bytes.into());
792        }
793
794        let _ = write.finish().await.unwrap();
795
796        let meta = object_store.head(&path).await.unwrap();
797
798        assert_eq!(meta.size, all_bytes.len() as u64);
799
800        assert_eq!(
801            object_store
802                .get(&path)
803                .await
804                .unwrap()
805                .bytes()
806                .await
807                .unwrap(),
808            Bytes::from(all_bytes)
809        );
810
811        // Case abort
812        let path: Path = "data/test_abort.txt".into();
813        let mut upload = object_store.put_multipart(&path).await.unwrap();
814        upload.put_part(vec![1; 1024].into()).await.unwrap();
815        upload.abort().await.unwrap();
816
817        let res = object_store.head(&path).await;
818        let err = res.unwrap_err();
819
820        assert!(matches!(
821            err,
822            datafusion_object_store::Error::NotFound { .. }
823        ))
824    }
825
826    #[tokio::test]
827    async fn test_list() {
828        let object_store = create_test_object_store().await;
829        let path: Path = "data/".into();
830        let results = object_store.list(Some(&path)).collect::<Vec<_>>().await;
831        assert_eq!(results.len(), 2);
832        let mut locations = results
833            .iter()
834            .map(|x| x.as_ref().unwrap().location.as_ref())
835            .collect::<Vec<_>>();
836
837        let expected_files = vec![
838            (
839                "data/nested/test.txt",
840                Bytes::from_static(b"hello, world! I am nested."),
841            ),
842            ("data/test.txt", Bytes::from_static(b"hello, world!")),
843        ];
844
845        let expected_locations = expected_files.iter().map(|x| x.0).collect::<Vec<&str>>();
846
847        locations.sort();
848        assert_eq!(locations, expected_locations);
849
850        for (location, bytes) in expected_files {
851            let path: Path = location.into();
852            assert_eq!(
853                object_store
854                    .get(&path)
855                    .await
856                    .unwrap()
857                    .bytes()
858                    .await
859                    .unwrap(),
860                bytes
861            );
862        }
863    }
864
865    #[tokio::test]
866    async fn test_list_with_delimiter() {
867        let object_store = create_test_object_store().await;
868        let path: Path = "data/".into();
869        let result = object_store.list_with_delimiter(Some(&path)).await.unwrap();
870        assert_eq!(result.objects.len(), 1);
871        assert_eq!(result.common_prefixes.len(), 1);
872        assert_eq!(result.objects[0].location.as_ref(), "data/test.txt");
873        assert_eq!(result.common_prefixes[0].as_ref(), "data/nested");
874    }
875
876    #[tokio::test]
877    async fn test_list_with_offset() {
878        let object_store = create_test_object_store().await;
879        let path: Path = "data/".into();
880        let offset: Path = "data/nested/test.txt".into();
881        let result = object_store
882            .list_with_offset(Some(&path), &offset)
883            .collect::<Vec<_>>()
884            .await;
885        assert_eq!(result.len(), 1);
886        assert_eq!(
887            result[0].as_ref().unwrap().location.as_ref(),
888            "data/test.txt"
889        );
890    }
891
892    mod stat_counter {
893        use std::sync::atomic::{AtomicUsize, Ordering};
894
895        use super::*;
896
897        #[derive(Debug, Clone)]
898        pub struct StatCounterLayer {
899            count: Arc<AtomicUsize>,
900        }
901
902        impl StatCounterLayer {
903            pub fn new(count: Arc<AtomicUsize>) -> Self {
904                Self { count }
905            }
906        }
907
908        impl<A: opendal::raw::Access> opendal::raw::Layer<A> for StatCounterLayer {
909            type LayeredAccess = StatCounterAccessor<A>;
910
911            fn layer(&self, inner: A) -> Self::LayeredAccess {
912                StatCounterAccessor {
913                    inner,
914                    count: self.count.clone(),
915                }
916            }
917        }
918
919        #[derive(Debug, Clone)]
920        pub struct StatCounterAccessor<A> {
921            inner: A,
922            count: Arc<AtomicUsize>,
923        }
924
925        impl<A: opendal::raw::Access> opendal::raw::LayeredAccess for StatCounterAccessor<A> {
926            type Inner = A;
927            type Reader = A::Reader;
928            type Writer = A::Writer;
929            type Lister = A::Lister;
930            type Deleter = A::Deleter;
931
932            fn inner(&self) -> &Self::Inner {
933                &self.inner
934            }
935
936            async fn stat(
937                &self,
938                path: &str,
939                args: opendal::raw::OpStat,
940            ) -> opendal::Result<opendal::raw::RpStat> {
941                self.count.fetch_add(1, Ordering::SeqCst);
942                self.inner.stat(path, args).await
943            }
944
945            async fn read(
946                &self,
947                path: &str,
948                args: opendal::raw::OpRead,
949            ) -> opendal::Result<(opendal::raw::RpRead, Self::Reader)> {
950                self.inner.read(path, args).await
951            }
952
953            async fn write(
954                &self,
955                path: &str,
956                args: opendal::raw::OpWrite,
957            ) -> opendal::Result<(opendal::raw::RpWrite, Self::Writer)> {
958                self.inner.write(path, args).await
959            }
960
961            async fn delete(&self) -> opendal::Result<(opendal::raw::RpDelete, Self::Deleter)> {
962                self.inner.delete().await
963            }
964
965            async fn list(
966                &self,
967                path: &str,
968                args: opendal::raw::OpList,
969            ) -> opendal::Result<(opendal::raw::RpList, Self::Lister)> {
970                self.inner.list(path, args).await
971            }
972
973            async fn copy(
974                &self,
975                from: &str,
976                to: &str,
977                args: opendal::raw::OpCopy,
978            ) -> opendal::Result<opendal::raw::RpCopy> {
979                self.inner.copy(from, to, args).await
980            }
981
982            async fn rename(
983                &self,
984                from: &str,
985                to: &str,
986                args: opendal::raw::OpRename,
987            ) -> opendal::Result<opendal::raw::RpRename> {
988                self.inner.rename(from, to, args).await
989            }
990        }
991    }
992
993    #[tokio::test]
994    async fn test_get_ranges_no_stat() {
995        use std::sync::atomic::{AtomicUsize, Ordering};
996
997        // Create a stat counter and operator with tracking layer
998        let stat_count = Arc::new(AtomicUsize::new(0));
999        let op = Operator::new(opendal::services::Memory::default())
1000            .unwrap()
1001            .layer(stat_counter::StatCounterLayer::new(stat_count.clone()))
1002            .finish();
1003        let store = OpendalStore::new(op);
1004
1005        // Create a test file
1006        let location = "test_get_range.txt".into();
1007        let value = Bytes::from_static(b"Hello, world!");
1008        store.put(&location, value.clone().into()).await.unwrap();
1009
1010        // Reset counter after put
1011        stat_count.store(0, Ordering::SeqCst);
1012
1013        // Test 1: get_ranges should NOT call stat()
1014        let range = 0..5;
1015        let ret = store
1016            .get_ranges(&location, std::slice::from_ref(&range))
1017            .await
1018            .unwrap();
1019        assert_eq!(vec![Bytes::from_static(b"Hello")], ret);
1020        assert_eq!(
1021            stat_count.load(Ordering::SeqCst),
1022            0,
1023            "get_ranges should not call stat()"
1024        );
1025
1026        // Reset counter
1027        stat_count.store(0, Ordering::SeqCst);
1028
1029        // Test 2: get_opts SHOULD call stat() to get metadata
1030        let opts = datafusion_object_store::GetOptions {
1031            range: Some(datafusion_object_store::GetRange::Bounded(0..5)),
1032            ..Default::default()
1033        };
1034        let ret = store.get_opts(&location, opts).await.unwrap();
1035        let data = ret.bytes().await.unwrap();
1036        assert_eq!(Bytes::from_static(b"Hello"), data);
1037        assert!(
1038            stat_count.load(Ordering::SeqCst) > 0,
1039            "get_opts should call stat() to get metadata"
1040        );
1041
1042        // Cleanup
1043        store.delete(&location).await.unwrap();
1044    }
1045}