1use super::{Schema, TopologicalSort, entity::index_table_ref};
2use crate::{ConnectionTrait, DbBackend, DbErr, EntityTrait, Statement};
3use sea_query::{
4 ForeignKeyCreateStatement, Index, IndexCreateStatement, IntoIden, TableAlterStatement,
5 TableCreateStatement, TableName, TableRef, extension::postgres::TypeCreateStatement,
6};
7
8pub struct SchemaBuilder {
10 helper: Schema,
11 entities: Vec<EntitySchemaInfo>,
12}
13
14pub struct EntitySchemaInfo {
16 table: TableCreateStatement,
17 enums: Vec<TypeCreateStatement>,
18 indexes: Vec<IndexCreateStatement>,
19 schema_name: Option<String>,
22}
23
24impl std::fmt::Debug for SchemaBuilder {
25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26 write!(f, "SchemaBuilder {{")?;
27 write!(f, " entities: [")?;
28 for (i, entity) in self.entities.iter().enumerate() {
29 if i > 0 {
30 write!(f, ", ")?;
31 }
32 entity.debug_print(f, &self.helper.backend)?;
33 }
34 write!(f, " ]")?;
35 write!(f, " }}")
36 }
37}
38
39impl std::fmt::Debug for EntitySchemaInfo {
40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 self.debug_print(f, &DbBackend::Sqlite)
42 }
43}
44
45impl SchemaBuilder {
46 pub fn new(schema: Schema) -> Self {
48 Self {
49 helper: schema,
50 entities: Default::default(),
51 }
52 }
53
54 pub fn register<E: EntityTrait>(mut self, entity: E) -> Self {
56 let entity = EntitySchemaInfo::new(entity, &self.helper);
57 if !self
58 .entities
59 .iter()
60 .any(|e| e.table.get_table_name() == entity.table.get_table_name())
61 {
62 self.entities.push(entity);
63 }
64 self
65 }
66
67 #[cfg(feature = "entity-registry")]
68 pub(crate) fn helper(&self) -> &Schema {
69 &self.helper
70 }
71
72 #[cfg(feature = "entity-registry")]
73 pub(crate) fn register_entity(&mut self, entity: EntitySchemaInfo) {
74 self.entities.push(entity);
75 }
76
77 #[cfg(feature = "schema-sync")]
89 #[cfg_attr(docsrs, doc(cfg(feature = "schema-sync")))]
90 pub async fn sync<C>(self, db: &C) -> Result<(), DbErr>
91 where
92 C: ConnectionTrait + sea_schema::Connection,
93 {
94 let _existing =
95 match db.get_database_backend() {
96 #[cfg(feature = "sqlx-mysql")]
97 DbBackend::MySql => {
98 use sea_schema::{mysql::discovery::SchemaDiscovery, probe::SchemaProbe};
99
100 let current_schema: String = db
101 .query_one(
102 sea_query::SelectStatement::new()
103 .expr(sea_schema::mysql::MySql::get_current_schema()),
104 )
105 .await?
106 .ok_or_else(|| DbErr::RecordNotFound("Can't get current schema".into()))?
107 .try_get_by_index(0)?;
108
109 let mut target_schemas = std::collections::BTreeSet::new();
111 for entity in &self.entities {
112 let schema = entity.schema_name.as_deref().unwrap_or(¤t_schema);
113 target_schemas.insert(schema.to_string());
114 }
115
116 let mut tables_by_schema = std::collections::HashMap::new();
117 for schema_name in &target_schemas {
118 let schema_discovery = SchemaDiscovery::new_no_exec(schema_name);
119 let schema = schema_discovery.discover_with(db).await.map_err(|err| {
120 DbErr::Query(crate::RuntimeErr::SqlxError(err.into()))
121 })?;
122
123 tables_by_schema.insert(
124 schema_name.clone(),
125 schema.tables.iter().map(|table| table.write()).collect(),
126 );
127 }
128
129 DiscoveredSchema {
130 current_schema,
131 tables_by_schema,
132 enums_by_schema: Default::default(),
133 }
134 }
135 #[cfg(feature = "sqlx-postgres")]
136 DbBackend::Postgres => {
137 use sea_schema::{postgres::discovery::SchemaDiscovery, probe::SchemaProbe};
138
139 let current_schema: String = db
140 .query_one(
141 sea_query::SelectStatement::new()
142 .expr(sea_schema::postgres::Postgres::get_current_schema()),
143 )
144 .await?
145 .ok_or_else(|| DbErr::RecordNotFound("Can't get current schema".into()))?
146 .try_get_by_index(0)?;
147
148 let mut target_schemas = std::collections::BTreeSet::new();
150 for entity in &self.entities {
151 let schema = entity.schema_name.as_deref().unwrap_or(¤t_schema);
152 target_schemas.insert(schema.to_string());
153 }
154
155 let mut tables_by_schema = std::collections::HashMap::new();
156 let mut enums_by_schema = std::collections::HashMap::new();
157 for schema_name in &target_schemas {
158 let schema_discovery = SchemaDiscovery::new_no_exec(schema_name);
159 let schema = schema_discovery.discover_with(db).await.map_err(|err| {
160 DbErr::Query(crate::RuntimeErr::SqlxError(err.into()))
161 })?;
162
163 tables_by_schema.insert(
164 schema_name.clone(),
165 schema.tables.iter().map(|table| table.write()).collect(),
166 );
167 enums_by_schema.insert(
168 schema_name.clone(),
169 schema.enums.iter().map(|def| def.write()).collect(),
170 );
171 }
172
173 DiscoveredSchema {
174 current_schema,
175 tables_by_schema,
176 enums_by_schema,
177 }
178 }
179 #[cfg(feature = "sqlx-sqlite")]
180 DbBackend::Sqlite => {
181 use sea_schema::sqlite::{SqliteDiscoveryError, discovery::SchemaDiscovery};
182 let schema = SchemaDiscovery::discover_with(db)
183 .await
184 .map_err(|err| {
185 DbErr::Query(match err {
186 SqliteDiscoveryError::SqlxError(err) => {
187 crate::RuntimeErr::SqlxError(err.into())
188 }
189 _ => crate::RuntimeErr::Internal(format!("{err:?}")),
190 })
191 })?
192 .merge_indexes_into_table();
193 let mut tables_by_schema = std::collections::HashMap::new();
194 tables_by_schema.insert(
195 String::new(),
196 schema.tables.iter().map(|table| table.write()).collect(),
197 );
198 DiscoveredSchema {
199 current_schema: String::new(),
200 tables_by_schema,
201 enums_by_schema: Default::default(),
202 }
203 }
204 #[cfg(feature = "rusqlite")]
205 DbBackend::Sqlite => {
206 use sea_schema::sqlite::{SqliteDiscoveryError, discovery::SchemaDiscovery};
207 let schema = SchemaDiscovery::discover_with(db)
208 .map_err(|err| {
209 DbErr::Query(match err {
210 SqliteDiscoveryError::RusqliteError(err) => {
211 crate::RuntimeErr::Rusqlite(err.into())
212 }
213 _ => crate::RuntimeErr::Internal(format!("{err:?}")),
214 })
215 })?
216 .merge_indexes_into_table();
217 let mut tables_by_schema = std::collections::HashMap::new();
218 tables_by_schema.insert(
219 String::new(),
220 schema.tables.iter().map(|table| table.write()).collect(),
221 );
222 DiscoveredSchema {
223 current_schema: String::new(),
224 tables_by_schema,
225 enums_by_schema: Default::default(),
226 }
227 }
228 #[allow(unreachable_patterns)]
229 other => {
230 return Err(DbErr::BackendNotSupported {
231 db: other.as_str(),
232 ctx: "SchemaBuilder::sync",
233 });
234 }
235 };
236
237 #[allow(unreachable_code)]
238 let mut created_enums: Vec<Statement> = Default::default();
239
240 #[allow(unreachable_code)]
241 for table_name in self.sorted_tables() {
242 if let Some(entity) = self
243 .entities
244 .iter()
245 .find(|entity| table_name == get_table_name(entity.table.get_table_name()))
246 {
247 entity.sync(db, &_existing, &mut created_enums).await?;
248 }
249 }
250
251 Ok(())
252 }
253
254 pub async fn apply<C: ConnectionTrait>(self, db: &C) -> Result<(), DbErr> {
258 let mut created_enums: Vec<Statement> = Default::default();
259
260 for table_name in self.sorted_tables() {
261 if let Some(entity) = self
262 .entities
263 .iter()
264 .find(|entity| table_name == get_table_name(entity.table.get_table_name()))
265 {
266 entity.apply(db, &mut created_enums).await?;
267 }
268 }
269
270 Ok(())
271 }
272
273 #[allow(dead_code)]
278 #[cfg(all(feature = "schema-sync", feature = "sqlx-sqlite"))]
279 fn _assert_sync_future_is_send(self, db: &crate::DatabaseConnection) {
280 fn assert_send<T: Send>(_: &T) {}
281 assert_send(&self.sync(db));
282 }
283
284 fn sorted_tables(&self) -> Vec<TableName> {
285 let mut sorter = TopologicalSort::<TableName>::new();
286
287 for entity in self.entities.iter() {
288 let table_name = get_table_name(entity.table.get_table_name());
289 sorter.insert(table_name);
290 }
291 for entity in self.entities.iter() {
292 let self_table = get_table_name(entity.table.get_table_name());
293 for fk in entity.table.get_foreign_key_create_stmts().iter() {
294 let fk = fk.get_foreign_key();
295 let ref_table = get_table_name(fk.get_ref_table());
296 if self_table != ref_table {
297 sorter.add_dependency(ref_table, self_table.clone());
299 }
300 }
301 }
302 let mut sorted = Vec::new();
303 while let Some(i) = sorter.pop() {
304 sorted.push(i);
305 }
306 if sorted.len() != self.entities.len() {
307 for entity in self.entities.iter() {
309 let table_name = get_table_name(entity.table.get_table_name());
310 if !sorted.contains(&table_name) {
311 sorted.push(table_name);
312 }
313 }
314 }
315
316 sorted
317 }
318}
319
320struct DiscoveredSchema {
321 current_schema: String,
323 tables_by_schema: std::collections::HashMap<String, Vec<TableCreateStatement>>,
325 enums_by_schema: std::collections::HashMap<String, Vec<TypeCreateStatement>>,
327}
328
329impl DiscoveredSchema {
330 fn find_table(
339 &self,
340 entity_schema: Option<&str>,
341 entity_table_name: &TableName,
342 ) -> Option<&TableCreateStatement> {
343 let schema = entity_schema.unwrap_or(&self.current_schema);
344 let schema_tables = self.tables_by_schema.get(schema)?;
345 let bare_entity_name = TableName(None, entity_table_name.1.clone());
348 schema_tables
349 .iter()
350 .find(|tbl| get_table_name(tbl.get_table_name()) == bare_entity_name)
351 }
352
353 fn find_enums(&self, entity_schema: Option<&str>) -> &[TypeCreateStatement] {
354 let schema = entity_schema.unwrap_or(&self.current_schema);
355 self.enums_by_schema
356 .get(schema)
357 .map(|v| v.as_slice())
358 .unwrap_or(&[])
359 }
360}
361
362impl EntitySchemaInfo {
363 pub fn new<E: EntityTrait>(entity: E, helper: &Schema) -> Self {
365 Self {
366 table: helper.create_table_from_entity(entity),
367 enums: helper.create_enum_from_entity(entity),
368 indexes: helper.create_index_from_entity(entity),
369 schema_name: entity.schema_name().map(|s| s.to_string()),
370 }
371 }
372
373 async fn apply<C: ConnectionTrait>(
374 &self,
375 db: &C,
376 created_enums: &mut Vec<Statement>,
377 ) -> Result<(), DbErr> {
378 for stmt in self.enums.iter() {
379 let new_stmt = db.get_database_backend().build(stmt);
380 if !created_enums.iter().any(|s| s == &new_stmt) {
381 db.execute(stmt).await?;
382 created_enums.push(new_stmt);
383 }
384 }
385 db.execute(&self.table).await?;
386 for stmt in self.indexes.iter() {
387 db.execute(stmt).await?;
388 }
389 Ok(())
390 }
391
392 #[allow(dead_code)]
394 async fn sync<C: ConnectionTrait>(
395 &self,
396 db: &C,
397 existing: &DiscoveredSchema,
398 created_enums: &mut Vec<Statement>,
399 ) -> Result<(), DbErr> {
400 let db_backend = db.get_database_backend();
401
402 let existing_enums = existing.find_enums(self.schema_name.as_deref());
404 for stmt in self.enums.iter() {
405 let mut has_enum = false;
406 let new_stmt = db_backend.build(stmt);
407 for existing_enum in existing_enums {
408 if db_backend.build(existing_enum) == new_stmt {
409 has_enum = true;
410 break;
412 }
413 }
414 if !has_enum && !created_enums.iter().any(|s| s == &new_stmt) {
415 db.execute(stmt).await?;
416 created_enums.push(new_stmt);
417 }
418 }
419 let table_name = get_table_name(self.table.get_table_name());
420 let existing_table = existing.find_table(self.schema_name.as_deref(), &table_name);
422 if let Some(existing_table) = existing_table {
423 for column_def in self.table.get_columns() {
424 let existing_column = existing_table
425 .get_columns()
426 .iter()
427 .find(|c| c.get_column_name() == column_def.get_column_name());
428 let Some(existing_column) = existing_column else {
429 let mut renamed_from = "";
430 if let Some(comment) = &column_def.get_column_spec().comment
431 && let Some((_, suffix)) = comment.rsplit_once("renamed_from \"")
432 && let Some((prefix, _)) = suffix.split_once('"')
433 {
434 renamed_from = prefix;
435 }
436 if renamed_from.is_empty() {
437 db.execute(
438 TableAlterStatement::new()
439 .table(self.table.get_table_name().expect("Checked above").clone())
440 .add_column(column_def.to_owned()),
441 )
442 .await?;
443 } else {
444 db.execute(
445 TableAlterStatement::new()
446 .table(self.table.get_table_name().expect("Checked above").clone())
447 .rename_column(
448 renamed_from.to_owned(),
449 column_def.get_column_name(),
450 ),
451 )
452 .await?;
453 }
454 continue;
455 };
456 if let (Some(desired), Some(current)) = (
459 column_def.get_column_type(),
460 existing_column.get_column_type(),
461 ) {
462 let backend = db.get_database_backend();
463 let desired_sql = render_column_type(backend, desired);
464 let current_sql = render_column_type(backend, current);
465 if desired_sql != current_sql {
466 tracing::warn!(
467 "schema sync: column `{}`.`{}` is `{}` in the database but the entity \
468 defines `{}`; sync is non-destructive and will not alter it (apply the \
469 change with a migration)",
470 table_name.1.to_string(),
471 column_def.get_column_name(),
472 current_sql,
473 desired_sql,
474 );
475 }
476 }
477 }
478 if db.get_database_backend() != DbBackend::Sqlite {
479 for foreign_key in self.table.get_foreign_key_create_stmts().iter() {
480 let mut key_exists = false;
481 for existing_key in existing_table.get_foreign_key_create_stmts().iter() {
482 if compare_foreign_key(foreign_key, existing_key) {
483 key_exists = true;
484 break;
485 }
486 }
487 if !key_exists {
488 db.execute(foreign_key).await?;
489 }
490 }
491 }
492 } else {
493 db.execute(&self.table).await?;
494 }
495 for stmt in self.indexes.iter() {
496 let mut has_index = false;
497 if let Some(existing_table) = existing_table {
498 for existing_index in existing_table.get_indexes() {
499 if existing_index.get_index_spec().get_column_names()
500 == stmt.get_index_spec().get_column_names()
501 {
502 has_index = true;
503 break;
504 }
505 }
506 }
507 if !has_index {
508 let mut stmt = stmt.clone();
510 stmt.if_not_exists();
511 db.execute(&stmt).await?;
512 }
513 }
514 if let Some(existing_table) = existing_table {
515 for column_def in self.table.get_columns() {
518 if column_def.get_column_spec().unique {
519 let col_name = column_def.get_column_name();
520 let col_exists = existing_table
521 .get_columns()
522 .iter()
523 .any(|c| c.get_column_name() == col_name);
524 if !col_exists {
525 continue;
528 }
529 let already_unique = existing_table.get_indexes().iter().any(|idx| {
530 if !idx.is_unique_key() {
531 return false;
532 }
533 let cols = idx.get_index_spec().get_column_names();
534 cols.len() == 1 && cols[0] == col_name
535 });
536 if !already_unique {
537 let table_name =
538 self.table.get_table_name().expect("table must have a name");
539 let tbl_str = table_name.sea_orm_table().to_string();
540 let table_ref = index_table_ref(table_name.clone(), db_backend);
541 db.execute(
542 Index::create()
543 .name(format!("idx-{tbl_str}-{col_name}"))
544 .table(table_ref)
545 .col(col_name.into_iden())
546 .unique()
547 .if_not_exists(),
548 )
549 .await?;
550 }
551 }
552 }
553 }
554 if let Some(existing_table) = existing_table {
555 for existing_index in existing_table.get_indexes() {
558 if existing_index.is_unique_key() {
559 let mut has_index = false;
560 for stmt in self.indexes.iter() {
561 if existing_index.get_index_spec().get_column_names()
562 == stmt.get_index_spec().get_column_names()
563 {
564 has_index = true;
565 break;
566 }
567 }
568 if !has_index {
573 let index_cols = existing_index.get_index_spec().get_column_names();
574 if index_cols.len() == 1 {
575 for column_def in self.table.get_columns() {
576 if column_def.get_column_name() == index_cols[0]
577 && column_def.get_column_spec().unique
578 {
579 has_index = true;
580 break;
581 }
582 }
583 }
584 }
585 if !has_index
586 && let Some(drop_existing) = existing_index
587 .get_index_spec()
588 .get_name()
589 .map(|s| s.to_owned())
590 {
591 if db_backend == DbBackend::Postgres {
592 db.execute(
597 TableAlterStatement::new()
598 .table(
599 self.table.get_table_name().expect("Checked above").clone(),
600 )
601 .drop_constraint(drop_existing),
602 )
603 .await?;
604 } else {
605 db.execute(sea_query::Index::drop().name(drop_existing))
606 .await?;
607 }
608 }
609 }
610 }
611 }
612 Ok(())
613 }
614
615 fn debug_print(
616 &self,
617 f: &mut std::fmt::Formatter<'_>,
618 backend: &DbBackend,
619 ) -> std::fmt::Result {
620 write!(f, "EntitySchemaInfo {{")?;
621 write!(f, " table: {:?}", backend.build(&self.table).to_string())?;
622 write!(f, " enums: [")?;
623 for (i, stmt) in self.enums.iter().enumerate() {
624 if i > 0 {
625 write!(f, ", ")?;
626 }
627 write!(f, "{:?}", backend.build(stmt).to_string())?;
628 }
629 write!(f, " ]")?;
630 write!(f, " indexes: [")?;
631 for (i, stmt) in self.indexes.iter().enumerate() {
632 if i > 0 {
633 write!(f, ", ")?;
634 }
635 write!(f, "{:?}", backend.build(stmt).to_string())?;
636 }
637 write!(f, " ]")?;
638 write!(f, " }}")
639 }
640}
641
642fn get_table_name(table_ref: Option<&TableRef>) -> TableName {
643 match table_ref {
644 Some(TableRef::Table(table_name, _)) => table_name.clone(),
645 None => panic!("Expect TableCreateStatement is properly built"),
646 _ => unreachable!("Unexpected {table_ref:?}"),
647 }
648}
649
650#[allow(dead_code)]
655fn render_column_type(backend: DbBackend, col_type: &sea_query::ColumnType) -> String {
656 use sea_query::backend::TableBuilder;
657 let mut sql = String::new();
658 match backend {
659 DbBackend::MySql => sea_query::MysqlQueryBuilder.prepare_column_type(col_type, &mut sql),
660 DbBackend::Postgres => {
661 sea_query::PostgresQueryBuilder.prepare_column_type(col_type, &mut sql)
662 }
663 DbBackend::Sqlite => sea_query::SqliteQueryBuilder.prepare_column_type(col_type, &mut sql),
664 }
665 sql
666}
667
668fn compare_foreign_key(a: &ForeignKeyCreateStatement, b: &ForeignKeyCreateStatement) -> bool {
669 let a = a.get_foreign_key();
670 let b = b.get_foreign_key();
671
672 a.get_name() == b.get_name()
673 || (a.get_ref_table() == b.get_ref_table()
674 && a.get_columns() == b.get_columns()
675 && a.get_ref_columns() == b.get_ref_columns())
676}