puffin/puffin_manager/
file_accessor.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 async_trait::async_trait;
16use common_base::range_read::{FileReader, SizeAwareRangeReader};
17use common_test_util::temp_dir::{TempDir, create_temp_dir};
18use futures::AsyncWrite;
19use tokio::fs::File;
20use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
21
22use crate::error::Result;
23
24/// `PuffinFileAccessor` is for opening readers and writers for puffin files.
25#[async_trait]
26#[auto_impl::auto_impl(Arc)]
27pub trait PuffinFileAccessor: Send + Sync + 'static {
28    type Reader: SizeAwareRangeReader + Sync;
29    type Writer: AsyncWrite + Unpin + Send;
30    type FileHandle: ToString + Clone + Send + Sync;
31
32    /// Opens a reader for the given puffin file handle.
33    async fn reader(&self, handle: &Self::FileHandle) -> Result<Self::Reader>;
34
35    /// Creates a writer for the given puffin file handle.
36    async fn writer(&self, handle: &Self::FileHandle) -> Result<Self::Writer>;
37}
38
39pub struct MockFileAccessor {
40    tempdir: TempDir,
41}
42
43impl MockFileAccessor {
44    pub fn new(prefix: &str) -> Self {
45        let tempdir = create_temp_dir(prefix);
46        Self { tempdir }
47    }
48}
49
50#[async_trait]
51impl PuffinFileAccessor for MockFileAccessor {
52    type Reader = FileReader;
53    type Writer = Compat<File>;
54    type FileHandle = String;
55
56    async fn reader(&self, handle: &String) -> Result<Self::Reader> {
57        Ok(FileReader::new(self.tempdir.path().join(handle))
58            .await
59            .unwrap())
60    }
61
62    async fn writer(&self, handle: &String) -> Result<Self::Writer> {
63        let p = self.tempdir.path().join(handle);
64        if let Some(p) = p.parent()
65            && !tokio::fs::try_exists(p).await.unwrap()
66        {
67            tokio::fs::create_dir_all(p).await.unwrap();
68        }
69        let f = tokio::fs::File::create(p).await.unwrap();
70        Ok(f.compat())
71    }
72}
puffin/puffin_manager/file_accessor.rs

puffin/puffin_manager/
file_accessor.rs
