Skip to content

Commit

Permalink
feat: Into Response for CacheResponse, create_cache_key accept u8 array
Browse files Browse the repository at this point in the history
  • Loading branch information
kane50613 committed Dec 6, 2024
1 parent e2856f1 commit e2936f1
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 39 deletions.
14 changes: 14 additions & 0 deletions src/db.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use axum::response::Response;
use reqwest::header::CONTENT_TYPE;
use sqlx::Statement;
use std::sync::LazyLock;

Expand Down Expand Up @@ -25,6 +27,18 @@ pub struct CacheResponse {
pub content_type: Option<String>,
}

impl From<CacheResponse> for Response {
fn from(val: CacheResponse) -> Self {
let mut builder = Response::builder().status(val.status);

if let Some(content_type) = val.content_type {
builder = builder.header(CONTENT_TYPE, content_type);
}

builder.body(val.body.unwrap_or_default().into()).unwrap()
}
}

pub struct Database<'a> {
connection: SqlitePool,
queries: OnceCell<Queries<'a>>,
Expand Down
35 changes: 10 additions & 25 deletions src/handlers/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use reqwest::{

use crate::{
db::{CacheResponse, DB},
hash::{create_cache_key, CacheKeyPayload},
hash::create_cache_key,
};

pub static HTTP: LazyLock<Client> = LazyLock::new(Client::default);
Expand All @@ -30,17 +30,14 @@ pub async fn api_handler(request: Request) -> Response {
.unwrap();
};

let cache_key = create_cache_key(CacheKeyPayload {
method: head.method.to_string(),
url: path.to_string(),
authorization: head
.headers
.get(AUTHORIZATION)
.map(|x| x.to_str().unwrap().to_string()),
});
let cache_key = create_cache_key(
head.method.as_str().as_bytes(),
path.as_str().as_bytes(),
head.headers.get(AUTHORIZATION).map(|x| x.as_bytes()),
);

if let Some(cache_response) = DB.get(cache_key).await {
return response_from_cache(cache_response);
return cache_response.into();
}

let url = Uri::builder()
Expand All @@ -59,7 +56,7 @@ pub async fn api_handler(request: Request) -> Response {
.await
.unwrap();

let status = response.status();
let status = response.status().as_u16();

let content_type = response.headers_mut().remove(CONTENT_TYPE);
let body = response.bytes().await.unwrap();
Expand All @@ -68,22 +65,10 @@ pub async fn api_handler(request: Request) -> Response {
key: cache_key,
content_type: content_type.map(|x| x.to_str().unwrap().to_string()),
body: Some(body.to_vec()),
status: status.as_u16(),
status,
};

DB.set(&cache_payload).await;

response_from_cache(cache_payload)
}

fn response_from_cache(cache_response: CacheResponse) -> Response {
let mut builder = Response::builder().status(cache_response.status);

if let Some(content_type) = &cache_response.content_type {
builder = builder.header(CONTENT_TYPE, content_type);
}

builder
.body(cache_response.body.unwrap_or_default().into())
.unwrap()
cache_payload.into()
}
6 changes: 5 additions & 1 deletion src/handlers/invalidate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@ use crate::{
};

pub async fn invalidate_handler(body: Json<CacheKeyPayload>) {
let key = create_cache_key(body.0);
let key = create_cache_key(
body.method.as_bytes(),
body.url.as_bytes(),
body.authorization.as_ref().map(|header| header.as_bytes()),
);

DB.delete(key).await;
}
21 changes: 8 additions & 13 deletions src/hash.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,16 @@ pub struct CacheKeyPayload {
pub authorization: Option<String>,
}

pub fn create_cache_key(payload: CacheKeyPayload) -> i64 {
let method_bytes = payload.method.as_bytes();
let url_bytes = payload.url.as_bytes();
let authorization_bytes = payload
.authorization
.as_deref()
.unwrap_or_default()
.as_bytes();

pub fn create_cache_key(method: &[u8], url: &[u8], authorization: Option<&[u8]>) -> i64 {
let mut buffer =
Vec::with_capacity(method_bytes.len() + url_bytes.len() + authorization_bytes.len());
Vec::with_capacity(method.len() + url.len() + authorization.map_or(0, |x| x.len()));

buffer.extend_from_slice(method);
buffer.extend_from_slice(url);

buffer.extend_from_slice(method_bytes);
buffer.extend_from_slice(url_bytes);
buffer.extend_from_slice(authorization_bytes);
if let Some(authorization) = authorization {
buffer.extend_from_slice(authorization);
}

xxh3::xxh3_64(&buffer) as i64
}

0 comments on commit e2936f1

Please sign in to comment.