Skip to main content

djangors_cache/
lib.rs

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