vrc_log/
cache.rs

1use std::{collections::HashMap, path::PathBuf};
2
3use anyhow::Result;
4use itertools::Itertools;
5use rusqlite::{Connection as RusqliteConnection, named_params, params_from_iter};
6use tokio_rusqlite_new::Connection;
7
8use crate::vrchat::VRCHAT_LOW_PATH;
9
10pub struct Cache {
11    connection: Connection,
12}
13
14pub type AvatarIDWithProvider<S> = (S, u32);
15
16impl Cache {
17    /// # Errors
18    /// Will return `Err` if `sqlite::open` errors
19    pub async fn new() -> Result<Self> {
20        debug!("Trying to open SQLite cache database.");
21        Self::new_at_location(&VRCHAT_LOW_PATH.join("avatars.sqlite")).await
22    }
23
24    /// # Errors
25    /// Will return `Err` if `sqlite::open` errors
26    pub async fn new_at_location(path: &PathBuf) -> Result<Self> {
27        debug!("Trying to open SQLite cache database.");
28        let connection = Connection::open(path).await?;
29
30        connection
31            .call(|connection| Self::setup_database(connection))
32            .await?;
33
34        Ok(Self { connection })
35    }
36
37    /// # Errors
38    /// Will return `Err` if `sqlite::open` errors
39    pub async fn new_in_memory() -> Result<Self> {
40        let connection = Connection::open_in_memory().await?;
41        connection
42            .call(|connection| Self::setup_database(connection))
43            .await?;
44        Ok(Self { connection })
45    }
46
47    fn setup_database(connection: &RusqliteConnection) -> Result<(), rusqlite::Error> {
48        let query = "CREATE TABLE avatars (
49                    id TEXT PRIMARY KEY,
50                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
51                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
52                    provider_bits INT DEFAULT 0
53                )";
54
55        debug!("Trying to create avatars table...");
56        if connection.execute(query, []).is_err() {
57            debug!("The avatars table already exists.");
58
59            let mut statement = connection.prepare("PRAGMA table_info(avatars)")?;
60            let columns = statement
61                .query_map([], |row| row.get::<_, String>(1))?
62                .collect::<Result<Vec<_>, _>>()?;
63
64            if !columns.contains(&"created_at".to_string()) {
65                debug!("Trying to create the created_at column.");
66                #[rustfmt::skip]
67                        connection.execute("
68                            ALTER TABLE avatars
69                            ADD COLUMN created_at DATETIME
70                        ", [])?;
71            }
72
73            debug!("Updating all rows with missing created_at");
74            #[rustfmt::skip]
75                    connection.execute("
76                        UPDATE avatars
77                        SET created_at = CURRENT_TIMESTAMP
78                        WHERE created_at IS NULL
79                    ", [])?;
80
81            if !columns.contains(&"updated_at".to_string()) {
82                debug!("Trying to create the updated_at column.");
83                #[rustfmt::skip]
84                        connection.execute("
85                            ALTER TABLE avatars
86                            ADD COLUMN updated_at DATETIME
87                        ", [])?;
88            }
89            if !columns.contains(&"provider_bits".to_string()) {
90                debug!("Trying to create the provider_bits column.");
91                #[rustfmt::skip]
92                        connection.execute("
93                            ALTER TABLE avatars
94                            ADD COLUMN provider_bits INT DEFAULT 0
95                        ", [])?;
96            }
97
98            debug!("Updating all rows with missing updated_at");
99            #[rustfmt::skip]
100                    connection.execute("
101                        UPDATE avatars
102                        SET updated_at = datetime('now', '-31 days')
103                        WHERE updated_at IS NULL
104                    ", [])?;
105        }
106
107        // Speed up queries on large databases
108        debug!("Trying to create an updated_at index.");
109        #[rustfmt::skip]
110                connection.execute("
111                    CREATE INDEX IF NOT EXISTS idx_avatars_updated_at
112                    ON avatars(updated_at)
113                ", [])?;
114
115        debug!("Trying to create an id index.");
116        #[rustfmt::skip]
117                connection.execute(
118                    "
119                    CREATE INDEX IF NOT EXISTS idx_avatars_id
120                    ON avatars(id)
121                ", [])?;
122
123        if let Ok(mut statement) = connection.prepare("SELECT COUNT(*) FROM avatars")
124            && let Ok(count) = statement.query_row([], |row| row.get::<_, i64>(0))
125        {
126            info!("{} Cached Avatars", count);
127        }
128
129        Ok(())
130    }
131
132    /// # Errors
133    /// Will return `Err` if `Connection::call(...)` errors
134    pub async fn store_avatar_ids_with_providers<
135        S: ToString,
136        I: IntoIterator<Item = AvatarIDWithProvider<S>>,
137    >(
138        &self,
139        insertables: I,
140    ) -> Result<()> {
141        let query = "
142            INSERT INTO avatars (id, provider_bits, created_at, updated_at)
143            VALUES (:id, :providers, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
144            ON CONFLICT (id) DO UPDATE
145                SET updated_at = CURRENT_TIMESTAMP,
146                    provider_bits = :providers
147        ";
148
149        let insertables: Vec<_> = insertables
150            .into_iter()
151            .map(|a| (a.0.to_string(), a.1))
152            .collect();
153
154        self.connection
155            .call(|c| -> Result<(), rusqlite::Error> {
156                let tx = c.transaction()?;
157                for (id, providers) in insertables {
158                    let _ = tx.execute(
159                        query,
160                        named_params! {
161                            ":id": id,
162                            ":providers": providers
163                        },
164                    );
165                }
166                tx.commit()
167            })
168            .await
169            .map_err(anyhow::Error::from)
170    }
171
172    const CHUNK_SIZE: usize = 950; // sqlite parameter limit is 999
173
174    /// # Errors
175    /// Will return `Err` if `Connection::call(...)` errors
176    pub async fn check_all_ids<I: IntoIterator<Item = String>>(
177        &self,
178        ids: I,
179    ) -> Result<HashMap<String, u32>> {
180        let ids: Vec<_> = ids.into_iter().collect();
181        self.connection
182            .call(|c| -> Result<_, rusqlite::Error> {
183                let mut output = HashMap::new();
184
185                for chunk in &ids.into_iter().chunks(Self::CHUNK_SIZE) {
186                    let chunk: Vec<String> = chunk.collect();
187                    for id in &chunk {
188                        output.insert(id.clone(), 0);
189                    }
190                    let found_ids = Self::check_batch_ids(c, chunk.into_iter())?;
191                    output.extend(found_ids);
192                }
193
194                Ok(output)
195            })
196            .await
197            .map_err(|e| anyhow::anyhow!(e))
198    }
199
200    fn check_batch_ids<I: Iterator<Item = String>>(
201        conn: &RusqliteConnection,
202        chunk: I,
203    ) -> std::result::Result<HashMap<String, u32>, rusqlite::Error> {
204        let chunk: Vec<_> = chunk.collect();
205        assert!(chunk.len() <= Self::CHUNK_SIZE);
206
207        let placeholders = std::iter::repeat_n("?", chunk.len())
208            .collect::<Vec<_>>()
209            .join(",");
210        let sql = format!(
211            "SELECT id, provider_bits FROM avatars WHERE id IN ({placeholders}) AND updated_at >= datetime('now', '-30 days')"
212        );
213
214        let mut stmt = conn.prepare(&sql)?;
215        stmt.query_map(params_from_iter(chunk.iter()), |row| {
216            Ok((row.get::<_, String>(0)?, row.get::<_, u32>(1)?))
217        })?
218        .collect::<Result<HashMap<_, _>, _>>()
219    }
220}
221
222// #[cfg(test)]
223mod tests {
224    use super::Cache;
225    /// Helper to create a cache for tests
226    #[allow(dead_code)]
227    async fn cache() -> Cache {
228        Cache::new_in_memory().await.unwrap()
229    }
230
231    #[tokio::test]
232    async fn creates_database_and_table() {
233        let cache = cache().await;
234
235        // Simple sanity check: storing should not fail
236        cache
237            .store_avatar_ids_with_providers(vec![("avatar_1", 1u32)].into_iter())
238            .await
239            .unwrap();
240    }
241
242    #[tokio::test]
243    async fn inserts_and_reads_avatar_ids() {
244        let cache = cache().await;
245
246        cache
247            .store_avatar_ids_with_providers(
248                vec![("avatar_a", 1u32), ("avatar_b", 2u32)].into_iter(),
249            )
250            .await
251            .unwrap();
252
253        let result = cache
254            .check_all_ids(vec!["avatar_a".into(), "avatar_b".into()])
255            .await
256            .unwrap();
257
258        assert_eq!(result.len(), 2);
259        assert_eq!(result["avatar_a"], 1);
260        assert_eq!(result["avatar_b"], 2);
261    }
262
263    #[tokio::test]
264    async fn returns_none_for_missing_ids() {
265        let cache = cache().await;
266
267        let result = cache
268            .check_all_ids(vec!["missing_avatar".into()])
269            .await
270            .unwrap();
271
272        assert_eq!(result["missing_avatar"], 0);
273    }
274
275    #[tokio::test]
276    async fn updates_provider_bits_on_conflict() {
277        let cache = cache().await;
278
279        cache
280            .store_avatar_ids_with_providers(vec![("avatar_x", 1u32)].into_iter())
281            .await
282            .unwrap();
283
284        cache
285            .store_avatar_ids_with_providers(vec![("avatar_x", 42u32)].into_iter())
286            .await
287            .unwrap();
288
289        let result = cache.check_all_ids(vec!["avatar_x".into()]).await.unwrap();
290
291        assert_eq!(result["avatar_x"], 42);
292    }
293
294    #[tokio::test]
295    async fn respects_chunking_limits() {
296        let cache = cache().await;
297
298        #[allow(clippy::cast_possible_truncation)]
299        let ids: Vec<(String, u32)> = (0..(Cache::CHUNK_SIZE + 10))
300            .map(|i| (format!("avatar_{i}"), i as u32))
301            .collect();
302
303        cache
304            .store_avatar_ids_with_providers(ids.iter().map(|(id, p)| (id.as_str(), *p)))
305            .await
306            .unwrap();
307
308        let result = cache
309            .check_all_ids(ids.iter().map(|(id, _)| id.clone()))
310            .await
311            .unwrap();
312
313        assert_eq!(result.len(), ids.len());
314
315        for (id, provider) in ids {
316            assert_eq!(result[&id], provider);
317        }
318    }
319
320    #[tokio::test]
321    async fn ignores_entries_older_than_30_days() {
322        let cache = cache().await;
323
324        // Insert normally
325        cache
326            .store_avatar_ids_with_providers(vec![("old_avatar", 1u32)].into_iter())
327            .await
328            .unwrap();
329
330        // Manually age the entry
331        cache
332            .connection
333            .call(|c| {
334                c.execute(
335                    "UPDATE avatars
336                     SET updated_at = datetime('now', '-31 days')
337                     WHERE id = 'old_avatar'",
338                    [],
339                )
340            })
341            .await
342            .unwrap();
343
344        let result = cache
345            .check_all_ids(vec!["old_avatar".into()])
346            .await
347            .unwrap();
348
349        // Exists, but filtered out by age
350        assert_eq!(result["old_avatar"], 0);
351    }
352}
vrc_log/cache.rs

vrc_log/
cache.rs
