djangors_cache/
lib.rs

1#![deny(missing_docs)]
2//! Raw-byte caches, common backends, and explicitly opted-in response caching.
3
4use async_trait::async_trait;
5use chrono::Utc;
6use http_body_util::Full;
7use hyper::{Request, Response};
8use moka::future::Cache as MokaCache;
9use serde::{de::DeserializeOwned, Serialize};
10use std::{convert::Infallible, future::Future, sync::Arc, time::Duration};
11use thiserror::Error;
12use tower::{Layer, Service};
13
14/// Errors encountered during cache access or value serialization.
15#[derive(Debug, Error)]
16pub enum CacheError {
17    /// An error reported by the underlying cache backend.
18    #[error("cache backend error: {0}")]
19    Backend(String),
20    /// An error serializing or deserializing cached values.
21    #[error("cache serialization error: {0}")]
22    Serialization(String),
23}
24
25/// An object-safe cache of raw byte values.
26#[async_trait]
27pub trait Cache: Send + Sync {
28    /// Retrieves the raw byte value associated with `key` if present and unexpired.
29    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError>;
30    /// Stores `value` under `key` with an optional time-to-live (`ttl`).
31    async fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>)
32        -> Result<(), CacheError>;
33    /// Deletes any entry associated with `key`.
34    async fn delete(&self, key: &str) -> Result<(), CacheError>;
35}
36
37/// Extension trait providing convenient `get_or_set` helpers on any [`Cache`].
38#[async_trait]
39pub trait CacheExt: Cache {
40    /// Retrieves `key` if cached; otherwise executes `f`, caches the result under `key`, and returns it.
41    async fn get_or_set<F, Fut>(
42        &self,
43        key: &str,
44        ttl: Option<Duration>,
45        f: F,
46    ) -> Result<Vec<u8>, CacheError>
47    where
48        F: FnOnce() -> Fut + Send,
49        Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
50    {
51        if let Some(value) = self.get(key).await? {
52            return Ok(value);
53        }
54        let value = f().await?;
55        self.set(key, value.clone(), ttl).await?;
56        Ok(value)
57    }
58
59    /// Retrieves and deserializes `key` from JSON; otherwise executes `f`, serializes to JSON, caches, and returns `T`.
60    async fn get_or_set_json<T, F, Fut>(
61        &self,
62        key: &str,
63        ttl: Option<Duration>,
64        f: F,
65    ) -> Result<T, CacheError>
66    where
67        T: Serialize + DeserializeOwned + Send,
68        F: FnOnce() -> Fut + Send,
69        Fut: Future<Output = Result<T, CacheError>> + Send,
70    {
71        let bytes = self
72            .get_or_set(key, ttl, || async {
73                let value = f().await?;
74                serde_json::to_vec(&value).map_err(|e| CacheError::Serialization(e.to_string()))
75            })
76            .await?;
77        serde_json::from_slice(&bytes).map_err(|e| CacheError::Serialization(e.to_string()))
78    }
79}
80impl<T: Cache + ?Sized> CacheExt for T {}
81
82/// Helper function to retrieve or compute a template fragment or raw byte cache entry.
83pub async fn get_or_set_fragment<F, Fut>(
84    cache: &dyn Cache,
85    key: &str,
86    ttl: Option<Duration>,
87    f: F,
88) -> Result<Vec<u8>, CacheError>
89where
90    F: FnOnce() -> Fut + Send,
91    Fut: Future<Output = Result<Vec<u8>, CacheError>> + Send,
92{
93    cache.get_or_set(key, ttl, f).await
94}
95
96/// An in-memory cache backend powered by Moka.
97#[derive(Clone)]
98pub struct InMemoryCache {
99    inner: MokaCache<String, MemoryEntry>,
100}
101#[derive(Clone)]
102struct MemoryEntry {
103    value: Vec<u8>,
104    expires_at: Option<std::time::Instant>,
105}
106impl InMemoryCache {
107    /// Creates a new `InMemoryCache` instance with `max_capacity` elements limit.
108    pub fn new(max_capacity: u64) -> Self {
109        Self {
110            inner: MokaCache::builder().max_capacity(max_capacity).build(),
111        }
112    }
113}
114impl Default for InMemoryCache {
115    fn default() -> Self {
116        Self::new(10_000)
117    }
118}
119#[async_trait]
120impl Cache for InMemoryCache {
121    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
122        Ok(self.inner.get(key).await.and_then(|entry| {
123            if entry
124                .expires_at
125                .is_none_or(|e| e > std::time::Instant::now())
126            {
127                Some(entry.value)
128            } else {
129                None
130            }
131        }))
132    }
133    async fn set(
134        &self,
135        key: &str,
136        value: Vec<u8>,
137        ttl: Option<Duration>,
138    ) -> Result<(), CacheError> {
139        self.inner
140            .insert(
141                key.to_owned(),
142                MemoryEntry {
143                    value,
144                    expires_at: ttl.map(|t| std::time::Instant::now() + t),
145                },
146            )
147            .await;
148        Ok(())
149    }
150    async fn delete(&self, key: &str) -> Result<(), CacheError> {
151        self.inner.invalidate(key).await;
152        Ok(())
153    }
154}
155
156/// A database-backed cache backend storing entries in `djangors_cache_entries`.
157#[derive(Clone)]
158pub struct DatabaseCache {
159    db: djangors_db::Database,
160}
161impl DatabaseCache {
162    /// Creates a new `DatabaseCache` backend using `db`.
163    pub fn new(db: djangors_db::Database) -> Self {
164        Self { db }
165    }
166    async fn ensure_table(&self) -> Result<(), CacheError> {
167        let mut conn = self.db.conn();
168        let dialect = conn.dialect();
169        let bytea = dialect.bytea_type();
170        let sql = format!(
171            "CREATE TABLE IF NOT EXISTS djangors_cache_entries (key TEXT PRIMARY KEY, value {bytea} NOT NULL, expires_at TIMESTAMPTZ)"
172        );
173        conn.execute(&sql, &[])
174            .await
175            .map(|_| ())
176            .map_err(|e| CacheError::Backend(e.to_string()))
177    }
178}
179#[async_trait]
180impl Cache for DatabaseCache {
181    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
182        self.ensure_table().await?;
183        let mut conn = self.db.conn();
184        let p1 = conn.dialect().placeholder(1);
185        let sql = format!("SELECT value, expires_at FROM djangors_cache_entries WHERE key = {p1}");
186        let params = vec![djangors_db::BindValue::Text(key.to_string())];
187        let row_opt = conn
188            .fetch_optional(&sql, &params)
189            .await
190            .map_err(|e| CacheError::Backend(e.to_string()))?;
191        if let Some(row) = row_opt {
192            let value = row
193                .try_bytes(0)
194                .map_err(|e| CacheError::Backend(e.to_string()))?
195                .unwrap_or_default();
196            let expiry = row
197                .try_datetime(1)
198                .map_err(|e| CacheError::Backend(e.to_string()))?;
199            if expiry.is_none_or(|e| e > Utc::now()) {
200                return Ok(Some(value));
201            }
202            self.delete(key).await?;
203        }
204        Ok(None)
205    }
206    async fn set(
207        &self,
208        key: &str,
209        value: Vec<u8>,
210        ttl: Option<Duration>,
211    ) -> Result<(), CacheError> {
212        self.ensure_table().await?;
213        let expires = ttl.map(|t| Utc::now() + chrono::Duration::from_std(t).unwrap_or_default());
214        let mut conn = self.db.conn();
215        let dialect = conn.dialect();
216        let p1 = dialect.placeholder(1);
217        let p2 = dialect.placeholder(2);
218        let p3 = dialect.placeholder(3);
219        let sql = format!(
220            "INSERT INTO djangors_cache_entries (key, value, expires_at) VALUES ({p1}, {p2}, {p3}) \
221             ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expires_at = EXCLUDED.expires_at"
222        );
223        let expires_param = match expires {
224            Some(dt) => djangors_db::BindValue::DateTime(dt),
225            None => djangors_db::BindValue::Null(djangors_db::NullKind::DateTime),
226        };
227        let params = vec![
228            djangors_db::BindValue::Text(key.to_string()),
229            djangors_db::BindValue::Bytes(value),
230            expires_param,
231        ];
232        conn.execute(&sql, &params)
233            .await
234            .map(|_| ())
235            .map_err(|e| CacheError::Backend(e.to_string()))
236    }
237    async fn delete(&self, key: &str) -> Result<(), CacheError> {
238        self.ensure_table().await?;
239        let mut conn = self.db.conn();
240        let p1 = conn.dialect().placeholder(1);
241        let sql = format!("DELETE FROM djangors_cache_entries WHERE key = {p1}");
242        let params = vec![djangors_db::BindValue::Text(key.to_string())];
243        conn.execute(&sql, &params)
244            .await
245            .map(|_| ())
246            .map_err(|e| CacheError::Backend(e.to_string()))
247    }
248}
249
250#[cfg(feature = "redis")]
251/// A Redis-backed cache implementation.
252#[derive(Clone)]
253pub struct RedisCache {
254    client: redis::Client,
255}
256#[cfg(feature = "redis")]
257impl RedisCache {
258    /// Creates a new `RedisCache` connecting to the given Redis URL string.
259    pub fn new(url: &str) -> Result<Self, CacheError> {
260        redis::Client::open(url)
261            .map(|client| Self { client })
262            .map_err(|e| CacheError::Backend(e.to_string()))
263    }
264    async fn connection(&self) -> Result<redis::aio::MultiplexedConnection, CacheError> {
265        self.client
266            .get_multiplexed_async_connection()
267            .await
268            .map_err(|e| CacheError::Backend(e.to_string()))
269    }
270}
271#[cfg(feature = "redis")]
272#[async_trait]
273impl Cache for RedisCache {
274    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
275        let mut c = self.connection().await?;
276        redis::cmd("GET")
277            .arg(key)
278            .query_async(&mut c)
279            .await
280            .map_err(|e| CacheError::Backend(e.to_string()))
281    }
282    async fn set(
283        &self,
284        key: &str,
285        value: Vec<u8>,
286        ttl: Option<Duration>,
287    ) -> Result<(), CacheError> {
288        let mut c = self.connection().await?;
289        let mut cmd = redis::cmd("SET");
290        cmd.arg(key).arg(value);
291        if let Some(t) = ttl {
292            cmd.arg("PX").arg(t.as_millis() as u64);
293        }
294        cmd.query_async::<()>(&mut c)
295            .await
296            .map_err(|e| CacheError::Backend(e.to_string()))
297    }
298    async fn delete(&self, key: &str) -> Result<(), CacheError> {
299        let mut c = self.connection().await?;
300        redis::cmd("DEL")
301            .arg(key)
302            .query_async::<()>(&mut c)
303            .await
304            .map_err(|e| CacheError::Backend(e.to_string()))
305    }
306}
307
308/// Request/response extension marker indicating that a response is cacheable by [`CacheLayer`].
309#[derive(Clone, Copy, Debug)]
310pub struct CacheableResponse;
311
312/// Tower middleware layer for caching responses marked with [`CacheableResponse`].
313#[derive(Clone)]
314pub struct CacheLayer {
315    cache: Arc<dyn Cache>,
316    ttl: Option<Duration>,
317}
318impl CacheLayer {
319    /// Creates a new `CacheLayer` using `cache` and optional default `ttl`.
320    pub fn new(cache: Arc<dyn Cache>, ttl: Option<Duration>) -> Self {
321        Self { cache, ttl }
322    }
323}
324impl<S> Layer<S> for CacheLayer
325where
326    S: Service<
327            Request<Full<bytes::Bytes>>,
328            Response = Response<Full<bytes::Bytes>>,
329            Error = Infallible,
330        > + Clone
331        + Send
332        + 'static,
333    S::Future: Send + 'static,
334{
335    type Service = CacheService<S>;
336    fn layer(&self, inner: S) -> Self::Service {
337        CacheService {
338            inner,
339            cache: self.cache.clone(),
340            ttl: self.ttl,
341        }
342    }
343}
344
345/// Tower service middleware that intercepts GET requests and serves/caches responses.
346#[derive(Clone)]
347pub struct CacheService<S> {
348    inner: S,
349    cache: Arc<dyn Cache>,
350    ttl: Option<Duration>,
351}
352impl<S> Service<Request<Full<bytes::Bytes>>> for CacheService<S>
353where
354    S: Service<
355            Request<Full<bytes::Bytes>>,
356            Response = Response<Full<bytes::Bytes>>,
357            Error = Infallible,
358        > + Clone
359        + Send
360        + 'static,
361    S::Future: Send + 'static,
362{
363    type Response = Response<Full<bytes::Bytes>>;
364    type Error = Infallible;
365    type Future =
366        std::pin::Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
367    fn poll_ready(
368        &mut self,
369        cx: &mut std::task::Context<'_>,
370    ) -> std::task::Poll<Result<(), Self::Error>> {
371        self.inner.poll_ready(cx)
372    }
373    fn call(&mut self, req: Request<Full<bytes::Bytes>>) -> Self::Future {
374        let mut inner = self.inner.clone();
375        let cache = self.cache.clone();
376        let ttl = self.ttl;
377        Box::pin(async move {
378            if req.method() != hyper::Method::GET {
379                return inner.call(req).await;
380            }
381            let key = format!("{}?{}", req.uri().path(), req.uri().query().unwrap_or(""));
382            if let Some(bytes) = cache.get(&key).await.unwrap_or(None) {
383                if let Ok(cached) = serde_json::from_slice::<CachedResponse>(&bytes) {
384                    return Ok(cached.into_response());
385                }
386            }
387            let response = inner.call(req).await?;
388            if response.extensions().get::<CacheableResponse>().is_some()
389                && !response.headers().contains_key(hyper::header::SET_COOKIE)
390            {
391                let cached = CachedResponse::from_response(&response);
392                if let Ok(bytes) = serde_json::to_vec(&cached) {
393                    let _ = cache.set(&key, bytes, ttl).await;
394                }
395            }
396            Ok(response)
397        })
398    }
399}
400#[derive(serde::Serialize, serde::Deserialize)]
401struct CachedResponse {
402    status: u16,
403    headers: Vec<(String, Vec<u8>)>,
404    body: Vec<u8>,
405}
406impl CachedResponse {
407    fn from_response(r: &Response<Full<bytes::Bytes>>) -> Self {
408        Self {
409            status: r.status().as_u16(),
410            headers: r
411                .headers()
412                .iter()
413                .map(|(n, v)| (n.to_string(), v.as_bytes().to_vec()))
414                .collect(),
415            body: r.body().clone().into_inner().unwrap_or_default().to_vec(),
416        }
417    }
418    fn into_response(self) -> Response<Full<bytes::Bytes>> {
419        let mut r = Response::builder()
420            .status(self.status)
421            .body(Full::new(bytes::Bytes::from(self.body)))
422            .unwrap();
423        for (n, v) in self.headers {
424            if let (Ok(n), Ok(v)) = (
425                n.parse::<hyper::header::HeaderName>(),
426                hyper::header::HeaderValue::from_bytes(&v),
427            ) {
428                r.headers_mut().append(n, v);
429            }
430        }
431        r.extensions_mut().insert(CacheableResponse);
432        r
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use bytes::Bytes;
440    use std::{
441        convert::Infallible,
442        sync::{
443            atomic::{AtomicUsize, Ordering},
444            Arc,
445        },
446    };
447    use tower::{service_fn, ServiceExt};
448
449    async fn exercise(cache: &dyn Cache) {
450        assert_eq!(cache.get("missing").await.unwrap(), None);
451        cache
452            .set("key", b"value".to_vec(), Some(Duration::from_secs(5)))
453            .await
454            .unwrap();
455        assert_eq!(cache.get("key").await.unwrap(), Some(b"value".to_vec()));
456        assert_eq!(
457            cache
458                .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
459                .await
460                .unwrap(),
461            b"value"
462        );
463        cache.delete("key").await.unwrap();
464        assert_eq!(
465            cache
466                .get_or_set("key", None, || async { Ok(b"new".to_vec()) })
467                .await
468                .unwrap(),
469            b"new"
470        );
471        cache
472            .set("ttl", b"gone".to_vec(), Some(Duration::from_millis(50)))
473            .await
474            .unwrap();
475        tokio::time::sleep(Duration::from_millis(80)).await;
476        assert_eq!(cache.get("ttl").await.unwrap(), None);
477    }
478    #[tokio::test]
479    async fn in_memory_cache_operations_and_ttl() {
480        exercise(&InMemoryCache::default()).await;
481    }
482    #[tokio::test]
483    async fn cache_layer_requires_response_opt_in() {
484        let count = Arc::new(AtomicUsize::new(0));
485        let c = count.clone();
486        let handler = service_fn(move |_req: Request<Full<Bytes>>| {
487            let c = c.clone();
488            async move {
489                c.fetch_add(1, Ordering::SeqCst);
490                let mut r = Response::new(Full::new(Bytes::from_static(b"cached")));
491                r.extensions_mut().insert(CacheableResponse);
492                Ok::<_, Infallible>(r)
493            }
494        });
495        let mut svc = CacheLayer::new(Arc::new(InMemoryCache::default()), None).layer(handler);
496        let request = || {
497            Request::builder()
498                .method("GET")
499                .uri("/fragment?x=1")
500                .body(Full::new(Bytes::new()))
501                .unwrap()
502        };
503        assert_eq!(
504            svc.ready()
505                .await
506                .unwrap()
507                .call(request())
508                .await
509                .unwrap()
510                .body()
511                .clone()
512                .into_inner(),
513            Some(Bytes::from_static(b"cached"))
514        );
515        svc.ready().await.unwrap().call(request()).await.unwrap();
516        assert_eq!(count.load(Ordering::SeqCst), 1);
517    }
518    #[cfg(feature = "redis")]
519    #[tokio::test]
520    async fn redis_cache_operations_and_ttl() {
521        exercise(&RedisCache::new("redis://127.0.0.1:6379").unwrap()).await;
522    }
523    #[tokio::test]
524    async fn database_cache_operations_and_ttl() {
525        let url = std::env::var("DATABASE_URL")
526            .unwrap_or_else(|_| "postgres://postgres:postgres@localhost/djangors_test".into());
527        let db = djangors_db::Database::connect(&djangors_db::DatabaseConfig::new(url))
528            .await
529            .unwrap();
530        let cache = DatabaseCache::new(db);
531        sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
532            .execute(cache.db.pool())
533            .await
534            .unwrap();
535        exercise(&cache).await;
536        sqlx::query("DROP TABLE IF EXISTS djangors_cache_entries")
537            .execute(cache.db.pool())
538            .await
539            .unwrap();
540    }
541}
djangors_cache/lib.rs

djangors_cache/
lib.rs
