url_cleaner_engine/glue/
caching.rs

1//! Caching to allow for only expanding redirects the first time you encounter them.
2
3use std::sync::Mutex;
4use std::str::FromStr;
5use std::cell::OnceCell;
6
7use thiserror::Error;
8use serde::{Serialize, Deserialize};
9use diesel::prelude::*;
10#[expect(unused_imports, reason = "Used in docs.")]
11use diesel::query_builder::SqlQuery;
12
13use crate::util::*;
14
15diesel::table! {
16    /// The table containing cache entries.
17    cache (id) {
18        /// The entry's unique ID.
19        id -> Integer,
20        /// The "category" of the entry.
21        category -> Text,
22        /// The "key" of the entry.
23        key -> Text,
24        /// The value of the entry.
25        value -> Nullable<Text>,
26    }
27}
28
29/// The Sqlite command to initialize the cache database.
30pub const DB_INIT_COMMAND: &str = r#"CREATE TABLE cache (
31    id INTEGER NOT NULL PRIMARY KEY,
32    category TEXT NOT NULL,
33    "key" TEXT NOT NULL,
34    value TEXT,
35    UNIQUE(category, "key") ON CONFLICT REPLACE
36)"#;
37
38/// An entry in the cache database.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Queryable, Selectable)]
40#[diesel(table_name = cache)]
41pub struct CacheEntry {
42    /// The ID of the entry.
43    pub id: i32,
44    /// The category of the entry.
45    pub category: String,
46    /// The key of the entry.
47    pub key: String,
48    /// The value of the entry.
49    pub value: Option<String>
50}
51
52/// A new entry for the cache database.
53#[derive(Debug, PartialEq, Eq, Serialize, Insertable)]
54#[diesel(table_name = cache)]
55pub struct NewCacheEntry<'a> {
56    /// The category of the new entry.
57    pub category: &'a str,
58    /// The key of the new entry.
59    pub key: &'a str,
60    /// The value of the new entry.
61    pub value: Option<&'a str>
62}
63
64/// A [`Mutex`]ed [`InnerCache`].
65/// # Examples
66/// ```
67/// use url_cleaner_engine::glue::*;
68///
69/// // Note the immutability.
70/// let cache = Cache::new(CachePath::Memory);
71///
72/// assert_eq!(cache.read("category", "key").unwrap(), None);
73/// cache.write("category", "key", None).unwrap();
74/// assert_eq!(cache.read("category", "key").unwrap(), Some(None));
75/// cache.write("category", "key", Some("value")).unwrap();
76/// assert_eq!(cache.read("category", "key").unwrap(), Some(Some("value".into())));
77/// ```
78#[derive(Debug, Default)]
79pub struct Cache(pub Mutex<InnerCache>);
80
81impl Cache {
82    /// Create a new unconnected [`Self`].
83    #[allow(dead_code, reason = "Public API.")]
84    pub fn new(path: CachePath) -> Self {
85        path.into()
86    }
87}
88
89impl From<InnerCache> for Cache {
90    fn from(value: InnerCache) -> Self {
91        Self(Mutex::new(value))
92    }
93}
94
95impl From<CachePath> for Cache {
96    fn from(value: CachePath) -> Self {
97        Cache::from(InnerCache::from(value))
98    }
99}
100
101/// A lazily connected connection to the cache database.
102/// # Examples
103/// ```
104/// use url_cleaner_engine::glue::*;
105///
106/// // Note the mutability.
107/// let mut cache = InnerCache::new(CachePath::Memory);
108///
109///
110/// assert_eq!(cache.read("category", "key").unwrap(), None);
111/// cache.write("category", "key", None).unwrap();
112/// assert_eq!(cache.read("category", "key").unwrap(), Some(None));
113/// cache.write("category", "key", Some("value")).unwrap();
114/// assert_eq!(cache.read("category", "key").unwrap(), Some(Some("value".into())));
115/// ```
116#[derive(Default)]
117pub struct InnerCache {
118    /// The path of the database.
119    path: CachePath,
120    /// The connection to the database.
121    connection: OnceCell<SqliteConnection>
122}
123
124impl InnerCache {
125    /// Create a new unconnected [`Self`].
126    pub fn new(path: CachePath) -> Self {
127        path.into()
128    }
129}
130
131impl PartialEq for InnerCache {
132    fn eq(&self, other: &Self) -> bool {
133        self.path == other.path
134    }
135}
136impl Eq for InnerCache {}
137
138impl From<CachePath> for InnerCache {
139    fn from(value: CachePath) -> Self {
140        Self {
141            path: value,
142            connection: Default::default()
143        }
144    }
145}
146
147/// The path of a cache database.
148#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, Suitability)]
149#[serde(remote = "Self")]
150pub enum CachePath {
151    /// Stores the database in memory, wiping it on program exit.
152    ///
153    /// Has the string representation of `:memory:`.
154    #[default]
155    Memory,
156    /// A filesystem/network/whatever path.
157    Path(String)
158}
159
160crate::util::string_or_struct_magic!(CachePath);
161
162impl CachePath {
163    /// The cache's path as a [`str`].
164    ///
165    /// If `self` is [`Self::Memory`], returns `:memory:`.
166    /// # Examples
167    /// ```
168    /// use url_cleaner_engine::glue::*;
169    ///
170    /// assert_eq!(CachePath::Memory                          .as_str(), ":memory:");
171    /// assert_eq!(CachePath::Path(       "abc.sqlite".into()).as_str(), "abc.sqlite");
172    /// assert_eq!(CachePath::Path("file://abc.sqlite".into()).as_str(), "file://abc.sqlite");
173    /// ```
174    pub fn as_str(&self) -> &str {
175        match self {
176            Self::Memory => ":memory:",
177            Self::Path(x) => x
178        }
179    }
180}
181
182impl AsRef<str> for CachePath {
183    fn as_ref(&self) -> &str {
184        self.as_str()
185    }
186}
187
188impl FromStr for CachePath {
189    type Err = std::convert::Infallible;
190
191    fn from_str(s: &str) -> Result<Self, Self::Err> {
192        Ok(s.into())
193    }
194}
195
196impl From<&str> for CachePath {
197    fn from(value: &str) -> Self {
198        value.to_string().into()
199    }
200}
201
202impl From<String> for CachePath {
203    fn from(value: String) -> Self {
204        match &*value {
205            ":memory:" => Self::Memory,
206            _ => Self::Path(value)
207        }
208    }
209}
210
211impl ::core::fmt::Debug for InnerCache {
212    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
213        f.debug_struct("InnerCache")
214            .field("path", &self.path)
215            .field("connection", if self.connection.get().is_some() {&"OnceCell(..)"} else {&"OnceCell(<uninit>)"})
216            .finish()
217    }
218}
219
220/// The enum of errors [`Cache::read`] and [`InnerCache::read`] can return.
221#[derive(Debug, Error)]
222pub enum ReadFromCacheError {
223    /// Returned when a call to [`Mutex::lock`] returns an error.
224    #[error("{0}")]
225    MutexPoisonError(String),
226    /// Returned when a [`diesel::result::Error`] is encountered.
227    #[error(transparent)]
228    DieselError(#[from] diesel::result::Error),
229    /// Returned when a [`ConnectCacheError`] is encountered.
230    #[error(transparent)]
231    ConnectCacheError(#[from] ConnectCacheError)
232}
233
234/// The enum of errors [`Cache::read`] and [`InnerCache::read`] can return.
235#[derive(Debug, Error)]
236pub enum WriteToCacheError {
237    /// Returned when a call to [`Mutex::lock`] returns an error.
238    #[error("{0}")]
239    MutexPoisonError(String),
240    /// Returned when a [`diesel::result::Error`] is encountered.
241    #[error(transparent)]
242    DieselError(#[from] diesel::result::Error),
243    /// Returned when a [`ConnectCacheError`] is encountered.
244    #[error(transparent)]
245    ConnectCacheError(#[from] ConnectCacheError)
246}
247
248impl Cache {
249    /// Reads from the cache.
250    /// # Errors
251    /// If the call to [`InnerCache::read`] returns an error, that error is returned.
252    pub fn read(&self, category: &str, key: &str) -> Result<Option<Option<String>>, ReadFromCacheError> {
253        self.0.lock().map_err(|e| ReadFromCacheError::MutexPoisonError(e.to_string()))?.read(category, key)
254    }
255
256    /// Writes to the cache.
257    ///
258    /// If an entry for the `category` and `key` already exists, overwrites it.
259    /// # Errors
260    /// If the call to [`InnerCache::write`] returns an error, that error is returned.
261    pub fn write(&self, category: &str, key: &str, value: Option<&str>) -> Result<(), WriteToCacheError> {
262        self.0.lock().map_err(|e| WriteToCacheError::MutexPoisonError(e.to_string()))?.write(category, key, value)
263    }
264}
265
266/// The enum of errors that [`InnerCache::connect`] can return.
267#[derive(Debug, Error)]
268pub enum ConnectCacheError {
269    /// Returned when a [`diesel::ConnectionError`] is encountered.
270    #[error(transparent)]
271    ConnectionError(#[from] diesel::ConnectionError),
272    /// Returned when a [`std::io::Error`] is encountered.
273    #[error(transparent)]
274    IoError(#[from] std::io::Error),
275    /// Returned when a [`diesel::result::Error`] is encountered.
276    #[error(transparent)]
277    DieselError(#[from] diesel::result::Error)
278}
279
280impl InnerCache {
281    /// Gets the [`CachePath`] of the connection.
282    pub fn path(&self) -> &CachePath {
283        &self.path
284    }
285
286    /// Gets the connection itself, if `self` has been connected via [`Self::connect`] yet.
287    pub fn connection(&mut self) -> Option<&mut SqliteConnection> {
288        self.connection.get_mut()
289    }
290
291    /// Returns the connection, connecting if not already connected.
292    /// # Errors
293    /// If the call to [`std::fs::exists`] to check if the database exists returns an error, that error is returned.
294    ///
295    /// If the call to [`std::fs::File::create_new`] to create the database returns an error, that error is returned.
296    ///
297    /// If the call to [`SqliteConnection::establish`] to connect to the database returns an error, that error is returned.
298    ///
299    /// If the call to [`SqlQuery::execute`] to initialize the database returns an error, that error is returned.
300    #[allow(clippy::missing_panics_doc, reason = "Doesn't panic, but should be replaced with OnceCell::get_or_try_init once that's stable.")]
301    pub fn connect(&mut self) -> Result<&mut SqliteConnection, ConnectCacheError> {
302        debug!(self, InnerCache::connect, self);
303        if self.connection.get().is_none() {
304            let mut needs_init = self.path == CachePath::Memory;
305            if let CachePath::Path(path) = &self.path {
306                if !std::fs::exists(path)? {
307                    needs_init = true;
308                    std::fs::File::create_new(path)?;
309                }
310            }
311            let mut connection = SqliteConnection::establish(self.path.as_str())?;
312            if needs_init {
313                diesel::sql_query(DB_INIT_COMMAND).execute(&mut connection)?;
314            }
315            self.connection.set(connection).map_err(|_| ()).expect("The connection to have just been confirmed unset.");
316        }
317        Ok(self.connection.get_mut().expect("The connection to have just been set."))
318    }
319
320    /// Disconnects from the database.
321    pub fn disconnect(&mut self) {
322        let _ = self.connection.take();
323    }
324
325    /// Reads from the database.
326    /// # Errors
327    /// If the call to [`Self::connect`] returns an error, that error is returned.
328    ///
329    /// If the call to [`RunQueryDsl::load`] returns an error, that error is returned.
330    pub fn read(&mut self, category: &str, key: &str) -> Result<Option<Option<String>>, ReadFromCacheError> {
331        debug!(self, InnerCache::read, self, category, key);
332        Ok(cache::dsl::cache
333            .filter(cache::dsl::category.eq(category))
334            .filter(cache::dsl::key.eq(key))
335            .limit(1)
336            .select(CacheEntry::as_select())
337            .load(self.connect()?)?
338            .first()
339            .map(|cache_entry| cache_entry.value.to_owned()))
340    }
341
342    /// Writes to the database, overwriting the entry the equivalent call to [`Self::read`] would return.
343    /// # Errors
344    /// If the call to [`Self::connect`] returns an error, that error is returned.
345    ///
346    /// If the call to [`RunQueryDsl::get_result`] returns an error, that error is returned.
347    pub fn write(&mut self, category: &str, key: &str, value: Option<&str>) -> Result<(), WriteToCacheError> {
348        debug!(self, InnerCache::write, self, category, key, value);
349        diesel::insert_into(cache::table)
350            .values(NewCacheEntry {category, key, value})
351            .execute(self.connect()?)?;
352        Ok(())
353    }
354}
355
356impl From<InnerCache> for (CachePath, OnceCell<SqliteConnection>) {
357    fn from(value: InnerCache) -> Self {
358        (value.path, value.connection)
359    }
360}