swiftide_integrations/duckdb/
node_cache.rs

1use anyhow::Context as _;
2use async_trait::async_trait;
3use swiftide_core::{
4    NodeCache,
5    indexing::{Chunk, Node},
6};
7
8use super::Duckdb;
9
10macro_rules! unwrap_or_log {
11    ($result:expr) => {
12        match $result {
13            Ok(value) => value,
14            Err(e) => {
15                tracing::error!("Error: {:#}", e);
16                debug_assert!(
17                    true,
18                    "Duckdb should not give errors unless in very weird situations; this is a bug: {:#}",
19                    e
20                );
21                return false;
22            }
23        }
24    };
25}
26
27#[async_trait]
28impl<T: Chunk> NodeCache for Duckdb<T> {
29    type Input = T;
30
31    async fn get(&self, node: &Node<T>) -> bool {
32        unwrap_or_log!(
33            self.lazy_create_cache()
34                .await
35                .context("failed to create cache table")
36        );
37
38        let sql = format!(
39            "SELECT EXISTS(SELECT 1 FROM {} WHERE uuid = ?)",
40            &self.cache_table
41        );
42
43        let lock = self.connection.lock().unwrap();
44        let mut stmt = unwrap_or_log!(
45            lock.prepare(&sql)
46                .context("Failed to prepare duckdb statement for persist")
47        );
48
49        let present = unwrap_or_log!(
50            stmt.query_map([self.node_key(node)], |row| row.get::<_, bool>(0))
51                .context("failed to query for documents")
52        )
53        .next()
54        .transpose();
55
56        unwrap_or_log!(present).unwrap_or(false)
57    }
58
59    async fn set(&self, node: &Node<T>) {
60        if let Err(err) = self
61            .lazy_create_cache()
62            .await
63            .context("failed to create cache table")
64        {
65            tracing::error!("Failed to create cache table: {:#}", err);
66            return;
67        }
68
69        let sql = format!(
70            "INSERT INTO {} (uuid, path) VALUES (?, ?) ON CONFLICT (uuid) DO NOTHING",
71            &self.cache_table
72        );
73
74        let lock = self.connection.lock().unwrap();
75        let mut stmt = match lock
76            .prepare(&sql)
77            .context("Failed to prepare duckdb statement for cache set")
78        {
79            Ok(stmt) => stmt,
80            Err(err) => {
81                tracing::error!(
82                    "Failed to prepare duckdb statement for cache set: {:#}",
83                    err
84                );
85                return;
86            }
87        };
88
89        if let Err(err) = stmt
90            .execute([self.node_key(node), node.path.to_string_lossy().into()])
91            .context("failed to insert into cache table")
92        {
93            tracing::error!("Failed to insert into cache table: {:#}", err);
94        }
95    }
96
97    async fn clear(&self) -> anyhow::Result<()> {
98        let sql = format!("DROP TABLE IF EXISTS {}", &self.cache_table);
99        let lock = self.connection.lock().unwrap();
100        let mut stmt = lock
101            .prepare(&sql)
102            .context("Failed to prepare duckdb statement for cache clear")?;
103
104        stmt.execute([]).context("failed to delete cache table")?;
105
106        Ok(())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use swiftide_core::indexing::TextNode;
114
115    fn setup_duckdb() -> Duckdb {
116        Duckdb::builder()
117            .connection(duckdb::Connection::open_in_memory().unwrap())
118            .build()
119            .unwrap()
120    }
121
122    #[tokio::test]
123    async fn test_get_set() {
124        let duckdb = setup_duckdb();
125        let node = TextNode::new("test_get_set");
126
127        assert!(!duckdb.get(&node).await);
128        duckdb.set(&node).await;
129        assert!(duckdb.get(&node).await);
130    }
131
132    #[tokio::test]
133    async fn test_clear() {
134        let duckdb = setup_duckdb();
135        let node = TextNode::new("test_clear");
136
137        duckdb.set(&node).await;
138        assert!(duckdb.get(&node).await);
139        duckdb.clear().await.unwrap();
140        assert!(!duckdb.get(&node).await);
141    }
142}
swiftide_integrations/duckdb/node_cache.rs

swiftide_integrations/duckdb/
node_cache.rs
