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

sea_orm_migration/
manager.rs
