Skip to main content

orchestral_runtime/bootstrap/
blob_store.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use bytes::Bytes;
5use futures_util::StreamExt;
6use tokio::sync::RwLock;
7
8use orchestral_core::io::{
9    BlobHead, BlobId, BlobIoError, BlobMeta, BlobRead, BlobStore, BlobWriteRequest,
10};
11
12#[derive(Clone)]
13struct InMemoryBlobObject {
14    meta: BlobMeta,
15    bytes: Vec<u8>,
16}
17
18#[derive(Default)]
19pub struct InMemoryBlobStore {
20    objects: RwLock<HashMap<String, InMemoryBlobObject>>,
21}
22
23#[async_trait]
24impl BlobStore for InMemoryBlobStore {
25    async fn write(&self, mut request: BlobWriteRequest) -> Result<BlobMeta, BlobIoError> {
26        let blob_id = BlobId::from(uuid::Uuid::new_v4().to_string());
27        let mut data: Vec<u8> = Vec::new();
28        while let Some(chunk) = request.body.next().await {
29            let chunk = chunk?;
30            data.extend_from_slice(&chunk);
31        }
32        if data.is_empty() {
33            return Err(BlobIoError::Invalid("empty blob payload".to_string()));
34        }
35        let now = chrono::Utc::now();
36        let meta = BlobMeta {
37            id: blob_id,
38            file_name: request.file_name.take(),
39            mime_type: request.mime_type.take(),
40            byte_size: data.len() as u64,
41            checksum_sha256: None,
42            metadata: if request.metadata.is_null() {
43                serde_json::json!({})
44            } else {
45                request.metadata
46            },
47            created_at: now,
48            updated_at: now,
49        };
50        self.objects.write().await.insert(
51            meta.id.to_string(),
52            InMemoryBlobObject {
53                meta: meta.clone(),
54                bytes: data,
55            },
56        );
57        Ok(meta)
58    }
59
60    async fn read(&self, blob_id: &BlobId) -> Result<BlobRead, BlobIoError> {
61        let obj = self
62            .objects
63            .read()
64            .await
65            .get(blob_id.as_str())
66            .cloned()
67            .ok_or_else(|| BlobIoError::NotFound(blob_id.to_string()))?;
68        let body = Box::pin(futures_util::stream::once(async move {
69            Ok(Bytes::from(obj.bytes))
70        }));
71        Ok(BlobRead {
72            meta: obj.meta,
73            body,
74        })
75    }
76
77    async fn head(&self, blob_id: &BlobId) -> Result<BlobHead, BlobIoError> {
78        let obj = self
79            .objects
80            .read()
81            .await
82            .get(blob_id.as_str())
83            .cloned()
84            .ok_or_else(|| BlobIoError::NotFound(blob_id.to_string()))?;
85        Ok(BlobHead {
86            byte_size: obj.meta.byte_size,
87            etag: None,
88            last_modified: Some(obj.meta.updated_at),
89        })
90    }
91
92    async fn delete(&self, blob_id: &BlobId) -> Result<bool, BlobIoError> {
93        Ok(self
94            .objects
95            .write()
96            .await
97            .remove(blob_id.as_str())
98            .is_some())
99    }
100}