sea_orm_migration/
manager.rs

1use super::{IntoSchemaManagerConnection, SchemaManagerConnection};
2use sea_orm::sea_query::{
3    extension::postgres::{TypeAlterStatement, TypeCreateStatement, TypeDropStatement},
4    ForeignKeyCreateStatement, ForeignKeyDropStatement, IndexCreateStatement, IndexDropStatement,
5    SelectStatement, TableAlterStatement, TableCreateStatement, TableDropStatement,
6    TableRenameStatement, TableTruncateStatement,
7};
8use sea_orm::{ConnectionTrait, DbBackend, DbErr, StatementBuilder};
9#[allow(unused_imports)]
10use sea_schema::probe::SchemaProbe;
11
12/// Helper struct for writing migration scripts in migration file
13pub struct SchemaManager<'c> {
14    conn: SchemaManagerConnection<'c>,
15}
16
17impl<'c> SchemaManager<'c> {
18    pub fn new<T>(conn: T) -> Self
19    where
20        T: IntoSchemaManagerConnection<'c>,
21    {
22        Self {
23            conn: conn.into_schema_manager_connection(),
24        }
25    }
26
27    pub async fn exec_stmt<S>(&self, stmt: S) -> Result<(), DbErr>
28    where
29        S: StatementBuilder,
30    {
31        let builder = self.conn.get_database_backend();
32        self.conn.execute(builder.build(&stmt)).await.map(|_| ())
33    }
34
35    pub fn get_database_backend(&self) -> DbBackend {
36        self.conn.get_database_backend()
37    }
38
39    pub fn get_connection(&self) -> &SchemaManagerConnection<'c> {
40        &self.conn
41    }
42}
43
44/// Schema Creation
45impl SchemaManager<'_> {
46    pub async fn create_table(&self, stmt: TableCreateStatement) -> Result<(), DbErr> {
47        self.exec_stmt(stmt).await
48    }
49
50    pub async fn create_index(&self, stmt: IndexCreateStatement) -> Result<(), DbErr> {
51        self.exec_stmt(stmt).await
52    }
53
54    pub async fn create_foreign_key(&self, stmt: ForeignKeyCreateStatement) -> Result<(), DbErr> {
55        self.exec_stmt(stmt).await
56    }
57
58    pub async fn create_type(&self, stmt: TypeCreateStatement) -> Result<(), DbErr> {
59        self.exec_stmt(stmt).await
60    }
61}
62
63/// Schema Mutation
64impl SchemaManager<'_> {
65    pub async fn alter_table(&self, stmt: TableAlterStatement) -> Result<(), DbErr> {
66        self.exec_stmt(stmt).await
67    }
68
69    pub async fn drop_table(&self, stmt: TableDropStatement) -> Result<(), DbErr> {
70        self.exec_stmt(stmt).await
71    }
72
73    pub async fn rename_table(&self, stmt: TableRenameStatement) -> Result<(), DbErr> {
74        self.exec_stmt(stmt).await
75    }
76
77    pub async fn truncate_table(&self, stmt: TableTruncateStatement) -> Result<(), DbErr> {
78        self.exec_stmt(stmt).await
79    }
80
81    pub async fn drop_index(&self, stmt: IndexDropStatement) -> Result<(), DbErr> {
82        self.exec_stmt(stmt).await
83    }
84
85    pub async fn drop_foreign_key(&self, stmt: ForeignKeyDropStatement) -> Result<(), DbErr> {
86        self.exec_stmt(stmt).await
87    }
88
89    pub async fn alter_type(&self, stmt: TypeAlterStatement) -> Result<(), DbErr> {
90        self.exec_stmt(stmt).await
91    }
92
93    pub async fn drop_type(&self, stmt: TypeDropStatement) -> Result<(), DbErr> {
94        self.exec_stmt(stmt).await
95    }
96}
97
98/// Schema Inspection.
99impl SchemaManager<'_> {
100    pub async fn has_table<T>(&self, table: T) -> Result<bool, DbErr>
101    where
102        T: AsRef<str>,
103    {
104        has_table(&self.conn, table).await
105    }
106
107    pub async fn has_column<T, C>(&self, _table: T, _column: C) -> Result<bool, DbErr>
108    where
109        T: AsRef<str>,
110        C: AsRef<str>,
111    {
112        let _stmt: SelectStatement = match self.conn.get_database_backend() {
113            #[cfg(feature = "sqlx-mysql")]
114            DbBackend::MySql => sea_schema::mysql::MySql.has_column(_table, _column),
115            #[cfg(feature = "sqlx-postgres")]
116            DbBackend::Postgres => sea_schema::postgres::Postgres.has_column(_table, _column),
117            #[cfg(feature = "sqlx-sqlite")]
118            DbBackend::Sqlite => sea_schema::sqlite::Sqlite.has_column(_table, _column),
119            #[allow(unreachable_patterns)]
120            other => panic!("{other:?} feature is off"),
121        };
122
123        #[allow(unreachable_code)]
124        let builder = self.conn.get_database_backend();
125        let res = self
126            .conn
127            .query_one(builder.build(&_stmt))
128            .await?
129            .ok_or_else(|| DbErr::Custom("Failed to check column exists".to_owned()))?;
130
131        res.try_get("", "has_column")
132    }
133
134    pub async fn has_index<T, I>(&self, _table: T, _index: I) -> Result<bool, DbErr>
135    where
136        T: AsRef<str>,
137        I: AsRef<str>,
138    {
139        let _stmt: SelectStatement = match self.conn.get_database_backend() {
140            #[cfg(feature = "sqlx-mysql")]
141            DbBackend::MySql => sea_schema::mysql::MySql.has_index(_table, _index),
142            #[cfg(feature = "sqlx-postgres")]
143            DbBackend::Postgres => sea_schema::postgres::Postgres.has_index(_table, _index),
144            #[cfg(feature = "sqlx-sqlite")]
145            DbBackend::Sqlite => sea_schema::sqlite::Sqlite.has_index(_table, _index),
146            #[allow(unreachable_patterns)]
147            other => panic!("{other:?} feature is off"),
148        };
149
150        #[allow(unreachable_code)]
151        let builder = self.conn.get_database_backend();
152        let res = self
153            .conn
154            .query_one(builder.build(&_stmt))
155            .await?
156            .ok_or_else(|| DbErr::Custom("Failed to check index exists".to_owned()))?;
157
158        res.try_get("", "has_index")
159    }
160}
161
162pub(crate) async fn has_table<C, T>(conn: &C, _table: T) -> Result<bool, DbErr>
163where
164    C: ConnectionTrait,
165    T: AsRef<str>,
166{
167    let _stmt: SelectStatement = match conn.get_database_backend() {
168        #[cfg(feature = "sqlx-mysql")]
169        DbBackend::MySql => sea_schema::mysql::MySql.has_table(_table),
170        #[cfg(feature = "sqlx-postgres")]
171        DbBackend::Postgres => sea_schema::postgres::Postgres.has_table(_table),
172        #[cfg(feature = "sqlx-sqlite")]
173        DbBackend::Sqlite => sea_schema::sqlite::Sqlite.has_table(_table),
174        #[allow(unreachable_patterns)]
175        other => panic!("{other:?} feature is off"),
176    };
177
178    #[allow(unreachable_code)]
179    let builder = conn.get_database_backend();
180    let res = conn
181        .query_one(builder.build(&_stmt))
182        .await?
183        .ok_or_else(|| DbErr::Custom("Failed to check table exists".to_owned()))?;
184
185    res.try_get("", "has_table")
186}
sea_orm_migration/manager.rs

sea_orm_migration/
manager.rs
