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) publication_authority_path: Option<PathBuf>,
32    pub(super) maintenance: Arc<Mutex<SqliteMaintenanceState>>,
33}
34
35impl SqliteGraphStore {
36    /// Opens a SQLite database and initializes the current schema.
37    pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
38        let path = path.as_ref().to_path_buf();
39        if let Some(parent) = path.parent() {
40            std::fs::create_dir_all(parent)?;
41        }
42
43        let connection = Connection::open(&path)?;
44        configure_writer_connection(&connection)?;
45        code::schema::retention_schema::upgrade_legacy_retention_activity_triggers(&connection)?;
46        if !marker::schema_initialization_is_current(&connection)? {
47            migration::prepare_existing_database(&connection)?;
48            initialization::initialize_schema_for_open(&connection)?;
49        }
50        code::schema::validate_existing_query_indexes(&connection)?;
51        let read_pool = ReadConnectionPool::open(&path)?;
52
53        Ok(Self {
54            connection: Arc::new(Mutex::new(connection)),
55            read_pool: Some(Arc::new(read_pool)),
56            database_path: Some(path),
57            publication_authority_path: None,
58            maintenance: Arc::new(Mutex::new(SqliteMaintenanceState::default())),
59        })
60    }
61
62    /// Opens an in-memory database for isolated tests.
63    pub fn open_in_memory() -> Result<Self, StorageError> {
64        let connection = Connection::open_in_memory()?;
65        configure_writer_connection(&connection)?;
66        initialization::initialize_schema(&connection)?;
67
68        Ok(Self {
69            connection: Arc::new(Mutex::new(connection)),
70            read_pool: None,
71            database_path: None,
72            publication_authority_path: None,
73            maintenance: Arc::new(Mutex::new(SqliteMaintenanceState::default())),
74        })
75    }
76
77    pub(in crate::storage) fn open_with_publication_authority(
78        path: impl AsRef<Path>,
79        authority_path: impl AsRef<Path>,
80    ) -> Result<Self, StorageError> {
81        let mut store = Self::open(path)?;
82        store.publication_authority_path = Some(authority_path.as_ref().to_path_buf());
83        Ok(store)
84    }
85
86    pub(in crate::storage) fn run<T, F>(&self, operation: F) -> StorageFuture<'_, T>
87    where
88        T: Send + 'static,
89        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
90    {
91        let connection = Arc::clone(&self.connection);
92
93        Box::pin(async move {
94            tokio::task::spawn_blocking(move || {
95                let mut guard = connection.lock().map_err(|_| StorageError::LockPoisoned)?;
96
97                operation(&mut guard)
98            })
99            .await?
100        })
101    }
102
103    pub(super) fn run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
104    where
105        T: Send + 'static,
106        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
107    {
108        if let Some(read_pool) = &self.read_pool {
109            let connections = read_pool.connections();
110            return Box::pin(async move {
111                tokio::task::spawn_blocking(move || {
112                    let mut guard = lock_any_read_connection(&connections)?;
113
114                    operation(&mut guard)
115                })
116                .await?
117            });
118        }
119
120        self.run(operation)
121    }
122
123    /// Runs related SELECTs in one deferred SQLite snapshot.
124    pub(super) fn run_read_snapshot<T, F>(&self, operation: F) -> StorageFuture<'_, T>
125    where
126        T: Send + 'static,
127        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
128    {
129        self.run_read(move |connection| {
130            if !connection.is_autocommit() {
131                return Err(StorageError::InvalidInput(
132                    "sqlite read snapshot requires an idle connection".to_owned(),
133                ));
134            }
135            connection.execute_batch("BEGIN DEFERRED TRANSACTION")?;
136            match operation(connection) {
137                Ok(output) => {
138                    if let Err(error) = connection.execute_batch("COMMIT") {
139                        let _ = connection.execute_batch("ROLLBACK");
140                        return Err(StorageError::from(error));
141                    }
142                    Ok(output)
143                }
144                Err(error) => {
145                    let _ = connection.execute_batch("ROLLBACK");
146                    Err(error)
147                }
148            }
149        })
150    }
151
152    pub(super) fn run_read_until<T, F>(
153        &self,
154        deadline: Instant,
155        timeout_message: &'static str,
156        operation: F,
157    ) -> StorageFuture<'_, T>
158    where
159        T: Send + 'static,
160        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
161    {
162        if let Some(read_pool) = &self.read_pool {
163            let connections = read_pool.connections();
164            return Box::pin(async move {
165                tokio::task::spawn_blocking(move || {
166                    let mut guard =
167                        lock_any_read_connection_until(&connections, deadline, timeout_message)?;
168
169                    operation(&mut guard)
170                })
171                .await?
172            });
173        }
174
175        let connection = Arc::clone(&self.connection);
176        Box::pin(async move {
177            tokio::task::spawn_blocking(move || {
178                let mut guard = lock_connection_until(&connection, deadline, timeout_message)?;
179
180                operation(&mut guard)
181            })
182            .await?
183        })
184    }
185
186    pub(super) fn try_run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
187    where
188        T: Send + 'static,
189        F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
190    {
191        if let Some(read_pool) = &self.read_pool {
192            let connections = read_pool.connections();
193            return Box::pin(async move {
194                tokio::task::spawn_blocking(move || {
195                    let mut guard = try_lock_any_read_connection(&connections)?;
196
197                    operation(&mut guard)
198                })
199                .await?
200            });
201        }
202
203        let connection = Arc::clone(&self.connection);
204        Box::pin(async move {
205            tokio::task::spawn_blocking(move || {
206                let mut guard = match connection.try_lock() {
207                    Ok(guard) => guard,
208                    Err(TryLockError::Poisoned(_)) => return Err(StorageError::LockPoisoned),
209                    Err(TryLockError::WouldBlock) => {
210                        return Err(StorageError::Busy(
211                            "sqlite write connection is currently occupied".to_owned(),
212                        ));
213                    }
214                };
215
216                operation(&mut guard)
217            })
218            .await?
219        })
220    }
221
222    pub(in crate::storage) fn import_code_repository_from_database(
223        &self,
224        source_path: PathBuf,
225        repository_id: String,
226        source_scope: Option<String>,
227    ) -> StorageFuture<'_, ()> {
228        self.run(move |connection| {
229            code::import_repository_from_database(
230                connection,
231                &source_path,
232                &repository_id,
233                source_scope.as_deref(),
234            )
235        })
236    }
237
238    pub(in crate::storage) fn code_repository_totals_excluding(
239        &self,
240        excluded_repository_ids: Vec<String>,
241    ) -> StorageFuture<'_, crate::domain::CodeRepositoryTotals> {
242        self.run_read(move |connection| {
243            code::repository_totals_excluding(connection, &excluded_repository_ids)
244        })
245    }
246
247    pub(in crate::storage) fn prune_code_repository_scopes_with_retained(
248        &self,
249        request: crate::storage::CodeScopeRetentionRequest,
250        extra_retained_scopes: Vec<String>,
251    ) -> StorageFuture<'_, crate::domain::CodeScopeRetentionSummary> {
252        self.run(move |connection| {
253            code::prune_scopes_with_retained(connection, request, extra_retained_scopes)
254        })
255    }
256
257    pub(in crate::storage) fn complete_code_repository_retention(
258        &self,
259        repository_id: String,
260        cutoff_ms: u64,
261    ) -> StorageFuture<'_, bool> {
262        self.run(move |connection| {
263            code::complete_repository_retention(connection, &repository_id, cutoff_ms)
264        })
265    }
266
267    pub(in crate::storage) fn repository_retention_republished_initial_scope(
268        &self,
269        repository_id: String,
270        initial_scope: String,
271        cutoff_ms: u64,
272        cutoff_publication_generation: u64,
273    ) -> StorageFuture<'_, Option<String>> {
274        self.run(move |connection| {
275            code::repository_retention_republished_initial_scope(
276                connection,
277                &repository_id,
278                &initial_scope,
279                cutoff_ms,
280                cutoff_publication_generation,
281            )
282        })
283    }
284}
285
286#[cfg(test)]
287#[path = "mod_tests.rs"]
288mod mod_tests;