refactor: accept blob id from caller

toeverything · Dec 25, 2023 · 971a932 · 971a932
1 parent 49a6b7a
commit 971a932
Show file tree

Hide file tree

Showing 9 changed files with 23 additions and 60 deletions.
diff --git a/apps/keck/src/server/api/blobs.rs b/apps/keck/src/server/api/blobs.rs
@@ -102,10 +102,11 @@ pub async fn set_blob(
         })
         .filter_map(|data| future::ready(data.ok()));
 
+    let (id, blob) = stream_to_blob(body).await;
     if let Ok(id) = context
         .storage
         .blobs()
-        .put_blob_stream(Some(workspace.clone()), body)
+        .put_blob(Some(workspace.clone()), id, blob)
         .await
     {
         if has_error {

diff --git a/apps/keck/src/server/api/mod.rs b/apps/keck/src/server/api/mod.rs
@@ -16,7 +16,7 @@ use axum::{
 };
 use doc::doc_apis;
 use jwst_rpc::{BroadcastChannels, RpcContextImpl};
-use jwst_storage::{BlobStorageType, JwstStorage, JwstStorageResult};
+use jwst_storage::{stream_to_blob, BlobStorageType, JwstStorage, JwstStorageResult};
 use tokio::sync::RwLock;
 
 use super::*;

diff --git a/apps/keck/src/server/sync/blobs.rs b/apps/keck/src/server/sync/blobs.rs
@@ -94,12 +94,9 @@ impl Context {
             })
             .filter_map(|data| future::ready(data.ok()));
 
-        if let Ok(id) = self
-            .get_storage()
-            .blobs()
-            .put_blob_stream(workspace.clone(), stream)
-            .await
-        {
+        let (id, blob) = stream_to_blob(stream).await;
+
+        if let Ok(id) = self.get_storage().blobs().put_blob(workspace.clone(), id, blob).await {
             if has_error {
                 let _ = self.get_storage().blobs().delete_blob(workspace, id).await;
                 StatusCode::INTERNAL_SERVER_ERROR.into_response()

diff --git a/apps/keck/src/server/utils.rs b/apps/keck/src/server/utils.rs
@@ -1,3 +1,4 @@
 pub use jwst_logger::{debug, error, info, warn};
+pub use jwst_storage::stream_to_blob;
 pub use nanoid::nanoid;
 pub use serde::{Deserialize, Serialize};
diff --git a/libs/jwst-core/src/types/blob.rs b/libs/jwst-core/src/types/blob.rs
@@ -1,8 +1,6 @@
 use std::collections::HashMap;
 
-use bytes::Bytes;
 use chrono::NaiveDateTime;
-use futures::Stream;
 
 use super::*;
 
@@ -29,12 +27,7 @@ pub trait BlobStorage<E = JwstError> {
         id: String,
         params: Option<HashMap<String, String>>,
     ) -> JwstResult<BlobMetadata, E>;
-    async fn put_blob_stream(
-        &self,
-        workspace: Option<String>,
-        stream: impl Stream<Item = Bytes> + Send,
-    ) -> JwstResult<String, E>;
-    async fn put_blob(&self, workspace: Option<String>, blob: Vec<u8>) -> JwstResult<String, E>;
+    async fn put_blob(&self, workspace: Option<String>, id: String, blob: Vec<u8>) -> JwstResult<String, E>;
     async fn delete_blob(&self, workspace: Option<String>, id: String) -> JwstResult<bool, E>;
     async fn delete_workspace(&self, workspace_id: String) -> JwstResult<(), E>;
     async fn get_blobs_size(&self, workspaces: Vec<String>) -> JwstResult<i64, E>;

diff --git a/libs/jwst-storage/src/lib.rs b/libs/jwst-storage/src/lib.rs
@@ -8,15 +8,18 @@ use std::{path::PathBuf, sync::Arc, time::Duration};
 
 use async_trait::async_trait;
 use chrono::Utc;
-use futures::{Future, Stream};
+use futures::Future;
 use jwst_core::{DocStorage, JwstResult, Workspace};
 use jwst_logger::{debug, error, info, trace, warn};
 use path_ext::PathExt;
 use rate_limiter::{get_bucket, is_sqlite, Bucket};
 use sea_orm::{prelude::*, ConnectOptions, Database, DbErr, QuerySelect, Set};
 #[cfg(feature = "bucket")]
 pub use storage::blobs::MixedBucketDBParam;
-pub use storage::{blobs::BlobStorageType, JwstStorage};
+pub use storage::{
+    blobs::{stream_to_blob, BlobStorageType},
+    JwstStorage,
+};
 pub use types::{JwstStorageError, JwstStorageResult};
 
 #[inline]

diff --git a/libs/jwst-storage/src/storage/blobs/blob_storage.rs b/libs/jwst-storage/src/storage/blobs/blob_storage.rs
@@ -1,6 +1,4 @@
-use jwst_core::{Base64Engine, URL_SAFE_ENGINE};
 use sea_orm::FromQueryResult;
-use sha2::{Digest, Sha256};
 
 use super::*;
 
@@ -208,33 +206,12 @@ impl BlobStorage<JwstStorageError> for BlobDBStorage {
         }
     }
 
-    async fn put_blob_stream(
-        &self,
-        workspace: Option<String>,
-        stream: impl Stream<Item = Bytes> + Send,
-    ) -> JwstStorageResult<String> {
-        let _lock = self.bucket.write().await;
-        let workspace = workspace.unwrap_or("__default__".into());
-
-        let (hash, blob) = get_hash(stream).await;
-
-        if self.insert(&workspace, &hash, &blob).await.is_ok() {
-            Ok(hash)
-        } else {
-            Err(JwstStorageError::WorkspaceNotFound(workspace))
-        }
-    }
-
-    async fn put_blob(&self, workspace: Option<String>, blob: Vec<u8>) -> JwstStorageResult<String> {
+    async fn put_blob(&self, workspace: Option<String>, id: String, blob: Vec<u8>) -> JwstStorageResult<String> {
         let _lock = self.bucket.write().await;
         let workspace = workspace.unwrap_or("__default__".into());
-        let mut hasher = Sha256::new();
-
-        hasher.update(&blob);
-        let hash = URL_SAFE_ENGINE.encode(hasher.finalize());
 
-        if self.insert(&workspace, &hash, &blob).await.is_ok() {
-            Ok(hash)
+        if self.insert(&workspace, &id, &blob).await.is_ok() {
+            Ok(id)
         } else {
             Err(JwstStorageError::WorkspaceNotFound(workspace))
         }

diff --git a/libs/jwst-storage/src/storage/blobs/mod.rs b/libs/jwst-storage/src/storage/blobs/mod.rs
@@ -14,12 +14,12 @@ mod utils;
 #[cfg(test)]
 pub use blob_storage::blobs_storage_test;
 pub use blob_storage::BlobDBStorage;
-use bytes::Bytes;
 use jwst_core::{BlobMetadata, BlobStorage};
 use jwst_storage_migration::Alias;
 use thiserror::Error;
 use tokio::task::JoinError;
-use utils::{get_hash, InternalBlobMetadata};
+pub use utils::stream_to_blob;
+use utils::InternalBlobMetadata;
 
 use super::{entities::prelude::*, *};
 
@@ -105,23 +105,14 @@ impl BlobStorage<JwstStorageError> for JwstBlobStorage {
         }
     }
 
-    async fn put_blob_stream(
+    async fn put_blob(
         &self,
         workspace: Option<String>,
-        stream: impl Stream<Item = Bytes> + Send,
+        id: String,
+        blob: Vec<u8>,
     ) -> JwstResult<String, JwstStorageError> {
         match self {
-            JwstBlobStorage::Raw(db) => db.put_blob_stream(workspace, stream).await,
-            #[cfg(feature = "image")]
-            JwstBlobStorage::Auto(db) => db.put_blob_stream(workspace, stream).await,
-            #[cfg(feature = "bucket")]
-            JwstBlobStorage::Bucket(db) => db.put_blob_stream(workspace, stream).await,
-        }
-    }
-
-    async fn put_blob(&self, workspace: Option<String>, blob: Vec<u8>) -> JwstResult<String, JwstStorageError> {
-        match self {
-            JwstBlobStorage::Raw(db) => db.put_blob(workspace, blob).await,
+            JwstBlobStorage::Raw(db) => db.put_blob(workspace, id, blob).await,
             #[cfg(feature = "image")]
             JwstBlobStorage::Auto(db) => db.put_blob(workspace, blob).await,
             #[cfg(feature = "bucket")]

diff --git a/libs/jwst-storage/src/storage/blobs/utils.rs b/libs/jwst-storage/src/storage/blobs/utils.rs
@@ -8,7 +8,7 @@ use jwst_core::{Base64Engine, BlobMetadata, URL_SAFE_ENGINE};
 use sea_orm::FromQueryResult;
 use sha2::{Digest, Sha256};
 
-pub async fn get_hash(stream: impl Stream<Item = Bytes> + Send) -> (String, Vec<u8>) {
+pub async fn stream_to_blob(stream: impl Stream<Item = Bytes> + Send) -> (String, Vec<u8>) {
     let mut hasher = Sha256::new();
 
     let buffer = stream