1use 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#[derive(Clone)]
84pub struct OpendalStore {
85 info: Arc<OperatorInfo>,
86 inner: Operator,
87}
88
89impl OpendalStore {
90 pub fn new(op: Operator) -> Self {
92 Self {
93 info: op.info().into(),
94 inner: op,
95 }
96 }
97
98 pub fn info(&self) -> &OperatorInfo {
100 self.info.as_ref()
101 }
102
103 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 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 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 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 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 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
582struct 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 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 let (tx, rx) = oneshot::channel();
617 let last_rx = std::mem::replace(&mut self.next_notify, rx);
619
620 async move {
621 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 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 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 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 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 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 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 stat_count.store(0, Ordering::SeqCst);
1012
1013 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 stat_count.store(0, Ordering::SeqCst);
1028
1029 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 store.delete(&location).await.unwrap();
1044 }
1045}