Skip to main content

relay_knowledge/storage/sqlite/store/
mod.rs

1mod implementations;
2
3use std::{
4    path::{Path, PathBuf},
5    sync::{Arc, Mutex, TryLockError},
6    time::Instant,
7};
8
9use rusqlite::Connection;
10
11use crate::storage::{StorageError, StorageFuture};
12
13use super::{
14    code,
15    connection_runtime::{
16        maintenance::{SqliteMaintenanceState, configure_writer_connection},
17        read_pool::{
18            ReadConnectionPool, lock_any_read_connection, lock_any_read_connection_until,
19            lock_connection_until, try_lock_any_read_connection,
20        },
21    },
22    schema::{initialization, marker, migration},
23};
24
25/// SQLite implementation of graph facts, mutation log, and index metadata.
26#[derive(Debug, Clone)]
27pub struct SqliteGraphStore {
28    pub(super) connection: Arc<Mutex<Connection>>,
29    pub(super) read_pool: Option<Arc<ReadConnectionPool>>,
30    pub(super) database_path: Option<PathBuf>,
31    pub(super) maintenance: Arc<Mutex<SqliteMaintenanceState>>,
32}
33
34impl SqliteGraphStore {
35    /// Opens a SQLite database and initializes the current schema.
36    pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
37        let path = path.as_ref().to_path_buf();
38        if let Some(parent) = path.parent() {
39            std::fs::create_dir_all(parent)?;
40        }
41
42        let connection = Connection::open(&path)?;
43        configure_writer_connection(&connection)?;
44        if !marker::schema_initialization_is_current(&connection)? {
45            migration::prepare_existing_database(&connection)?;
46            initialization::initialize_schema(&connection)?;
47            marker::mark_schema_initialization_current(&connection)?;
48        }
49        let read_pool = ReadConnectionPool::open(&path)?;
50
51        Ok(Self {
52            connection: Arc::new(Mutex::new(connection)),
53            read_pool: Some(Arc::new(read_pool)),
54            database_path: Some(path),
55            maintenance: Arc::new(Mutex::new(SqliteMaintenanceState::default())),
56        })
57    }
58
59    /// Opens an in-memory database for isolated tests.
60    pub fn open_in_memory() -> Result<Self, StorageError> {
61        let connection = Connection::open_in_memory()?;
62        configure_writer_connection(&connection)?;
63        initialization::initialize_schema(&connection)?;
64
65        Ok(Self {
66            connection: Arc::new(Mutex::new(connection)),
67            read_pool: None,
68            database_path: None,
69            maintenance: Arc::new(Mutex::new(SqliteMaintenanceState::default())),
70        })
71    }
72
73    pub(in crate::storage) fn run<T, F>(&self, operation: F) -> StorageFuture<'_, T>
74    where
75        T: Send + 'static,
76        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
77    {
78        let connection = Arc::clone(&self.connection);
79
80        Box::pin(async move {
81            tokio::task::spawn_blocking(move || {
82                let mut guard = connection.lock().map_err(|_| StorageError::LockPoisoned)?;
83
84                operation(&mut guard)
85            })
86            .await?
87        })
88    }
89
90    pub(super) fn run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
91    where
92        T: Send + 'static,
93        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
94    {
95        if let Some(read_pool) = &self.read_pool {
96            let connections = read_pool.connections();
97            return Box::pin(async move {
98                tokio::task::spawn_blocking(move || {
99                    let mut guard = lock_any_read_connection(&connections)?;
100
101                    operation(&mut guard)
102                })
103                .await?
104            });
105        }
106
107        self.run(operation)
108    }
109
110    pub(super) fn run_read_until<T, F>(
111        &self,
112        deadline: Instant,
113        timeout_message: &'static str,
114        operation: F,
115    ) -> StorageFuture<'_, T>
116    where
117        T: Send + 'static,
118        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
119    {
120        if let Some(read_pool) = &self.read_pool {
121            let connections = read_pool.connections();
122            return Box::pin(async move {
123                tokio::task::spawn_blocking(move || {
124                    let mut guard =
125                        lock_any_read_connection_until(&connections, deadline, timeout_message)?;
126
127                    operation(&mut guard)
128                })
129                .await?
130            });
131        }
132
133        let connection = Arc::clone(&self.connection);
134        Box::pin(async move {
135            tokio::task::spawn_blocking(move || {
136                let mut guard = lock_connection_until(&connection, deadline, timeout_message)?;
137
138                operation(&mut guard)
139            })
140            .await?
141        })
142    }
143
144    pub(super) fn try_run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
145    where
146        T: Send + 'static,
147        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
148    {
149        if let Some(read_pool) = &self.read_pool {
150            let connections = read_pool.connections();
151            return Box::pin(async move {
152                tokio::task::spawn_blocking(move || {
153                    let mut guard = try_lock_any_read_connection(&connections)?;
154
155                    operation(&mut guard)
156                })
157                .await?
158            });
159        }
160
161        let connection = Arc::clone(&self.connection);
162        Box::pin(async move {
163            tokio::task::spawn_blocking(move || {
164                let mut guard = match connection.try_lock() {
165                    Ok(guard) => guard,
166                    Err(TryLockError::Poisoned(_)) => return Err(StorageError::LockPoisoned),
167                    Err(TryLockError::WouldBlock) => {
168                        return Err(StorageError::Busy(
169                            "sqlite write connection is currently occupied".to_owned(),
170                        ));
171                    }
172                };
173
174                operation(&mut guard)
175            })
176            .await?
177        })
178    }
179
180    pub(in crate::storage) fn import_code_repository_from_database(
181        &self,
182        source_path: PathBuf,
183        repository_id: String,
184        source_scope: Option<String>,
185    ) -> StorageFuture<'_, ()> {
186        self.run(move |connection| {
187            code::import_repository_from_database(
188                connection,
189                &source_path,
190                &repository_id,
191                source_scope.as_deref(),
192            )
193        })
194    }
195
196    pub(in crate::storage) fn code_repository_totals_excluding(
197        &self,
198        excluded_repository_ids: Vec<String>,
199    ) -> StorageFuture<'_, crate::domain::CodeRepositoryTotals> {
200        self.run_read(move |connection| {
201            code::repository_totals_excluding(connection, &excluded_repository_ids)
202        })
203    }
204
205    pub(in crate::storage) fn prune_code_repository_scopes_with_retained(
206        &self,
207        request: crate::storage::CodeScopeRetentionRequest,
208        extra_retained_scopes: Vec<String>,
209    ) -> StorageFuture<'_, crate::domain::CodeScopeRetentionSummary> {
210        self.run(move |connection| {
211            code::prune_scopes_with_retained(connection, request, extra_retained_scopes)
212        })
213    }
214}
215
216#[cfg(test)]
217mod mod_tests;