mecomp_storage/db/
mod.rs

1#[cfg(feature = "db")]
2pub mod crud;
3#[cfg(feature = "db")]
4pub mod health;
5#[cfg(feature = "db")]
6pub(crate) mod queries;
7pub mod schemas;
8
9#[cfg(feature = "db")]
10use surrealdb::{Surreal, engine::local::Db};
11
12#[cfg(feature = "db")]
13#[cfg(not(tarpaulin_include))]
14static DB_DIR: once_cell::sync::OnceCell<std::path::PathBuf> = once_cell::sync::OnceCell::new();
15#[cfg(feature = "db")]
16#[cfg(not(tarpaulin_include))]
17static TEMP_DB_DIR: once_cell::sync::Lazy<tempfile::TempDir> = once_cell::sync::Lazy::new(|| {
18    tempfile::tempdir().expect("Failed to create temporary directory")
19});
20
21/// NOTE: if you change this, you must go through the schemas and update the index analyzer names
22pub const FULL_TEXT_SEARCH_ANALYZER_NAME: &str = "custom_analyzer";
23
24/// Set the path to the database.
25///
26/// # Errors
27///
28/// This function will return an error if the path cannot be set.
29#[cfg(feature = "db")]
30#[allow(clippy::missing_inline_in_public_items)]
31pub fn set_database_path(path: std::path::PathBuf) -> Result<(), crate::errors::Error> {
32    DB_DIR
33        .set(path)
34        .map_err(crate::errors::Error::DbPathSetError)?;
35    log::info!("Primed database path");
36    Ok(())
37}
38
39/// Initialize the database with the necessary tables.
40///
41/// # Errors
42///
43/// This function will return an error if the database cannot be initialized.
44#[cfg(feature = "db")]
45#[allow(clippy::missing_inline_in_public_items)]
46pub async fn init_database() -> surrealdb::Result<Surreal<Db>> {
47    let db = Surreal::new(DB_DIR
48        .get().cloned()
49        .unwrap_or_else(|| {
50            log::warn!("DB_DIR not set, defaulting to a temporary directory `{}`, this is likely a bug because `init_database` should be called before `db`", TEMP_DB_DIR.path().display());
51            TEMP_DB_DIR.path()
52            .to_path_buf()
53        })).await?;
54
55    db.use_ns("mecomp").use_db("music").await?;
56
57    register_custom_analyzer(&db).await?;
58    surrealqlx::register_tables!(
59        &db,
60        schemas::album::Album,
61        schemas::artist::Artist,
62        schemas::song::Song,
63        schemas::collection::Collection,
64        schemas::playlist::Playlist,
65        schemas::dynamic::DynamicPlaylist
66    )?;
67    #[cfg(feature = "analysis")]
68    surrealqlx::register_tables!(&db, schemas::analysis::Analysis)?;
69
70    queries::relations::define_relation_tables(&db).await?;
71
72    Ok(db)
73}
74
75#[cfg(feature = "db")]
76pub(crate) async fn register_custom_analyzer<C>(db: &Surreal<C>) -> surrealdb::Result<()>
77where
78    C: surrealdb::Connection,
79{
80    use queries::define_analyzer;
81    use surrealdb::sql::Tokenizer;
82
83    db.query(define_analyzer(
84        FULL_TEXT_SEARCH_ANALYZER_NAME,
85        Some(Tokenizer::Class),
86        &[
87            "ascii",
88            "lowercase",
89            "edgengram(1, 10)",
90            "snowball(English)",
91        ],
92    ))
93    .await?;
94
95    Ok(())
96}
97
98#[cfg(test)]
99mod test {
100    use super::schemas::{
101        album::Album, artist::Artist, collection::Collection, dynamic::DynamicPlaylist,
102        playlist::Playlist, song::Song,
103    };
104    use super::*;
105
106    use surrealdb::engine::local::Mem;
107    use surrealqlx::traits::Table;
108
109    #[tokio::test]
110    async fn test_register_tables() -> anyhow::Result<()> {
111        // use an in-memory db for testing
112        let db = Surreal::new::<Mem>(()).await?;
113        db.use_ns("test").use_db("test").await?;
114
115        // register the custom analyzer
116        register_custom_analyzer(&db).await?;
117
118        // first we init all the table to ensure that the queries made by the macro work without error
119        <Album as Table>::init_table(&db).await?;
120        <Artist as Table>::init_table(&db).await?;
121        <Song as Table>::init_table(&db).await?;
122        <Collection as Table>::init_table(&db).await?;
123        <Playlist as Table>::init_table(&db).await?;
124        <DynamicPlaylist as Table>::init_table(&db).await?;
125
126        // then we init the relation tables
127        queries::relations::define_relation_tables(&db).await?;
128
129        // then we try initializing one of the tables again to ensure that initialization won't mess with existing tables/data
130        <Album as Table>::init_table(&db).await?;
131
132        Ok(())
133    }
134}
135
136#[cfg(test)]
137mod minimal_reproduction {
138    //! This module contains minimal reproductions of issues from MECOMPs past.
139    //! They exist to ensure that the issues are indeed fixed.
140    use serde::{Deserialize, Serialize};
141    use surrealdb::{RecordId, Surreal, engine::local::Mem, method::Stats};
142
143    use crate::db::queries::generic::{Count, count};
144
145    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
146    struct User {
147        id: RecordId,
148        name: String,
149        age: i32,
150        favorite_numbers: [i32; 7],
151    }
152
153    static SCHEMA_SQL: &str = r"
154    BEGIN;
155    DEFINE TABLE users SCHEMAFULL;
156    COMMIT;
157    BEGIN;
158    DEFINE FIELD id ON users TYPE record;
159    DEFINE FIELD name ON users TYPE string;
160    DEFINE FIELD age ON users TYPE int;
161    DEFINE FIELD favorite_numbers ON users TYPE array<int>;
162    COMMIT;
163    BEGIN;
164    DEFINE INDEX users_name_unique_index ON users FIELDS name UNIQUE;
165    DEFINE INDEX users_age_normal_index ON users FIELDS age;
166    DEFINE INDEX users_favorite_numbers_vector_index ON users FIELDS favorite_numbers MTREE DIMENSION 7;
167    ";
168
169    #[tokio::test]
170    async fn minimal_reproduction() {
171        let db = Surreal::new::<Mem>(()).await.unwrap();
172        db.use_ns("test").use_db("test").await.unwrap();
173
174        db.query(SCHEMA_SQL).await.unwrap();
175
176        let cnt: Option<Count> = db
177            // new syntax
178            .query(count("users"))
179            .await
180            .unwrap()
181            .take(0)
182            .unwrap();
183
184        assert_eq!(cnt, Some(Count::new(0)));
185
186        let john_id = RecordId::from(("users", "0"));
187        let john = User {
188            id: john_id.clone(),
189            name: "John".to_string(),
190            age: 42,
191            favorite_numbers: [1, 2, 3, 4, 5, 6, 7],
192        };
193
194        let sally_id = RecordId::from(("users", "1"));
195        let sally = User {
196            id: sally_id.clone(),
197            name: "Sally".to_string(),
198            age: 24,
199            favorite_numbers: [8, 9, 10, 11, 12, 13, 14],
200        };
201
202        let result: Option<User> = db
203            .create(john_id.clone())
204            .content(john.clone())
205            .await
206            .unwrap();
207
208        assert_eq!(result, Some(john.clone()));
209
210        let result: Option<User> = db
211            .create(sally_id.clone())
212            .content(sally.clone())
213            .await
214            .unwrap();
215
216        assert_eq!(result, Some(sally.clone()));
217
218        let result: Option<User> = db.select(john_id).await.unwrap();
219
220        assert_eq!(result, Some(john.clone()));
221
222        const NUMBER_OF_USERS: usize = 100;
223        // create like 100 more users
224        for i in 2..NUMBER_OF_USERS {
225            let user_id = RecordId::from(("users", i.to_string()));
226            let user = User {
227                id: user_id.clone(),
228                name: format!("User {}", i),
229                age: i as i32,
230                favorite_numbers: [i as i32; 7],
231            };
232            let _: Option<User> = db.create(user_id.clone()).content(user).await.unwrap();
233        }
234
235        let mut resp_new = db
236            // new syntax
237            .query("SELECT count() FROM users GROUP ALL")
238            .with_stats()
239            .await
240            .unwrap();
241        dbg!(&resp_new);
242        let res = resp_new.take(0).unwrap();
243        let cnt: Option<Count> = res.1.unwrap();
244        assert_eq!(cnt, Some(Count::new(NUMBER_OF_USERS)));
245        let stats_new: Stats = res.0;
246
247        let mut resp_old = db
248            // old syntax
249            .query("RETURN array::len((SELECT * FROM users))")
250            .with_stats()
251            .await
252            .unwrap();
253        dbg!(&resp_old);
254        let res = resp_old.take(0).unwrap();
255        let cnt: Option<usize> = res.1.unwrap();
256        assert_eq!(cnt, Some(NUMBER_OF_USERS));
257        let stats_old: Stats = res.0;
258
259        // just a check to ensure the new syntax is faster
260        assert!(stats_new.execution_time.unwrap() < stats_old.execution_time.unwrap());
261
262        let result: Vec<User> = db.delete("users").await.unwrap();
263
264        assert_eq!(result.len(), NUMBER_OF_USERS);
265        assert!(result.contains(&john), "Result does not contain 'john'");
266        assert!(result.contains(&sally), "Result does not contain 'sally'");
267    }
268}
mecomp_storage/db/mod.rs

mecomp_storage/db/
mod.rs
