Skip to main content

iota_storage/object_store/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::sync::Arc;
6
7use anyhow::{Result, anyhow};
8use async_trait::async_trait;
9use bytes::Bytes;
10use futures::stream::BoxStream;
11use object_store::{DynObjectStore, ObjectMeta, ObjectStore, ObjectStoreExt, path::Path};
12pub mod http;
13pub mod util;
14
15#[async_trait]
16pub trait ObjectStoreGetExt: std::fmt::Display + Send + Sync + 'static {
17    /// Return the bytes at given path in object store
18    async fn get_bytes(&self, src: &Path) -> Result<Bytes>;
19
20    /// Return whether an object exists at the given path.
21    async fn exists(&self, src: &Path) -> Result<bool>;
22}
23
24macro_rules! as_ref_get_ext_impl {
25    ($type:ty) => {
26        #[async_trait]
27        impl ObjectStoreGetExt for $type {
28            async fn get_bytes(&self, src: &Path) -> Result<Bytes> {
29                self.as_ref().get_bytes(src).await
30            }
31
32            async fn exists(&self, src: &Path) -> Result<bool> {
33                self.as_ref().exists(src).await
34            }
35        }
36    };
37}
38
39as_ref_get_ext_impl!(Arc<dyn ObjectStoreGetExt>);
40as_ref_get_ext_impl!(Box<dyn ObjectStoreGetExt>);
41
42macro_rules! as_ref_get_impl {
43    ($type:ty) => {
44        #[async_trait]
45        impl ObjectStoreGetExt for $type {
46            async fn get_bytes(&self, src: &Path) -> Result<Bytes> {
47                self.get(src)
48                    .await
49                    .map_err(|e| anyhow!("Failed to get file {src} with error: {e:?}"))?
50                    .bytes()
51                    .await
52                    .map_err(|e| {
53                        anyhow!(
54                            "Failed to collect GET result for file {src} into bytes with error: {e:?}")
55                    })
56            }
57
58            async fn exists(&self, src: &Path) -> Result<bool> {
59                match self.head(src).await {
60                    Ok(_) => Ok(true),
61                    Err(object_store::Error::NotFound { .. }) => Ok(false),
62                    Err(e) => Err(anyhow!("Failed to check if file {src} exists with error: {e:?}")),
63                }
64            }
65        }
66    };
67}
68
69as_ref_get_impl!(Arc<dyn ObjectStore>);
70as_ref_get_impl!(Box<dyn ObjectStore>);
71
72#[async_trait]
73pub trait ObjectStoreListExt: Send + Sync + 'static {
74    /// List the objects at the given path in object store
75    async fn list_objects(
76        &self,
77        src: Option<&Path>,
78    ) -> BoxStream<'_, object_store::Result<ObjectMeta>>;
79}
80
81macro_rules! as_ref_list_ext_impl {
82    ($type:ty) => {
83        #[async_trait]
84        impl ObjectStoreListExt for $type {
85            async fn list_objects(
86                &self,
87                src: Option<&Path>,
88            ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
89                self.as_ref().list_objects(src).await
90            }
91        }
92    };
93}
94
95as_ref_list_ext_impl!(Arc<dyn ObjectStoreListExt>);
96as_ref_list_ext_impl!(Box<dyn ObjectStoreListExt>);
97
98#[async_trait]
99impl ObjectStoreListExt for Arc<DynObjectStore> {
100    async fn list_objects(
101        &self,
102        src: Option<&Path>,
103    ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
104        self.list(src)
105    }
106}
107
108#[async_trait]
109pub trait ObjectStorePutExt: Send + Sync + 'static {
110    /// Write the bytes at the given location in object store
111    async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()>;
112}
113
114macro_rules! as_ref_put_ext_impl {
115    ($type:ty) => {
116        #[async_trait]
117        impl ObjectStorePutExt for $type {
118            async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()> {
119                self.as_ref().put_bytes(src, bytes).await
120            }
121        }
122    };
123}
124
125as_ref_put_ext_impl!(Arc<dyn ObjectStorePutExt>);
126as_ref_put_ext_impl!(Box<dyn ObjectStorePutExt>);
127
128#[async_trait]
129impl ObjectStorePutExt for Arc<DynObjectStore> {
130    async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()> {
131        self.put(src, bytes.into()).await?;
132        Ok(())
133    }
134}
135
136#[async_trait]
137pub trait ObjectStoreDeleteExt: Send + Sync + 'static {
138    /// Delete the object at the given location in object store
139    async fn delete_object(&self, src: &Path) -> Result<()>;
140}
141
142macro_rules! as_ref_delete_ext_impl {
143    ($type:ty) => {
144        #[async_trait]
145        impl ObjectStoreDeleteExt for $type {
146            async fn delete_object(&self, src: &Path) -> Result<()> {
147                self.as_ref().delete_object(src).await
148            }
149        }
150    };
151}
152
153as_ref_delete_ext_impl!(Arc<dyn ObjectStoreDeleteExt>);
154as_ref_delete_ext_impl!(Box<dyn ObjectStoreDeleteExt>);
155
156#[async_trait]
157
158impl ObjectStoreDeleteExt for Arc<DynObjectStore> {
159    async fn delete_object(&self, src: &Path) -> Result<()> {
160        self.delete(src).await?;
161        Ok(())
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use std::sync::Arc;
168
169    use object_store::{ObjectStore, memory::InMemory, path::Path};
170
171    use crate::object_store::{ObjectStoreGetExt, ObjectStorePutExt};
172
173    #[tokio::test]
174    async fn test_dyn_object_store_exists() -> anyhow::Result<()> {
175        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
176        let path = Path::from("file1");
177        store
178            .put_bytes(&path, bytes::Bytes::from_static(b"Lorem ipsum"))
179            .await?;
180
181        assert!(store.exists(&path).await?);
182        assert!(!store.exists(&Path::from("missing")).await?);
183        Ok(())
184    }
185}