object_store/layers/
mock.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::fmt::Debug;
16use std::sync::Arc;
17
18use derive_builder::Builder;
19pub use oio::*;
20pub use opendal::raw::{
21    Access, Layer, LayeredAccess, OpDelete, OpList, OpRead, OpWrite, RpDelete, RpList, RpRead,
22    RpWrite, oio,
23};
24use opendal::raw::{OpCopy, RpCopy};
25pub use opendal::{Buffer, Error, ErrorKind, Metadata, Result};
26
27pub type MockWriterFactory = Arc<dyn Fn(&str, OpWrite, oio::Writer) -> oio::Writer + Send + Sync>;
28pub type MockReaderFactory = Arc<dyn Fn(&str, OpRead, oio::Reader) -> oio::Reader + Send + Sync>;
29pub type MockListerFactory = Arc<dyn Fn(&str, OpList, oio::Lister) -> oio::Lister + Send + Sync>;
30pub type MockDeleterFactory = Arc<dyn Fn(oio::Deleter) -> oio::Deleter + Send + Sync>;
31pub type CopyInterceptor = Arc<dyn Fn(&str, &str, OpCopy) -> Option<Result<RpCopy>> + Send + Sync>;
32
33#[derive(Builder)]
34pub struct MockLayer {
35    #[builder(setter(strip_option), default)]
36    writer_factory: Option<MockWriterFactory>,
37    #[builder(setter(strip_option), default)]
38    reader_factory: Option<MockReaderFactory>,
39    #[builder(setter(strip_option), default)]
40    lister_factory: Option<MockListerFactory>,
41    #[builder(setter(strip_option), default)]
42    deleter_factory: Option<MockDeleterFactory>,
43    #[builder(setter(strip_option), default)]
44    copy_interceptor: Option<CopyInterceptor>,
45}
46
47impl Clone for MockLayer {
48    fn clone(&self) -> Self {
49        Self {
50            writer_factory: self.writer_factory.clone(),
51            reader_factory: self.reader_factory.clone(),
52            lister_factory: self.lister_factory.clone(),
53            deleter_factory: self.deleter_factory.clone(),
54            copy_interceptor: self.copy_interceptor.clone(),
55        }
56    }
57}
58
59impl<A: Access> Layer<A> for MockLayer {
60    type LayeredAccess = MockAccessor<A>;
61
62    fn layer(&self, inner: A) -> Self::LayeredAccess {
63        MockAccessor {
64            inner,
65            writer_factory: self.writer_factory.clone(),
66            reader_factory: self.reader_factory.clone(),
67            lister_factory: self.lister_factory.clone(),
68            deleter_factory: self.deleter_factory.clone(),
69            copy_interceptor: self.copy_interceptor.clone(),
70        }
71    }
72}
73
74pub struct MockAccessor<A> {
75    inner: A,
76    writer_factory: Option<MockWriterFactory>,
77    reader_factory: Option<MockReaderFactory>,
78    lister_factory: Option<MockListerFactory>,
79    deleter_factory: Option<MockDeleterFactory>,
80    copy_interceptor: Option<CopyInterceptor>,
81}
82
83impl<A: Debug> Debug for MockAccessor<A> {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        f.debug_struct("MockAccessor")
86            .field("inner", &self.inner)
87            .finish()
88    }
89}
90
91pub struct MockReader {
92    inner: oio::Reader,
93}
94
95impl oio::Read for MockReader {
96    async fn read(&mut self) -> Result<Buffer> {
97        self.inner.read().await
98    }
99}
100
101pub struct MockWriter {
102    inner: oio::Writer,
103}
104
105impl oio::Write for MockWriter {
106    async fn write(&mut self, bs: Buffer) -> Result<()> {
107        self.inner.write(bs).await
108    }
109
110    async fn close(&mut self) -> Result<Metadata> {
111        self.inner.close().await
112    }
113
114    async fn abort(&mut self) -> Result<()> {
115        self.inner.abort().await
116    }
117}
118
119pub struct MockLister {
120    inner: oio::Lister,
121}
122
123impl oio::List for MockLister {
124    async fn next(&mut self) -> Result<Option<oio::Entry>> {
125        self.inner.next().await
126    }
127}
128
129pub struct MockDeleter {
130    inner: oio::Deleter,
131}
132
133impl oio::Delete for MockDeleter {
134    fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
135        self.inner.delete(path, args)
136    }
137
138    async fn flush(&mut self) -> Result<usize> {
139        self.inner.flush().await
140    }
141}
142
143impl<A: Access> LayeredAccess for MockAccessor<A> {
144    type Inner = A;
145    type Reader = MockReader;
146    type Writer = MockWriter;
147    type Lister = MockLister;
148    type Deleter = MockDeleter;
149
150    fn inner(&self) -> &Self::Inner {
151        &self.inner
152    }
153
154    async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
155        if let Some(reader_factory) = self.reader_factory.as_ref() {
156            let (rp_read, reader) = self.inner.read(path, args.clone()).await?;
157            let reader = reader_factory(path, args, Box::new(reader));
158            Ok((rp_read, MockReader { inner: reader }))
159        } else {
160            self.inner.read(path, args).await.map(|(rp_read, reader)| {
161                (
162                    rp_read,
163                    MockReader {
164                        inner: Box::new(reader),
165                    },
166                )
167            })
168        }
169    }
170
171    async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
172        if let Some(writer_factory) = self.writer_factory.as_ref() {
173            let (rp_write, writer) = self.inner.write(path, args.clone()).await?;
174            let writer = writer_factory(path, args, Box::new(writer));
175            Ok((rp_write, MockWriter { inner: writer }))
176        } else {
177            self.inner
178                .write(path, args)
179                .await
180                .map(|(rp_write, writer)| {
181                    (
182                        rp_write,
183                        MockWriter {
184                            inner: Box::new(writer),
185                        },
186                    )
187                })
188        }
189    }
190
191    async fn delete(&self) -> Result<(RpDelete, Self::Deleter)> {
192        if let Some(deleter_factory) = self.deleter_factory.as_ref() {
193            let (rp_delete, deleter) = self.inner.delete().await?;
194            let deleter = deleter_factory(Box::new(deleter));
195            Ok((rp_delete, MockDeleter { inner: deleter }))
196        } else {
197            self.inner.delete().await.map(|(rp_delete, deleter)| {
198                (
199                    rp_delete,
200                    MockDeleter {
201                        inner: Box::new(deleter),
202                    },
203                )
204            })
205        }
206    }
207
208    async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Lister)> {
209        if let Some(lister_factory) = self.lister_factory.as_ref() {
210            let (rp_list, lister) = self.inner.list(path, args.clone()).await?;
211            let lister = lister_factory(path, args, Box::new(lister));
212            Ok((rp_list, MockLister { inner: lister }))
213        } else {
214            self.inner.list(path, args).await.map(|(rp_list, lister)| {
215                (
216                    rp_list,
217                    MockLister {
218                        inner: Box::new(lister),
219                    },
220                )
221            })
222        }
223    }
224
225    async fn copy(&self, from: &str, to: &str, args: OpCopy) -> Result<RpCopy> {
226        let Some(copy_interceptor) = self.copy_interceptor.as_ref() else {
227            return self.inner.copy(from, to, args).await;
228        };
229
230        let Some(result) = copy_interceptor(from, to, args.clone()) else {
231            return self.inner.copy(from, to, args).await;
232        };
233
234        result
235    }
236}
object_store/layers/mock.rs

object_store/layers/
mock.rs
