vv_agent/workspace/s3/
backend.rs

1use std::any::Any;
2use std::io::{Error, ErrorKind};
3use std::sync::Arc;
4
5use futures_util::TryStreamExt;
6use object_store::aws::AmazonS3Builder;
7use object_store::memory::InMemory;
8use object_store::path::Path as ObjectPath;
9use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
10use tokio::runtime::Runtime;
11
12use crate::workspace::{
13    glob_match, non_empty_option, normalize_workspace_path, normalized_glob_pattern,
14    object_store_error_to_io, suffix_with_dot, FileInfo, WorkspaceBackend,
15};
16
17use super::config::S3WorkspaceConfig;
18use super::paths::{list_prefix, object_key, relative_key};
19use super::runtime::{block_on_object_store, build_runtime};
20
21#[derive(Clone)]
22pub struct S3WorkspaceBackend {
23    store: Arc<dyn ObjectStore>,
24    prefix: String,
25    runtime: Arc<Runtime>,
26}
27
28impl std::fmt::Debug for S3WorkspaceBackend {
29    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        formatter
31            .debug_struct("S3WorkspaceBackend")
32            .field("store", &self.store.to_string())
33            .field("prefix", &self.prefix)
34            .finish_non_exhaustive()
35    }
36}
37
38impl Default for S3WorkspaceBackend {
39    fn default() -> Self {
40        Self::from_object_store(InMemory::new(), "")
41            .expect("in-memory S3 workspace backend must initialize")
42    }
43}
44
45impl S3WorkspaceBackend {
46    pub fn new(bucket: impl Into<String>) -> std::io::Result<Self> {
47        Self::from_config(S3WorkspaceConfig::new(bucket))
48    }
49
50    pub fn from_config(config: S3WorkspaceConfig) -> std::io::Result<Self> {
51        let bucket = config.bucket.trim();
52        if bucket.is_empty() {
53            return Err(Error::new(
54                ErrorKind::InvalidInput,
55                "S3 bucket cannot be empty",
56            ));
57        }
58        let mut builder = AmazonS3Builder::from_env().with_bucket_name(bucket.to_string());
59        if let Some(endpoint) = non_empty_option(config.endpoint_url) {
60            if endpoint.starts_with("http://") {
61                builder = builder.with_allow_http(true);
62            }
63            builder = builder.with_endpoint(endpoint);
64        }
65        if let Some(region) = non_empty_option(config.region_name) {
66            builder = builder.with_region(region);
67        }
68        if let Some(access_key_id) = non_empty_option(config.aws_access_key_id) {
69            builder = builder.with_access_key_id(access_key_id);
70        }
71        if let Some(secret_access_key) = non_empty_option(config.aws_secret_access_key) {
72            builder = builder.with_secret_access_key(secret_access_key);
73        }
74        if let Some(token) = non_empty_option(config.aws_session_token) {
75            builder = builder.with_token(token);
76        }
77        let virtual_hosted = match config.addressing_style.trim().to_ascii_lowercase().as_str() {
78            "" | "virtual" => true,
79            "path" => false,
80            other => {
81                return Err(Error::new(
82                    ErrorKind::InvalidInput,
83                    format!("unsupported S3 addressing_style: {other}"),
84                ))
85            }
86        };
87        builder = builder.with_virtual_hosted_style_request(virtual_hosted);
88        let store = builder.build().map_err(object_store_error_to_io)?;
89        Self::from_object_store(store, config.prefix)
90    }
91
92    pub fn from_object_store(
93        store: impl ObjectStore + 'static,
94        prefix: impl Into<String>,
95    ) -> std::io::Result<Self> {
96        Ok(Self {
97            store: Arc::new(store),
98            prefix: normalize_workspace_path(&prefix.into()),
99            runtime: Arc::new(build_runtime()?),
100        })
101    }
102
103    fn object_key(&self, path: &str) -> String {
104        object_key(&self.prefix, path)
105    }
106
107    fn relative_key(&self, key: &str) -> Option<String> {
108        relative_key(&self.prefix, key)
109    }
110
111    fn list_prefix(&self) -> Option<ObjectPath> {
112        list_prefix(&self.prefix)
113    }
114
115    fn block_on<T>(
116        &self,
117        future: impl std::future::Future<Output = object_store::Result<T>>,
118    ) -> std::io::Result<T> {
119        block_on_object_store(&self.runtime, future)
120    }
121}
122
123impl WorkspaceBackend for S3WorkspaceBackend {
124    fn as_any(&self) -> &dyn Any {
125        self
126    }
127
128    fn list_files(&self, base: &str, glob: &str) -> std::io::Result<Vec<String>> {
129        let base = normalize_workspace_path(base);
130        let pattern = if base.is_empty() {
131            normalized_glob_pattern(glob)
132        } else {
133            format!("{base}/{}", normalized_glob_pattern(glob))
134        };
135        let prefix = self.list_prefix();
136        let objects = self.block_on(async {
137            self.store
138                .list(prefix.as_ref())
139                .try_collect::<Vec<_>>()
140                .await
141        })?;
142        let mut files = objects
143            .into_iter()
144            .filter_map(|object| self.relative_key(object.location.as_ref()))
145            .filter(|path| !path.is_empty() && !path.ends_with('/'))
146            .filter(|path| glob_match(path, &pattern))
147            .collect::<Vec<_>>();
148        files.sort();
149        Ok(files)
150    }
151
152    fn read_text(&self, path: &str) -> std::io::Result<String> {
153        let bytes = self.read_bytes(path)?;
154        Ok(String::from_utf8_lossy(&bytes).to_string())
155    }
156
157    fn read_bytes(&self, path: &str) -> std::io::Result<Vec<u8>> {
158        let key = ObjectPath::from(self.object_key(path));
159        self.block_on(async { Ok(self.store.get(&key).await?.bytes().await?.to_vec()) })
160    }
161
162    fn write_text(&self, path: &str, content: &str, append: bool) -> std::io::Result<usize> {
163        let key = ObjectPath::from(self.object_key(path));
164        let content = if append {
165            match self.read_text(path) {
166                Ok(existing) => existing + content,
167                Err(error) if error.kind() == ErrorKind::NotFound => content.to_string(),
168                Err(error) => return Err(error),
169            }
170        } else {
171            content.to_string()
172        };
173        let len = content.len();
174        self.block_on(async {
175            self.store
176                .put(&key, PutPayload::from(content.into_bytes()))
177                .await
178                .map(|_| ())
179        })?;
180        Ok(len)
181    }
182
183    fn file_info(&self, path: &str) -> std::io::Result<Option<FileInfo>> {
184        let key = ObjectPath::from(self.object_key(path));
185        let metadata = match self.block_on(async { self.store.head(&key).await }) {
186            Ok(metadata) => metadata,
187            Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None),
188            Err(error) => return Err(error),
189        };
190        let object_path = metadata.location.as_ref();
191        Ok(Some(FileInfo {
192            path: self
193                .relative_key(object_path)
194                .unwrap_or_else(|| normalize_workspace_path(path)),
195            is_file: true,
196            is_dir: false,
197            size: metadata.size,
198            modified_at: metadata.last_modified.to_rfc3339(),
199            suffix: suffix_with_dot(path),
200        }))
201    }
202
203    fn exists(&self, path: &str) -> bool {
204        self.file_info(path).ok().flatten().is_some()
205    }
206
207    fn is_file(&self, path: &str) -> bool {
208        self.exists(path)
209    }
210
211    fn mkdir(&self, _path: &str) -> std::io::Result<()> {
212        Ok(())
213    }
214}
vv_agent/workspace/s3/backend.rs

vv_agent/workspace/s3/
backend.rs
