1use sea_orm_migration::prelude::*;
2
3pub struct Schema;
4
5impl Schema {
6 pub async fn create<F>(manager: &SchemaManager<'_>, table_name: &str, callback: F) -> Result<(), DbErr>
7 where
8 F: FnOnce(&mut Blueprint),
9 {
10 let mut blueprint = Blueprint::new(table_name);
11 blueprint.is_alter = false;
12 callback(&mut blueprint);
13
14 let mut table_stmt = blueprint.table;
15
16 if blueprint.auto_id {
18 table_stmt.col(
19 ColumnDef::new(Alias::new("id"))
20 .integer()
21 .not_null()
22 .auto_increment()
23 .primary_key()
24 );
25 }
26
27 for mut col in blueprint.columns {
29 table_stmt.col(&mut col);
30 }
31
32 if blueprint.timestamps {
34 table_stmt.col(
35 ColumnDef::new(Alias::new("created_at"))
36 .date_time()
37 .not_null()
38 .default(Expr::current_timestamp())
39 );
40 table_stmt.col(
41 ColumnDef::new(Alias::new("updated_at"))
42 .date_time()
43 .not_null()
44 .default(Expr::current_timestamp())
45 );
46 }
47
48 manager.create_table(table_stmt).await?;
50
51 for index in blueprint.indices {
53 manager.create_index(index).await?;
54 }
55
56 Ok(())
57 }
58
59 pub async fn table<F>(manager: &SchemaManager<'_>, table_name: &str, callback: F) -> Result<(), DbErr>
60 where
61 F: FnOnce(&mut Blueprint),
62 {
63 let mut blueprint = Blueprint::new(table_name);
64 blueprint.is_alter = true;
65 blueprint.auto_id = false; blueprint.timestamps = false; callback(&mut blueprint);
68
69 let mut alter_stmt = blueprint.alter;
70 for mut col in blueprint.columns {
71 alter_stmt.add_column(&mut col);
72 }
73
74 manager.alter_table(alter_stmt).await?;
75
76 for index in blueprint.indices {
77 manager.create_index(index).await?;
78 }
79
80 Ok(())
81 }
82
83 pub async fn drop(manager: &SchemaManager<'_>, table_name: &str) -> Result<(), DbErr> {
84 manager.drop_table(
85 Table::drop()
86 .table(Alias::new(table_name))
87 .to_owned()
88 ).await
89 }
90}
91
92pub struct Blueprint {
93 pub table_name: String,
94 pub table: TableCreateStatement,
95 pub alter: TableAlterStatement,
96 pub is_alter: bool,
97 pub indices: Vec<IndexCreateStatement>,
98 pub columns: Vec<ColumnDef>,
99 pub auto_id: bool,
100 pub timestamps: bool,
101}
102
103impl Blueprint {
104 pub fn new(table_name: &str) -> Self {
105 let mut table = Table::create();
106 table.table(Alias::new(table_name)).if_not_exists();
107
108 let mut alter = Table::alter();
109 alter.table(Alias::new(table_name));
110
111 Self {
112 table_name: table_name.to_string(),
113 table,
114 alter,
115 is_alter: false,
116 indices: Vec::new(),
117 columns: Vec::new(),
118 auto_id: true,
119 timestamps: true,
120 }
121 }
122
123 pub fn no_id(&mut self) -> &mut Self {
124 self.auto_id = false;
125 self
126 }
127
128 pub fn no_timestamps(&mut self) -> &mut Self {
129 self.timestamps = false;
130 self
131 }
132
133 pub fn id(&mut self) -> &mut Self {
134 self.auto_id = true;
135 self
136 }
137
138 pub fn string<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
139 let mut col_def = ColumnDef::new(Alias::new(name));
140 col_def.string();
141 ColumnBuilder::new(self, name, col_def)
142 }
143
144 pub fn text<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
145 let mut col_def = ColumnDef::new(Alias::new(name));
146 col_def.text();
147 ColumnBuilder::new(self, name, col_def)
148 }
149
150 pub fn integer<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
151 let mut col_def = ColumnDef::new(Alias::new(name));
152 col_def.integer();
153 ColumnBuilder::new(self, name, col_def)
154 }
155
156 pub fn big_integer<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
157 let mut col_def = ColumnDef::new(Alias::new(name));
158 col_def.big_integer();
159 ColumnBuilder::new(self, name, col_def)
160 }
161
162 pub fn float<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
163 let mut col_def = ColumnDef::new(Alias::new(name));
164 col_def.float();
165 ColumnBuilder::new(self, name, col_def)
166 }
167
168 pub fn double<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
169 let mut col_def = ColumnDef::new(Alias::new(name));
170 col_def.double();
171 ColumnBuilder::new(self, name, col_def)
172 }
173
174 pub fn decimal<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
175 let mut col_def = ColumnDef::new(Alias::new(name));
176 col_def.decimal();
177 ColumnBuilder::new(self, name, col_def)
178 }
179
180 pub fn char<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
181 let mut col_def = ColumnDef::new(Alias::new(name));
182 col_def.char();
183 ColumnBuilder::new(self, name, col_def)
184 }
185
186 pub fn boolean<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
187 let mut col_def = ColumnDef::new(Alias::new(name));
188 col_def.boolean();
189 ColumnBuilder::new(self, name, col_def)
190 }
191
192 pub fn date_time<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
193 let mut col_def = ColumnDef::new(Alias::new(name));
194 col_def.date_time();
195 ColumnBuilder::new(self, name, col_def)
196 }
197
198 pub fn timestamp<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
199 let mut col_def = ColumnDef::new(Alias::new(name));
200 col_def.timestamp();
201 ColumnBuilder::new(self, name, col_def)
202 }
203
204 pub fn uuid<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
205 let mut col_def = ColumnDef::new(Alias::new(name));
206 col_def.uuid();
207 ColumnBuilder::new(self, name, col_def)
208 }
209
210 pub fn json<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
211 let mut col_def = ColumnDef::new(Alias::new(name));
212 col_def.json();
213 ColumnBuilder::new(self, name, col_def)
214 }
215
216 pub fn json_binary<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
217 let mut col_def = ColumnDef::new(Alias::new(name));
218 col_def.json_binary();
219 ColumnBuilder::new(self, name, col_def)
220 }
221
222 pub fn binary<'a>(&'a mut self, name: &str) -> ColumnBuilder<'a> {
223 let mut col_def = ColumnDef::new(Alias::new(name));
224 col_def.binary();
225 ColumnBuilder::new(self, name, col_def)
226 }
227
228 pub fn timestamps(&mut self) -> &mut Self {
229 self.timestamps = true;
230 self
231 }
232
233 pub fn drop_column(&mut self, name: &str) -> &mut Self {
234 self.alter.drop_column(Alias::new(name));
235 self
236 }
237
238 pub fn foreign<'a>(&'a mut self, from_col: &str) -> ForeignKeyBuilder<'a> {
239 ForeignKeyBuilder::new(self, from_col)
240 }
241}
242
243pub struct ColumnBuilder<'a> {
244 blueprint: &'a mut Blueprint,
245 col_name: String,
246 col_def: Option<ColumnDef>,
247 is_indexed: bool,
248}
249
250impl<'a> ColumnBuilder<'a> {
251 pub fn new(blueprint: &'a mut Blueprint, col_name: &str, col_def: ColumnDef) -> Self {
252 Self {
253 blueprint,
254 col_name: col_name.to_string(),
255 col_def: Some(col_def),
256 is_indexed: false,
257 }
258 }
259
260 pub fn not_null(mut self) -> Self {
261 if let Some(ref mut col) = self.col_def {
262 col.not_null();
263 }
264 self
265 }
266
267 pub fn nullable(mut self) -> Self {
268 if let Some(ref mut col) = self.col_def {
269 col.null();
270 }
271 self
272 }
273
274 pub fn unique(mut self) -> Self {
275 if let Some(ref mut col) = self.col_def {
276 col.unique_key();
277 }
278 self
279 }
280
281 pub fn primary_key(mut self) -> Self {
282 if let Some(ref mut col) = self.col_def {
283 col.primary_key();
284 }
285 self
286 }
287
288 pub fn default<T>(mut self, value: T) -> Self
289 where
290 T: Into<SimpleExpr>,
291 {
292 if let Some(ref mut col) = self.col_def {
293 col.default(value);
294 }
295 self
296 }
297
298 pub fn index(mut self) -> Self {
299 self.is_indexed = true;
300 self
301 }
302}
303
304impl<'a> Drop for ColumnBuilder<'a> {
305 fn drop(&mut self) {
306 if let Some(col) = self.col_def.take() {
307 self.blueprint.columns.push(col);
308 if self.is_indexed {
309 let index_name = format!("{}_{}_index", self.blueprint.table_name, self.col_name);
310 let index_stmt = Index::create()
311 .name(&index_name)
312 .table(Alias::new(&self.blueprint.table_name))
313 .col(Alias::new(&self.col_name))
314 .to_owned();
315 self.blueprint.indices.push(index_stmt);
316 }
317 }
318 }
319}
320
321pub struct ForeignKeyBuilder<'a> {
322 blueprint: &'a mut Blueprint,
323 from_col: String,
324 to_col: Option<String>,
325 to_table: Option<String>,
326 on_delete: Option<ForeignKeyAction>,
327 on_update: Option<ForeignKeyAction>,
328}
329
330impl<'a> ForeignKeyBuilder<'a> {
331 pub fn new(blueprint: &'a mut Blueprint, from_col: &str) -> Self {
332 Self {
333 blueprint,
334 from_col: from_col.to_string(),
335 to_col: None,
336 to_table: None,
337 on_delete: None,
338 on_update: None,
339 }
340 }
341
342 pub fn references(mut self, to_col: &str) -> Self {
343 self.to_col = Some(to_col.to_string());
344 self
345 }
346
347 pub fn on(mut self, to_table: &str) -> Self {
348 self.to_table = Some(to_table.to_string());
349 self
350 }
351
352 pub fn on_delete(mut self, action: &str) -> Self {
353 let act = match action.to_lowercase().as_str() {
354 "cascade" => ForeignKeyAction::Cascade,
355 "restrict" => ForeignKeyAction::Restrict,
356 "set null" | "set_null" => ForeignKeyAction::SetNull,
357 "no action" | "no_action" => ForeignKeyAction::NoAction,
358 _ => ForeignKeyAction::Cascade,
359 };
360 self.on_delete = Some(act);
361 self
362 }
363
364 pub fn on_update(mut self, action: &str) -> Self {
365 let act = match action.to_lowercase().as_str() {
366 "cascade" => ForeignKeyAction::Cascade,
367 "restrict" => ForeignKeyAction::Restrict,
368 "set null" | "set_null" => ForeignKeyAction::SetNull,
369 "no action" | "no_action" => ForeignKeyAction::NoAction,
370 _ => ForeignKeyAction::Cascade,
371 };
372 self.on_update = Some(act);
373 self
374 }
375}
376
377impl<'a> Drop for ForeignKeyBuilder<'a> {
378 fn drop(&mut self) {
379 if let (Some(to_table), Some(to_col)) = (&self.to_table, &self.to_col) {
380 let mut fk = ForeignKey::create();
381
382 let fk_name = format!(
383 "fk_{}_{}_{}_{}",
384 self.blueprint.table_name, self.from_col, to_table, to_col
385 );
386
387 fk.name(&fk_name)
388 .from(Alias::new(&self.blueprint.table_name), Alias::new(&self.from_col))
389 .to(Alias::new(to_table), Alias::new(to_col));
390
391 if let Some(action) = self.on_delete {
392 fk.on_delete(action);
393 }
394 if let Some(action) = self.on_update {
395 fk.on_update(action);
396 }
397
398 self.blueprint.table.foreign_key(&mut fk);
399 }
400 }
401}