aster_forge_db_migration/
search.rs

1use sea_orm_migration::prelude::*;
2use sea_orm_migration::sea_orm::ConnectionTrait;
3use sea_orm_migration::sea_query::{
4    Alias, IntoIndexColumn, PostgresQueryBuilder, extension::postgres::Extension,
5};
6
7/// `SQLite` FTS5 virtual table and synchronization-trigger names used by one migration.
8pub struct SqliteFtsConfig<'a> {
9    pub virtual_table: &'a str,
10    pub source_table: &'a str,
11    pub columns: &'a [&'a str],
12    pub insert_trigger: &'a str,
13    pub delete_trigger: &'a str,
14    pub update_trigger: &'a str,
15}
16
17/// Builds the `SQLite` FTS5 table, backfill, and synchronization trigger statements.
18///
19/// # Errors
20///
21/// Returns an error when the configuration contains an invalid identifier or no indexed columns.
22pub fn sqlite_fts_up_statements(config: &SqliteFtsConfig<'_>) -> Result<Vec<String>, DbErr> {
23    validate_sqlite_fts_config(config)?;
24    let column_list = config.columns.join(", ");
25    let new_values = config
26        .columns
27        .iter()
28        .map(|column| format!("new.{column}"))
29        .collect::<Vec<_>>()
30        .join(", ");
31    let update_assignments = config
32        .columns
33        .iter()
34        .map(|column| format!("{column} = new.{column}"))
35        .collect::<Vec<_>>()
36        .join(", ");
37
38    Ok(vec![
39        format!(
40            "CREATE VIRTUAL TABLE IF NOT EXISTS {} USING fts5({}, tokenize='trigram')",
41            config.virtual_table, column_list
42        ),
43        format!("DELETE FROM {}", config.virtual_table),
44        format!(
45            "INSERT INTO {}(rowid, {}) SELECT id, {} FROM {}",
46            config.virtual_table, column_list, column_list, config.source_table
47        ),
48        format!(
49            "CREATE TRIGGER IF NOT EXISTS {} AFTER INSERT ON {} BEGIN \
50             INSERT INTO {}(rowid, {}) VALUES (new.id, {}); END",
51            config.insert_trigger,
52            config.source_table,
53            config.virtual_table,
54            column_list,
55            new_values,
56        ),
57        format!(
58            "CREATE TRIGGER IF NOT EXISTS {} AFTER DELETE ON {} BEGIN \
59             DELETE FROM {} WHERE rowid = old.id; END",
60            config.delete_trigger, config.source_table, config.virtual_table
61        ),
62        format!(
63            "CREATE TRIGGER IF NOT EXISTS {} AFTER UPDATE OF {} ON {} BEGIN \
64             UPDATE {} SET {} WHERE rowid = new.id; END",
65            config.update_trigger,
66            column_list,
67            config.source_table,
68            config.virtual_table,
69            update_assignments,
70        ),
71    ])
72}
73
74/// Builds the `SQLite` statements that remove FTS synchronization and the virtual table.
75///
76/// # Errors
77///
78/// Returns an error when the configuration contains an invalid identifier or no indexed columns.
79pub fn sqlite_fts_down_statements(config: &SqliteFtsConfig<'_>) -> Result<Vec<String>, DbErr> {
80    validate_sqlite_fts_config(config)?;
81    Ok(vec![
82        format!("DROP TRIGGER IF EXISTS {}", config.insert_trigger),
83        format!("DROP TRIGGER IF EXISTS {}", config.delete_trigger),
84        format!("DROP TRIGGER IF EXISTS {}", config.update_trigger),
85        format!("DROP TABLE IF EXISTS {}", config.virtual_table),
86    ])
87}
88
89/// Executes generated `SQLite` migration statements in order with caller-provided error context.
90///
91/// # Errors
92///
93/// Returns an error containing `error_context` when any generated SQL statement fails.
94pub async fn execute_sqlite_statements(
95    manager: &SchemaManager<'_>,
96    statements: impl IntoIterator<Item = String>,
97    error_context: &str,
98) -> Result<(), DbErr> {
99    let db = manager.get_connection();
100    for sql in statements {
101        db.execute_unprepared(&sql)
102            .await
103            .map_err(|error| DbErr::Custom(format!("{error_context}: {error}")))?;
104    }
105    Ok(())
106}
107
108/// Creates a `PostgreSQL` extension when it is not already installed.
109///
110/// # Errors
111///
112/// Returns an error when the generated extension statement fails.
113pub async fn ensure_postgres_extension(
114    manager: &SchemaManager<'_>,
115    extension_name: &str,
116) -> Result<(), DbErr> {
117    let sql = Extension::create()
118        .name(extension_name)
119        .if_not_exists()
120        .to_string(PostgresQueryBuilder);
121    manager.get_connection().execute_unprepared(&sql).await?;
122    Ok(())
123}
124
125/// Builds a `PostgreSQL` GIN trigram index statement.
126#[must_use]
127pub fn postgres_trigram_index(
128    index_name: &str,
129    table_name: &str,
130    column_name: &str,
131) -> IndexCreateStatement {
132    Index::create()
133        .if_not_exists()
134        .name(index_name)
135        .table(Alias::new(table_name))
136        .full_text()
137        .col(
138            Alias::new(column_name)
139                .into_index_column()
140                .with_operator_class("gin_trgm_ops"),
141        )
142        .to_owned()
143}
144
145/// Builds a portable `DROP INDEX IF EXISTS` statement for `PostgreSQL` migrations.
146#[must_use]
147pub fn postgres_drop_index(index_name: &str) -> IndexDropStatement {
148    Index::drop().if_exists().name(index_name).to_owned()
149}
150
151/// Builds `MySQL`'s ngram-backed full-text index statement.
152///
153/// # Errors
154///
155/// Returns an error when the index, table, or column identifiers are invalid, or when `columns` is
156/// empty.
157pub fn mysql_fulltext_index_sql(
158    index_name: &str,
159    table_name: &str,
160    columns: &[&str],
161) -> Result<String, DbErr> {
162    validate_identifier(index_name)?;
163    validate_identifier(table_name)?;
164    validate_columns(columns)?;
165    Ok(format!(
166        "CREATE FULLTEXT INDEX {index_name} ON {table_name} ({}) WITH PARSER ngram",
167        columns.join(", ")
168    ))
169}
170
171fn validate_sqlite_fts_config(config: &SqliteFtsConfig<'_>) -> Result<(), DbErr> {
172    for identifier in [
173        config.virtual_table,
174        config.source_table,
175        config.insert_trigger,
176        config.delete_trigger,
177        config.update_trigger,
178    ] {
179        validate_identifier(identifier)?;
180    }
181    validate_columns(config.columns)
182}
183
184fn validate_columns(columns: &[&str]) -> Result<(), DbErr> {
185    if columns.is_empty() {
186        return Err(DbErr::Custom(
187            "migration search column list must not be empty".to_string(),
188        ));
189    }
190    for column in columns {
191        validate_identifier(column)?;
192    }
193    Ok(())
194}
195
196fn validate_identifier(identifier: &str) -> Result<(), DbErr> {
197    if !identifier.is_empty()
198        && identifier
199            .bytes()
200            .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
201    {
202        return Ok(());
203    }
204    Err(DbErr::Custom(format!(
205        "invalid migration search identifier: {identifier:?}"
206    )))
207}
208
209#[cfg(test)]
210mod tests {
211    use sea_orm_migration::sea_query::PostgresQueryBuilder;
212
213    use super::*;
214
215    fn config<'a>(columns: &'a [&'a str]) -> SqliteFtsConfig<'a> {
216        SqliteFtsConfig {
217            virtual_table: "files_name_fts",
218            source_table: "files",
219            columns,
220            insert_trigger: "files_name_fts_ai",
221            delete_trigger: "files_name_fts_ad",
222            update_trigger: "files_name_fts_au",
223        }
224    }
225
226    #[test]
227    fn sqlite_fts_statements_cover_create_backfill_triggers_and_drop() {
228        let columns = ["name", "description"];
229        let up = sqlite_fts_up_statements(&config(&columns)).unwrap();
230        assert_eq!(up.len(), 6);
231        assert!(up[0].contains("USING fts5(name, description, tokenize='trigram')"));
232        assert!(up[2].contains("SELECT id, name, description FROM files"));
233        assert!(up[3].contains("VALUES (new.id, new.name, new.description)"));
234        assert!(up[5].contains("name = new.name, description = new.description"));
235
236        let down = sqlite_fts_down_statements(&config(&columns)).unwrap();
237        assert_eq!(down.len(), 4);
238        assert_eq!(down[3], "DROP TABLE IF EXISTS files_name_fts");
239    }
240
241    #[test]
242    fn search_statement_builders_reject_empty_and_unsafe_identifiers() {
243        assert!(sqlite_fts_up_statements(&config(&[])).is_err());
244        assert!(mysql_fulltext_index_sql("idx-name", "files", &["name"]).is_err());
245        assert!(
246            mysql_fulltext_index_sql("idx_name", "files; DROP TABLE users", &["name"]).is_err()
247        );
248        assert!(mysql_fulltext_index_sql("idx_name", "files", &["name`, secret"]).is_err());
249    }
250
251    #[test]
252    fn backend_search_index_builders_render_expected_sql() {
253        let postgres = postgres_trigram_index("idx_files_name", "files", "name")
254            .to_string(PostgresQueryBuilder);
255        assert!(postgres.contains("USING GIN"));
256        assert!(postgres.contains("gin_trgm_ops"));
257
258        let drop = postgres_drop_index("idx_files_name").to_string(PostgresQueryBuilder);
259        assert_eq!(drop, "DROP INDEX IF EXISTS \"idx_files_name\"");
260
261        let mysql = mysql_fulltext_index_sql("idx_files_name", "files", &["name"]).unwrap();
262        assert_eq!(
263            mysql,
264            "CREATE FULLTEXT INDEX idx_files_name ON files (name) WITH PARSER ngram"
265        );
266    }
267}