From e63403b4eea9c6e092776e1db0728dc99d6b37c4 Mon Sep 17 00:00:00 2001 From: Dylan Knutson Date: Thu, 25 Apr 2024 09:58:54 -0700 Subject: [PATCH] refactor store request/response, use IntoResponse --- src/handlers/store_handler.rs | 104 ++++++++++++++++++++++---- src/main.rs | 1 - src/request_response/mod.rs | 1 - src/request_response/store_request.rs | 41 ---------- src/sha256.rs | 8 +- src/shard/fn_get.rs | 36 +++++++++ src/shard/fn_store.rs | 78 +++++++++++++++++++ src/shard/mod.rs | 89 +++------------------- 8 files changed, 220 insertions(+), 138 deletions(-) delete mode 100644 src/request_response/mod.rs delete mode 100644 src/request_response/store_request.rs create mode 100644 src/shard/fn_get.rs create mode 100644 src/shard/fn_store.rs diff --git a/src/handlers/store_handler.rs b/src/handlers/store_handler.rs index dc83d3f..89be049 100644 --- a/src/handlers/store_handler.rs +++ b/src/handlers/store_handler.rs @@ -1,18 +1,80 @@ -use crate::{ - request_response::store_request::{StoreRequest, StoreResult}, - sha256::Sha256, - shards::Shards, -}; -use axum::Extension; -use axum_typed_multipart::TypedMultipart; +use crate::{sha256::Sha256, shard::StoreResult, shards::Shards}; +use axum::{body::Bytes, response::IntoResponse, Extension, Json}; +use axum_typed_multipart::{FieldData, TryFromMultipart, TypedMultipart}; +use axum::http::StatusCode; +use serde::Serialize; use tracing::error; +#[derive(TryFromMultipart)] +pub struct StoreRequest { + pub sha256: Option, + pub content_type: String, + pub data: FieldData, +} + +#[derive(Serialize, Debug, PartialEq)] +#[serde(rename_all = "snake_case", tag = "status")] +pub enum StoreResponse { + Created { + stored_size: usize, + data_size: usize, + }, + Exists { + stored_size: usize, + data_size: usize, + }, + Sha256Mismatch { + expected_sha256: String, + }, + InternalError { + message: String, + }, +} + +impl StoreResponse { + fn status_code(&self) -> StatusCode { + match self { + StoreResponse::Created { .. } => StatusCode::CREATED, + StoreResponse::Exists { .. } => StatusCode::OK, + StoreResponse::Sha256Mismatch { .. } => StatusCode::BAD_REQUEST, + StoreResponse::InternalError { .. } => StatusCode::INTERNAL_SERVER_ERROR, + } + } +} + +impl From for StoreResponse { + fn from(result: StoreResult) -> Self { + match result { + StoreResult::Created { + stored_size, + data_size, + } => StoreResponse::Created { + stored_size, + data_size, + }, + StoreResult::Exists { + stored_size, + data_size, + } => StoreResponse::Exists { + stored_size, + data_size, + }, + } + } +} + +impl IntoResponse for StoreResponse { + fn into_response(self) -> axum::response::Response { + (self.status_code(), Json(self)).into_response() + } +} + #[axum::debug_handler] pub async fn store_handler( Extension(shards): Extension, TypedMultipart(request): TypedMultipart, -) -> StoreResult { +) -> StoreResponse { let sha256 = Sha256::from_bytes(&request.data.contents); let shard = shards.shard_for(&sha256); @@ -23,12 +85,21 @@ pub async fn store_handler( "client sent mismatched sha256: (client) {} != (computed) {}", req_sha256, sha256_str ); - return StoreResult::Sha256Mismatch { sha256: sha256_str }; + + return StoreResponse::Sha256Mismatch { + expected_sha256: sha256_str, + }; } } - shard + match shard .store(sha256, request.content_type, request.data.contents) .await + { + Ok(store_result) => store_result.into(), + Err(err) => StoreResponse::InternalError { + message: err.to_string(), + }, + } } #[cfg(test)] @@ -42,7 +113,7 @@ mod test { Shards::new(vec![make_shard().await]).unwrap() } - async fn send_request(sha256: Option, data: Bytes) -> StoreResult { + async fn send_request(sha256: Option, data: Bytes) -> StoreResponse { store_handler( Extension(make_shards().await), TypedMultipart(StoreRequest { @@ -60,7 +131,8 @@ mod test { #[tokio::test] async fn test_store_handler() { let result = send_request(None, "hello, world!".as_bytes().into()).await; - assert!(matches!(result, StoreResult::Created { .. })); + assert_eq!(result.status_code(), StatusCode::CREATED); + assert!(matches!(result, StoreResponse::Created { .. })); } #[tokio::test] @@ -68,10 +140,11 @@ mod test { let not_hello_world = Sha256::from_bytes("not hello, world!".as_bytes()); let hello_world = Sha256::from_bytes("hello, world!".as_bytes()); let result = send_request(Some(not_hello_world), "hello, world!".as_bytes().into()).await; + assert_eq!(result.status_code(), StatusCode::BAD_REQUEST); assert_eq!( result, - StoreResult::Sha256Mismatch { - sha256: hello_world.hex_string() + StoreResponse::Sha256Mismatch { + expected_sha256: hello_world.hex_string() } ); } @@ -80,6 +153,7 @@ mod test { async fn test_store_handler_matching_sha256() { let hello_world = Sha256::from_bytes("hello, world!".as_bytes()); let result = send_request(Some(hello_world), "hello, world!".as_bytes().into()).await; - assert!(matches!(result, StoreResult::Created { .. })); + assert_eq!(result.status_code(), StatusCode::CREATED); + assert!(matches!(result, StoreResponse::Created { .. })); } } diff --git a/src/main.rs b/src/main.rs index 1507a8c..cb00039 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,4 @@ mod handlers; -mod request_response; mod sha256; mod shard; mod shards; diff --git a/src/request_response/mod.rs b/src/request_response/mod.rs deleted file mode 100644 index 4065c30..0000000 --- a/src/request_response/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod store_request; diff --git a/src/request_response/store_request.rs b/src/request_response/store_request.rs deleted file mode 100644 index 4208048..0000000 --- a/src/request_response/store_request.rs +++ /dev/null @@ -1,41 +0,0 @@ -use axum::{body::Bytes, http::StatusCode, response::IntoResponse, Json}; -use axum_typed_multipart::{FieldData, TryFromMultipart}; -use serde::Serialize; - -#[derive(TryFromMultipart)] -pub struct StoreRequest { - pub sha256: Option, - pub content_type: String, - pub data: FieldData, -} - -#[derive(Serialize, PartialEq, Debug)] -#[serde(tag = "status", rename_all = "snake_case")] -pub enum StoreResult { - Created { - stored_size: usize, - data_size: usize, - }, - Exists { - stored_size: usize, - data_size: usize, - }, - Sha256Mismatch { - sha256: String, - }, - InternalError { - message: String, - }, -} - -impl IntoResponse for StoreResult { - fn into_response(self) -> axum::response::Response { - let status_code = match &self { - StoreResult::Created { .. } => StatusCode::CREATED, - StoreResult::Exists { .. } => StatusCode::OK, - StoreResult::Sha256Mismatch { .. } => StatusCode::BAD_REQUEST, - StoreResult::InternalError { .. } => StatusCode::INTERNAL_SERVER_ERROR, - }; - (status_code, Json(self)).into_response() - } -} diff --git a/src/sha256.rs b/src/sha256.rs index e5d76d8..0ca39ec 100644 --- a/src/sha256.rs +++ b/src/sha256.rs @@ -17,7 +17,7 @@ impl Display for Sha256Error { } impl Error for Sha256Error {} -#[derive(Clone, Copy)] +#[derive(Clone, Copy, PartialEq, Eq)] pub struct Sha256([u8; 32]); impl Sha256 { pub fn from_hex_string(hex: &str) -> Result> { @@ -65,3 +65,9 @@ impl PartialEq for Sha256 { } } } + +impl core::fmt::Debug for Sha256 { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Sha256({})", self.hex_string()) + } +} diff --git a/src/shard/fn_get.rs b/src/shard/fn_get.rs new file mode 100644 index 0000000..c69a9eb --- /dev/null +++ b/src/shard/fn_get.rs @@ -0,0 +1,36 @@ +use super::*; + +impl Shard { + pub async fn get(&self, sha256: Sha256) -> Result, Box> { + self.conn + .call(move |conn| get_impl(conn, sha256).map_err(|e| e.into())) + .await + .map_err(|e| { + error!("get failed: {}", e); + e.into() + }) + } +} + +fn get_impl( + conn: &mut rusqlite::Connection, + sha256: Sha256, +) -> Result, rusqlite::Error> { + conn.query_row( + "SELECT content_type, compressed_size, created_at, data FROM entries WHERE sha256 = ?", + params![sha256.hex_string()], + |row| { + let content_type = row.get(0)?; + let stored_size = row.get(1)?; + let created_at = parse_created_at_str(row.get(2)?)?; + let data = row.get(3)?; + Ok(GetResult { + content_type, + stored_size, + created_at, + data, + }) + }, + ) + .optional() +} diff --git a/src/shard/fn_store.rs b/src/shard/fn_store.rs new file mode 100644 index 0000000..59fe488 --- /dev/null +++ b/src/shard/fn_store.rs @@ -0,0 +1,78 @@ +use super::*; + +#[derive(PartialEq, Debug)] +pub enum StoreResult { + Created { + stored_size: usize, + data_size: usize, + }, + Exists { + stored_size: usize, + data_size: usize, + }, +} + +impl Shard { + pub async fn store( + &self, + sha256: Sha256, + content_type: String, + data: Bytes, + ) -> Result> { + self.conn + .call(move |conn| store(conn, sha256, content_type, data).map_err(|e| e.into())) + .await + .map_err(|e| { + error!("store failed: {}", e); + e.into() + }) + } +} + +fn store( + conn: &mut rusqlite::Connection, + sha256: Sha256, + content_type: String, + data: Bytes, +) -> Result { + let sha256 = sha256.hex_string(); + + // check for existing entry + let maybe_existing: Option = conn + .query_row( + "SELECT uncompressed_size, compressed_size, created_at FROM entries WHERE sha256 = ?", + params![sha256], + |row| { + Ok(StoreResult::Exists { + stored_size: row.get(0)?, + data_size: row.get(1)?, + }) + }, + ) + .optional()?; + + if let Some(existing) = maybe_existing { + return Ok(existing); + } + + let created_at = chrono::Utc::now().to_rfc3339(); + let data_size = data.len(); + + conn.execute( + "INSERT INTO entries (sha256, content_type, compression, uncompressed_size, compressed_size, data, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + params![ + sha256, + content_type, + 0, + data_size, + data_size, + &data[..], + created_at, + ], + )?; + + Ok(StoreResult::Created { + stored_size: data_size, + data_size, + }) +} diff --git a/src/shard/mod.rs b/src/shard/mod.rs index dcd4949..be3c7c4 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -1,10 +1,11 @@ +mod fn_get; mod fn_migrate; +mod fn_store; pub mod shard_error; -use crate::{ - request_response::store_request::StoreResult, sha256::Sha256, shard::shard_error::ShardError, -}; +use crate::{sha256::Sha256, shard::shard_error::ShardError}; use axum::body::Bytes; +pub use fn_store::StoreResult; use rusqlite::{params, types::FromSql, OptionalExtension}; use std::error::Error; @@ -61,79 +62,6 @@ impl Shard { .map_err(|e| e.into()) } - pub async fn store(&self, sha256: Sha256, content_type: String, data: Bytes) -> StoreResult { - let sha256 = sha256.hex_string(); - self.conn.call(move |conn| { - // check for existing entry - let maybe_existing: Option = conn - .query_row( - "SELECT uncompressed_size, compressed_size, created_at FROM entries WHERE sha256 = ?", - params![sha256], - |row| Ok(StoreResult::Exists { - stored_size: row.get(0)?, - data_size: row.get(1)?, - }) - ) - .optional()?; - - if let Some(existing) = maybe_existing { - return Ok(existing); - } - - let created_at = chrono::Utc::now().to_rfc3339(); - let data_size = data.len(); - - conn.execute( - "INSERT INTO entries (sha256, content_type, compression, uncompressed_size, compressed_size, data, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", - params![ - sha256, - content_type, - 0, - data_size, - data_size, - &data[..], - created_at, - ], - )?; - - Ok(StoreResult::Created { stored_size: data_size, data_size }) - }) - .await.unwrap_or_else(|e| { - error!("store failed: {}", e); - StoreResult::InternalError { message: e.to_string() } - }) - } - - pub async fn get(&self, sha256: Sha256) -> Result, Box> { - self.conn - .call(move |conn| { - let get_result = conn - .query_row( - "SELECT content_type, compressed_size, created_at, data FROM entries WHERE sha256 = ?", - params![sha256.hex_string()], - |row| { - let content_type = row.get(0)?; - let stored_size = row.get(1)?; - let created_at = parse_created_at_str(row.get(2)?)?; - let data = row.get(3)?; - Ok(GetResult { - content_type, - stored_size, - created_at, - data, - }) - }, - ) - .optional()?; - Ok(get_result) - }) - .await - .map_err(|e| { - error!("get failed: {}", e); - e.into() - }) - } - pub async fn num_entries(&self) -> Result> { get_num_entries(&self.conn).await.map_err(|e| e.into()) } @@ -191,7 +119,8 @@ pub mod test { let sha256 = Sha256::from_bytes(data); let store_result = shard .store(sha256, "text/plain".to_string(), data.into()) - .await; + .await + .unwrap(); assert_eq!( store_result, super::StoreResult::Created { @@ -215,7 +144,8 @@ pub mod test { let store_result = shard .store(sha256, "text/plain".to_string(), data.into()) - .await; + .await + .unwrap(); assert_eq!( store_result, super::StoreResult::Created { @@ -226,7 +156,8 @@ pub mod test { let store_result = shard .store(sha256, "text/plain".to_string(), data.into()) - .await; + .await + .unwrap(); assert_eq!( store_result, super::StoreResult::Exists {