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

object_store/layers/
mock.rs
