aster_forge_db_migration/
search.rs1use 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
7pub 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
17pub 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
74pub 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
89pub 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
108pub 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#[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#[must_use]
147pub fn postgres_drop_index(index_name: &str) -> IndexDropStatement {
148 Index::drop().if_exists().name(index_name).to_owned()
149}
150
151pub 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}