From de056c8f0c53433ecade95141b5845f5b4b37d3f Mon Sep 17 00:00:00 2001 From: Dylan Knutson Date: Sun, 16 Jun 2024 18:26:58 -0700 Subject: [PATCH] add id to response --- src/compressor/compressor_trait.rs | 2 +- src/handlers/store_handler.rs | 125 ++++++++++++++++------------- src/sha256.rs | 7 ++ src/shard/fn_store.rs | 52 +++++++----- src/shard/mod.rs | 12 +-- src/shard/shard_struct.rs | 59 +++++--------- src/sql_types/entry_id.rs | 7 ++ src/sql_types/mod.rs | 2 + src/sql_types/utc_date_time.rs | 7 ++ 9 files changed, 152 insertions(+), 121 deletions(-) diff --git a/src/compressor/compressor_trait.rs b/src/compressor/compressor_trait.rs index 2a1d4ef..8e910f8 100644 --- a/src/compressor/compressor_trait.rs +++ b/src/compressor/compressor_trait.rs @@ -6,7 +6,7 @@ pub trait Compressor: Send { } #[cfg(test)] -pub(self) mod test { + mod test { use super::*; use crate::compressor::*; use rstest::*; diff --git a/src/handlers/store_handler.rs b/src/handlers/store_handler.rs index d921d77..d5e88a7 100644 --- a/src/handlers/store_handler.rs +++ b/src/handlers/store_handler.rs @@ -1,6 +1,6 @@ use crate::{ sha256::Sha256, - shard::{StoreArgs, StoreResult}, + shard::{StoreArgs, StoreResult, StoreResultKind}, shards::ShardsArc, }; use axum::http::StatusCode; @@ -19,60 +19,30 @@ pub struct StoreRequest { #[derive(Serialize, Debug, PartialEq)] #[serde(rename_all = "snake_case", tag = "status")] pub enum StoreResponse { - Created { - stored_size: usize, - data_size: usize, - created_at: String, - }, - Exists { - stored_size: usize, - data_size: usize, - created_at: String, - }, - Sha256Mismatch { - expected_sha256: String, - }, - InternalError { - message: String, - }, + Ok(StoreResult), + Sha256Mismatch { expected_sha256: String }, + InternalError { message: String }, +} + +impl From for StatusCode { + fn from(val: StoreResultKind) -> Self { + match val { + StoreResultKind::Created => StatusCode::CREATED, + StoreResultKind::Exists => StatusCode::OK, + } + } } impl StoreResponse { fn status_code(&self) -> StatusCode { match self { - StoreResponse::Created { .. } => StatusCode::CREATED, - StoreResponse::Exists { .. } => StatusCode::OK, + StoreResponse::Ok(StoreResult { kind, .. }) => (*kind).into(), 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, - created_at, - } => StoreResponse::Created { - stored_size, - data_size, - created_at: created_at.to_string(), - }, - StoreResult::Exists { - stored_size, - data_size, - created_at, - } => StoreResponse::Exists { - stored_size, - data_size, - created_at: created_at.to_string(), - }, - } - } -} - impl IntoResponse for StoreResponse { fn into_response(self) -> axum::response::Response { (self.status_code(), Json(self)).into_response() @@ -91,7 +61,7 @@ pub async fn store_handler( if sha256 != req_sha256 { let sha256_str = sha256.hex_string(); error!( - "client sent mismatched sha256: (client) {} != (computed) {}", + "sent mismatched sha256: (sent) {} != (computed) {}", req_sha256, sha256_str ); @@ -108,7 +78,7 @@ pub async fn store_handler( }) .await { - Ok(store_result) => store_result.into(), + Ok(store_result) => StoreResponse::Ok(store_result), Err(err) => StoreResponse::InternalError { message: err.to_string(), }, @@ -117,22 +87,21 @@ pub async fn store_handler( #[cfg(test)] pub mod test { - use crate::{compression_manager::test::make_compressor_with, shards::test::make_shards_with}; - use super::*; use crate::CompressionPolicy; + use crate::{compression_manager::test::make_compressor_with, shards::test::make_shards_with}; use axum::body::Bytes; use axum_typed_multipart::FieldData; use rstest::rstest; - async fn send_request>( - compression_policy: CompressionPolicy, + async fn send_request_with>( + shards: ShardsArc, sha256: Option, content_type: &str, data: D, ) -> StoreResponse { store_handler( - Extension(make_shards_with(make_compressor_with(compression_policy).into_arc()).await), + Extension(shards), TypedMultipart(StoreRequest { sha256: sha256.map(|s| s.hex_string()), content_type: content_type.to_string(), @@ -145,12 +114,31 @@ pub mod test { .await } + async fn send_request>( + compression_policy: CompressionPolicy, + sha256: Option, + content_type: &str, + data: D, + ) -> StoreResponse { + send_request_with( + make_shards(compression_policy).await, + sha256, + content_type, + data, + ) + .await + } + + async fn make_shards(compression_policy: CompressionPolicy) -> ShardsArc { + make_shards_with(make_compressor_with(compression_policy).into_arc()).await + } + #[tokio::test] async fn test_store_handler() { let result = send_request(CompressionPolicy::Auto, None, "text/plain", "hello, world!").await; assert_eq!(result.status_code(), StatusCode::CREATED, "{:?}", result); - assert!(matches!(result, StoreResponse::Created { .. })); + assert!(matches!(result, StoreResponse::Ok(..))); } #[tokio::test] @@ -184,7 +172,32 @@ pub mod test { ) .await; assert_eq!(result.status_code(), StatusCode::CREATED); - assert!(matches!(result, StoreResponse::Created { .. })); + assert!(matches!(result, StoreResponse::Ok(..))); + } + + #[tokio::test] + async fn test_id() { + let shards = make_shards(CompressionPolicy::Auto).await; + let data = "hello, world!"; + let sha256 = Sha256::from_bytes(data.as_bytes()); + let result = send_request_with(shards.clone(), Some(sha256), "text/plain", data).await; + assert_eq!(result.status_code(), StatusCode::CREATED); + let result = match result { + StoreResponse::Ok(result) => result, + _ => panic!("expected StoreResponse::Ok"), + }; + assert_eq!(result.sha256, sha256); + assert_eq!(result.kind, StoreResultKind::Created); + let id = result.id; + + let result = send_request_with(shards, Some(sha256), "text/plain", data).await; + assert_eq!(result.status_code(), StatusCode::OK); + let result = match result { + StoreResponse::Ok(result) => result, + _ => panic!("expected StoreResponse::Ok"), + }; + assert_eq!(result.kind, StoreResultKind::Exists); + assert_eq!(result.id, id); } fn make_assert_eq(value: T) -> impl Fn(T) { @@ -212,13 +225,15 @@ pub mod test { let result = send_request(compression_policy, None, content_type, vec![0; 1024]).await; assert_eq!(result.status_code(), StatusCode::CREATED); match result { - StoreResponse::Created { + StoreResponse::Ok(StoreResult { + kind, stored_size, data_size, .. - } => { + }) => { assert_stored_size(stored_size); assert_eq!(data_size, 1024); + assert_eq!(kind, StoreResultKind::Created); } _ => panic!("expected StoreResponse::Created"), }; diff --git a/src/sha256.rs b/src/sha256.rs index ca52277..2909db9 100644 --- a/src/sha256.rs +++ b/src/sha256.rs @@ -2,6 +2,7 @@ use rusqlite::{ types::{FromSql, ToSqlOutput}, ToSql, }; +use serde::Serialize; use sha2::Digest; use std::{ error::Error, @@ -50,6 +51,12 @@ impl Sha256 { } } +impl Serialize for Sha256 { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_bytes(&self.0) + } +} + impl ToSql for Sha256 { fn to_sql(&self) -> rusqlite::Result { Ok(ToSqlOutput::Borrowed(rusqlite::types::ValueRef::Blob( diff --git a/src/shard/fn_store.rs b/src/shard/fn_store.rs index 80d685b..ff789a6 100644 --- a/src/shard/fn_store.rs +++ b/src/shard/fn_store.rs @@ -1,23 +1,27 @@ +use serde::Serialize; + use super::*; use crate::{ compressible_data::CompressibleData, concat_lines, into_tokio_rusqlite_err, - sql_types::{CompressionId, UtcDateTime}, + sql_types::{CompressionId, EntryId, UtcDateTime}, AsyncBoxError, }; -#[derive(PartialEq, Debug)] -pub enum StoreResult { - Created { - stored_size: usize, - data_size: usize, - created_at: UtcDateTime, - }, - Exists { - stored_size: usize, - data_size: usize, - created_at: UtcDateTime, - }, +#[derive(PartialEq, Debug, Serialize, Copy, Clone)] +pub enum StoreResultKind { + Created, + Exists, +} + +#[derive(PartialEq, Debug, Serialize)] +pub struct StoreResult { + pub kind: StoreResultKind, + pub id: EntryId, + pub sha256: Sha256, + pub stored_size: usize, + pub data_size: usize, + pub created_at: UtcDateTime, } #[derive(Default)] @@ -75,16 +79,19 @@ fn find_with_sha256( ) -> Result, rusqlite::Error> { conn.query_row( concat_lines!( - "SELECT uncompressed_size, compressed_size, created_at", + "SELECT rowid, 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)?, - created_at: row.get(2)?, + Ok(StoreResult { + kind: StoreResultKind::Exists, + id: row.get(0)?, + sha256: *sha256, + stored_size: row.get(1)?, + data_size: row.get(2)?, + created_at: row.get(3)?, }) }, ) @@ -102,11 +109,12 @@ fn insert( let created_at = UtcDateTime::now(); let compressed_size = data.len(); - conn.execute( + let id = conn.query_row( concat_lines!( "INSERT INTO entries", " (sha256, content_type, dict_id, uncompressed_size, compressed_size, data, created_at)", "VALUES (?, ?, ?, ?, ?, ?, ?)", + "RETURNING rowid" ), params![ sha256, @@ -117,9 +125,13 @@ fn insert( data.as_ref(), created_at, ], + |row| row.get(0), )?; - Ok(StoreResult::Created { + Ok(StoreResult { + kind: StoreResultKind::Created, + id, + sha256: *sha256, stored_size: compressed_size, data_size: uncompressed_size, created_at, diff --git a/src/shard/mod.rs b/src/shard/mod.rs index 5b0d634..dd132b9 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -5,15 +5,15 @@ pub mod shard_error; mod shard_struct; pub use fn_get::{GetArgs, GetResult}; -pub use fn_store::{StoreArgs, StoreResult}; +pub use fn_store::{StoreArgs, StoreResult, StoreResultKind}; pub use shard_struct::Shard; - -#[cfg(test)] -pub mod test { - pub use super::shard_struct::test::*; -} use crate::{sha256::Sha256, shard::shard_error::ShardError}; use axum::body::Bytes; use rusqlite::{params, types::FromSql, OptionalExtension}; use tokio_rusqlite::Connection; use tracing::debug; + +#[cfg(test)] +pub mod test { + pub use super::shard_struct::test::*; +} diff --git a/src/shard/shard_struct.rs b/src/shard/shard_struct.rs index 99ae756..289f90a 100644 --- a/src/shard/shard_struct.rs +++ b/src/shard/shard_struct.rs @@ -64,10 +64,11 @@ async fn get_num_entries(conn: &Connection) -> Result { - // assert_eq!(stored_size, data.len()); - assert_eq!(data_size, data.len()); - assert!(created_at > chrono::Utc::now() - chrono::Duration::seconds(1)); - } - _ => panic!("expected StoreResult::Created"), - } + + assert_eq!(store_result.kind, StoreResultKind::Created); + assert_eq!(store_result.sha256, sha256); + assert_eq!(store_result.data_size, data.len()); + assert!(store_result.created_at > chrono::Utc::now() - chrono::Duration::seconds(1)); assert_eq!(shard.num_entries().await.unwrap(), 1); let get_result = shard.get(GetArgs { sha256 }).await.unwrap().unwrap(); @@ -148,42 +141,31 @@ pub mod test { sha256, content_type: "text/plain".to_string(), data: data.into(), - ..Default::default() }) .await .unwrap(); - assert!(matches!(store_result, StoreResult::Created { .. })); - let created_at = match store_result { - StoreResult::Created { - stored_size, - data_size, - created_at, - } => { - assert_eq!(stored_size, data.len()); - assert_eq!(data_size, data.len()); - created_at - } - _ => panic!("expected StoreResult::Created"), - }; + assert_eq!(store_result.kind, StoreResultKind::Created); + assert_eq!(store_result.sha256, sha256); + assert_eq!(store_result.stored_size, data.len()); + assert_eq!(store_result.data_size, data.len()); + let created_at = store_result.created_at; let store_result = shard .store(StoreArgs { sha256, content_type: "text/plain".to_string(), data: data.into(), - ..Default::default() }) .await .unwrap(); - assert_eq!( - store_result, - StoreResult::Exists { - data_size: data.len(), - stored_size: data.len(), - created_at - } - ); + + assert_eq!(store_result.kind, StoreResultKind::Exists); + assert_eq!(store_result.sha256, sha256); + assert_eq!(store_result.data_size, data.len(),); + assert_eq!(store_result.stored_size, data.len(),); + assert_eq!(store_result.created_at, created_at); + assert_eq!(shard.num_entries().await.unwrap(), 1); } @@ -213,11 +195,10 @@ pub mod test { sha256, content_type: content_type.clone(), data: data.clone().into(), - ..Default::default() }) .await .unwrap(); - assert!(matches!(store_result, StoreResult::Created { .. })); + assert_eq!(store_result.kind, StoreResultKind::Created); let get_result = shard.get(GetArgs { sha256 }).await.unwrap().unwrap(); assert_eq!(get_result.content_type, content_type); diff --git a/src/sql_types/entry_id.rs b/src/sql_types/entry_id.rs index 8f41d78..5ee0a0d 100644 --- a/src/sql_types/entry_id.rs +++ b/src/sql_types/entry_id.rs @@ -2,6 +2,7 @@ use rusqlite::{ types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef}, ToSql, }; +use serde::Serialize; #[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)] pub struct EntryId(pub i64); @@ -11,6 +12,12 @@ impl From for EntryId { } } +impl Serialize for EntryId { + fn serialize(&self, serializer: S) -> Result { + self.0.serialize(serializer) + } +} + impl FromSql for EntryId { fn column_result(value: ValueRef<'_>) -> FromSqlResult { Ok(value.as_i64()?.into()) diff --git a/src/sql_types/mod.rs b/src/sql_types/mod.rs index 346bf34..fa30745 100644 --- a/src/sql_types/mod.rs +++ b/src/sql_types/mod.rs @@ -1,5 +1,7 @@ mod compression_id; +mod entry_id; mod utc_date_time; pub use compression_id::CompressionId; +pub use entry_id::EntryId; pub use utc_date_time::UtcDateTime; diff --git a/src/sql_types/utc_date_time.rs b/src/sql_types/utc_date_time.rs index f1f55b1..b0e9d16 100644 --- a/src/sql_types/utc_date_time.rs +++ b/src/sql_types/utc_date_time.rs @@ -3,6 +3,7 @@ use rusqlite::{ types::{FromSql, FromSqlError, ToSqlOutput, ValueRef}, Result, ToSql, }; +use serde::Serialize; #[derive(PartialEq, Debug, PartialOrd)] pub struct UtcDateTime(DateTime); @@ -24,6 +25,12 @@ impl ToString for UtcDateTime { } } +impl Serialize for UtcDateTime { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(&self.to_string()) + } +} + impl PartialEq> for UtcDateTime { fn eq(&self, other: &DateTime) -> bool { self.0 == *other