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 std::sync::Arc;
16
17use async_trait::async_trait;
18use common_base::range_read::{FileReader, SizeAwareRangeReader};
19use common_test_util::temp_dir::{TempDir, create_temp_dir};
20use futures::AsyncWrite;
21use tokio::fs::File;
22use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
23
24use crate::error::Result;
25use crate::puffin_manager::file_key::{PuffinFileKey, PuffinStorageScope};
26
27/// `PuffinFileAccessor` is for opening readers and writers for puffin files.
28#[async_trait]
29#[auto_impl::auto_impl(Arc)]
30pub trait PuffinFileAccessor: Send + Sync + 'static {
31    type Reader: SizeAwareRangeReader + Sync;
32    type Writer: AsyncWrite + Unpin + Send;
33    type FileHandle: ToString + Clone + Send + Sync;
34
35    /// Returns stable identity using the same resolved path as reader and writer.
36    /// Keys include the physical file version; this operation performs no I/O.
37    fn file_key(&self, handle: &Self::FileHandle) -> PuffinFileKey<Self::FileHandle>;
38
39    /// Opens a reader for the given puffin file handle.
40    async fn reader(&self, handle: &Self::FileHandle) -> Result<Self::Reader>;
41
42    /// Creates a writer for the given puffin file handle.
43    async fn writer(&self, handle: &Self::FileHandle) -> Result<Self::Writer>;
44}
45
46pub struct MockFileAccessor {
47    tempdir: TempDir,
48    scope: Arc<PuffinStorageScope>,
49}
50
51impl MockFileAccessor {
52    pub fn new(prefix: &str) -> Self {
53        let tempdir = create_temp_dir(prefix);
54        let scope = Arc::new(PuffinStorageScope::new(
55            "file".to_string(),
56            String::new(),
57            tempdir.path().to_string_lossy().into_owned(),
58        ));
59        Self { tempdir, scope }
60    }
61}
62
63#[async_trait]
64impl PuffinFileAccessor for MockFileAccessor {
65    type Reader = FileReader;
66    type Writer = Compat<File>;
67    type FileHandle = String;
68
69    fn file_key(&self, handle: &Self::FileHandle) -> PuffinFileKey<Self::FileHandle> {
70        PuffinFileKey::new(self.scope.clone(), handle.clone(), handle.clone())
71    }
72
73    async fn reader(&self, handle: &String) -> Result<Self::Reader> {
74        Ok(FileReader::new(self.tempdir.path().join(handle))
75            .await
76            .unwrap())
77    }
78
79    async fn writer(&self, handle: &String) -> Result<Self::Writer> {
80        let p = self.tempdir.path().join(handle);
81        if let Some(p) = p.parent()
82            && !tokio::fs::try_exists(p).await.unwrap()
83        {
84            tokio::fs::create_dir_all(p).await.unwrap();
85        }
86        let f = tokio::fs::File::create(p).await.unwrap();
87        Ok(f.compat())
88    }
89}
puffin/puffin_manager/file_accessor.rs

puffin/puffin_manager/
file_accessor.rs
