Skip to main content

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::{OpCopier, 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    async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
135        self.inner.delete(path, args).await
136    }
137
138    async fn close(&mut self) -> Result<()> {
139        self.inner.close().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    type Copier = oio::Copier;
150
151    fn inner(&self) -> &Self::Inner {
152        &self.inner
153    }
154
155    async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
156        if let Some(reader_factory) = self.reader_factory.as_ref() {
157            let (rp_read, reader) = self.inner.read(path, args.clone()).await?;
158            let reader = reader_factory(path, args, Box::new(reader));
159            Ok((rp_read, MockReader { inner: reader }))
160        } else {
161            self.inner.read(path, args).await.map(|(rp_read, reader)| {
162                (
163                    rp_read,
164                    MockReader {
165                        inner: Box::new(reader),
166                    },
167                )
168            })
169        }
170    }
171
172    async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
173        if let Some(writer_factory) = self.writer_factory.as_ref() {
174            let (rp_write, writer) = self.inner.write(path, args.clone()).await?;
175            let writer = writer_factory(path, args, Box::new(writer));
176            Ok((rp_write, MockWriter { inner: writer }))
177        } else {
178            self.inner
179                .write(path, args)
180                .await
181                .map(|(rp_write, writer)| {
182                    (
183                        rp_write,
184                        MockWriter {
185                            inner: Box::new(writer),
186                        },
187                    )
188                })
189        }
190    }
191
192    async fn delete(&self) -> Result<(RpDelete, Self::Deleter)> {
193        if let Some(deleter_factory) = self.deleter_factory.as_ref() {
194            let (rp_delete, deleter) = self.inner.delete().await?;
195            let deleter = deleter_factory(Box::new(deleter));
196            Ok((rp_delete, MockDeleter { inner: deleter }))
197        } else {
198            self.inner.delete().await.map(|(rp_delete, deleter)| {
199                (
200                    rp_delete,
201                    MockDeleter {
202                        inner: Box::new(deleter),
203                    },
204                )
205            })
206        }
207    }
208
209    async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Lister)> {
210        if let Some(lister_factory) = self.lister_factory.as_ref() {
211            let (rp_list, lister) = self.inner.list(path, args.clone()).await?;
212            let lister = lister_factory(path, args, Box::new(lister));
213            Ok((rp_list, MockLister { inner: lister }))
214        } else {
215            self.inner.list(path, args).await.map(|(rp_list, lister)| {
216                (
217                    rp_list,
218                    MockLister {
219                        inner: Box::new(lister),
220                    },
221                )
222            })
223        }
224    }
225
226    async fn copy(
227        &self,
228        from: &str,
229        to: &str,
230        args: OpCopy,
231        opts: OpCopier,
232    ) -> Result<(RpCopy, Self::Copier)> {
233        if let Some(result) = self
234            .copy_interceptor
235            .as_ref()
236            .and_then(|copy_interceptor| copy_interceptor(from, to, args.clone()))
237        {
238            return result.map(|rp_copy| (rp_copy, Box::new(()) as oio::Copier));
239        }
240
241        self.inner
242            .copy(from, to, args, opts)
243            .await
244            .map(|(rp_copy, copier)| (rp_copy, Box::new(copier) as oio::Copier))
245    }
246}