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

vv_agent/workspace/s3/
backend.rs
