Skip to main content

aster_forge_db/
index.rs

1use sea_orm::sea_query::{Alias, Index, IndexDropStatement};
2use sea_orm::{ConnectionTrait, DatabaseBackend, DbErr, Statement};
3
4/// Drops an index when it exists across all supported database backends.
5///
6/// MySQL does not support `DROP INDEX IF EXISTS`, so this helper checks
7/// `information_schema.statistics` before issuing the backend-specific drop.
8pub async fn drop_index_if_exists<C>(
9    db: &C,
10    table_name: &str,
11    index_name: &str,
12) -> Result<(), DbErr>
13where
14    C: ConnectionTrait,
15{
16    let backend = db.get_database_backend();
17    if backend == DatabaseBackend::MySql && !mysql_index_exists(db, table_name, index_name).await? {
18        return Ok(());
19    }
20
21    db.execute(&drop_index_for_backend(backend, table_name, index_name))
22        .await?;
23    Ok(())
24}
25
26/// Renames a MySQL index when the source exists and the target does not.
27///
28/// This is useful when a table or business term is renamed while preserving
29/// the physical index definition. Calling it repeatedly is safe.
30pub async fn rename_mysql_index_if_exists<C>(
31    db: &C,
32    table_name: &str,
33    old_index_name: &str,
34    new_index_name: &str,
35) -> Result<(), DbErr>
36where
37    C: ConnectionTrait,
38{
39    validate_mysql_identifier(table_name)?;
40    validate_mysql_identifier(old_index_name)?;
41    validate_mysql_identifier(new_index_name)?;
42
43    if db.get_database_backend() != DatabaseBackend::MySql {
44        return Err(DbErr::Custom(
45            "rename_mysql_index_if_exists requires a MySQL connection".to_string(),
46        ));
47    }
48
49    if !mysql_index_exists(db, table_name, old_index_name).await?
50        || mysql_index_exists(db, table_name, new_index_name).await?
51    {
52        return Ok(());
53    }
54
55    db.execute_unprepared(&format!(
56        "ALTER TABLE `{table_name}` RENAME INDEX `{old_index_name}` TO `{new_index_name}`"
57    ))
58    .await?;
59    Ok(())
60}
61
62fn drop_index_for_backend(
63    backend: DatabaseBackend,
64    table_name: &str,
65    index_name: &str,
66) -> IndexDropStatement {
67    let mut statement = Index::drop();
68    statement
69        .name(index_name.to_owned())
70        .table(Alias::new(table_name));
71    if backend != DatabaseBackend::MySql {
72        statement.if_exists();
73    }
74    statement.to_owned()
75}
76
77async fn mysql_index_exists<C>(db: &C, table_name: &str, index_name: &str) -> Result<bool, DbErr>
78where
79    C: ConnectionTrait,
80{
81    let row = db
82        .query_one_raw(Statement::from_sql_and_values(
83            DatabaseBackend::MySql,
84            "SELECT 1 FROM information_schema.statistics \
85             WHERE table_schema = DATABASE() AND table_name = ? AND index_name = ? LIMIT 1",
86            [table_name.into(), index_name.into()],
87        ))
88        .await?;
89    Ok(row.is_some())
90}
91
92fn validate_mysql_identifier(identifier: &str) -> Result<(), DbErr> {
93    if !identifier.is_empty()
94        && identifier
95            .bytes()
96            .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
97    {
98        return Ok(());
99    }
100
101    Err(DbErr::Custom(format!(
102        "invalid MySQL migration identifier: {identifier:?}"
103    )))
104}
105
106#[cfg(test)]
107mod tests {
108    use sea_orm::sea_query::{MysqlQueryBuilder, PostgresQueryBuilder, SqliteQueryBuilder};
109    use sea_orm::{ConnectionTrait, Database, DatabaseBackend};
110
111    use super::{
112        drop_index_for_backend, drop_index_if_exists, rename_mysql_index_if_exists,
113        validate_mysql_identifier,
114    };
115
116    #[test]
117    fn drop_index_sql_respects_backend_capabilities() {
118        let mysql = drop_index_for_backend(DatabaseBackend::MySql, "example_table", "idx_example")
119            .to_string(MysqlQueryBuilder);
120        assert_eq!(mysql, "DROP INDEX `idx_example` ON `example_table`");
121
122        let postgres =
123            drop_index_for_backend(DatabaseBackend::Postgres, "example_table", "idx_example")
124                .to_string(PostgresQueryBuilder);
125        assert_eq!(postgres, "DROP INDEX IF EXISTS \"idx_example\"");
126
127        let sqlite =
128            drop_index_for_backend(DatabaseBackend::Sqlite, "example_table", "idx_example")
129                .to_string(SqliteQueryBuilder);
130        assert_eq!(sqlite, "DROP INDEX IF EXISTS \"idx_example\"");
131    }
132
133    #[test]
134    fn mysql_identifier_validation_rejects_raw_sql_fragments() {
135        assert!(validate_mysql_identifier("idx_example_2026").is_ok());
136        assert!(validate_mysql_identifier("").is_err());
137        assert!(validate_mysql_identifier("idx-example").is_err());
138        assert!(validate_mysql_identifier("idx` DROP TABLE users").is_err());
139    }
140
141    #[tokio::test]
142    async fn drop_index_if_exists_is_idempotent_on_sqlite() {
143        let db = Database::connect("sqlite::memory:")
144            .await
145            .expect("SQLite migration helper test database should connect");
146        db.execute_unprepared("CREATE TABLE example_table (id INTEGER PRIMARY KEY)")
147            .await
148            .expect("example table should be created");
149        db.execute_unprepared("CREATE INDEX idx_example ON example_table (id)")
150            .await
151            .expect("example index should be created");
152
153        drop_index_if_exists(&db, "example_table", "idx_example")
154            .await
155            .expect("existing index should be dropped");
156        drop_index_if_exists(&db, "example_table", "idx_example")
157            .await
158            .expect("missing index should be ignored");
159    }
160
161    #[tokio::test]
162    async fn mysql_index_rename_rejects_non_mysql_connections() {
163        let db = Database::connect("sqlite::memory:")
164            .await
165            .expect("SQLite migration helper test database should connect");
166        let error = rename_mysql_index_if_exists(
167            &db,
168            "example_table",
169            "idx_example_old",
170            "idx_example_new",
171        )
172        .await
173        .expect_err("MySQL-only index rename should reject SQLite");
174
175        assert!(error.to_string().contains("requires a MySQL connection"));
176    }
177}