Add references when adding/changing pg columns

Electron100 · Apr 4, 2024 · c693c05 · c693c05
1 parent d7d7a62
commit c693c05
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 12 deletions.
diff --git a/butane_core/src/db/pg.rs b/butane_core/src/db/pg.rs
@@ -578,7 +578,7 @@ fn create_table_constraints(table: &ATable) -> String {
         .columns
         .iter()
         .filter(|column| column.reference().is_some())
-        .map(|column| define_constraint(table, column))
+        .map(|column| define_constraint(&table.name, column))
         .collect::<Vec<String>>()
         .join("\n")
 }
@@ -602,7 +602,7 @@ fn define_column(col: &AColumn) -> Result<String> {
     ))
 }
 
-fn define_constraint(table: &ATable, column: &AColumn) -> String {
+fn define_constraint(table_name: &str, column: &AColumn) -> String {
     let reference = column
         .reference()
         .as_ref()
@@ -611,7 +611,7 @@ fn define_constraint(table: &ATable, column: &AColumn) -> String {
         ARef::Literal(literal) => {
             format!(
                 "ALTER TABLE {} ADD FOREIGN KEY ({}) REFERENCES {}({});",
-                helper::quote_reserved_word(&table.name),
+                helper::quote_reserved_word(&table_name),
                 helper::quote_reserved_word(column.name()),
                 helper::quote_reserved_word(literal.table_name()),
                 helper::quote_reserved_word(literal.column_name()),
@@ -658,12 +658,17 @@ fn drop_table(name: &str) -> String {
 
 fn add_column(tbl_name: &str, col: &AColumn) -> Result<String> {
     let default: SqlVal = helper::column_default(col)?;
-    Ok(format!(
+    let mut stmts = vec![format!(
         "ALTER TABLE {} ADD COLUMN {} DEFAULT {};",
         helper::quote_reserved_word(tbl_name),
         define_column(col)?,
         helper::sql_literal_value(default)?
-    ))
+    )];
+    if col.reference().is_some() {
+        stmts.push(define_constraint(tbl_name, col));
+    }
+    let result = stmts.join("\n");
+    Ok(result)
 }
 
 fn remove_column(tbl_name: &str, name: &str) -> String {
@@ -715,16 +720,18 @@ fn change_column(
         Some(col) => new_table.replace_column(col.clone()),
         None => new_table.remove_column(old.name()),
     }
-    let stmts: [&str; 4] = [
-        &create_table(&new_table, false)?,
-        &copy_table(old_table, &new_table),
-        &drop_table(&old_table.name),
-        &format!(
+    let mut stmts: Vec<String> = vec![
+        create_table(&new_table, false)?,
+        create_table_constraints(&new_table),
+        copy_table(old_table, &new_table),
+        drop_table(&old_table.name),
+        format!(
             "ALTER TABLE {} RENAME TO {};",
             helper::quote_reserved_word(&new_table.name),
             helper::quote_reserved_word(tbl_name)
         ),
     ];
+    stmts.retain(|stmt| !stmt.is_empty());
     let result = stmts.join("\n");
     new_table.name.clone_from(&old_table.name);
     current.replace_table(new_table);

diff --git a/examples/getting_started/.butane/migrations/20240115_023841384_dbconstraints/pg_down.sql b/examples/getting_started/.butane/migrations/20240115_023841384_dbconstraints/pg_down.sql
@@ -100,6 +100,7 @@ CREATE TABLE Post_tags_Many__butane_tmp (
 owner INTEGER NOT NULL,
 has TEXT NOT NULL
 );
+ALTER TABLE Post_tags_Many__butane_tmp ADD FOREIGN KEY (owner) REFERENCES Post(id);
 INSERT INTO Post_tags_Many__butane_tmp SELECT owner, has FROM Post_tags_Many;
 DROP TABLE Post_tags_Many;
 ALTER TABLE Post_tags_Many__butane_tmp RENAME TO Post_tags_Many;

diff --git a/examples/getting_started/.butane/migrations/20240115_023841384_dbconstraints/pg_up.sql b/examples/getting_started/.butane/migrations/20240115_023841384_dbconstraints/pg_up.sql
@@ -21,6 +21,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -33,6 +34,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -45,6 +47,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -57,6 +60,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -69,6 +73,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -81,6 +86,7 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
@@ -93,20 +99,24 @@ blog BIGINT NOT NULL,
 byline TEXT ,
 likes INTEGER NOT NULL
 );
+ALTER TABLE Post__butane_tmp ADD FOREIGN KEY (blog) REFERENCES Blog(id);
 INSERT INTO Post__butane_tmp SELECT id, title, body, published, blog, byline, likes FROM Post;
 DROP TABLE Post;
 ALTER TABLE Post__butane_tmp RENAME TO Post;
 CREATE TABLE Post_tags_Many__butane_tmp (
 owner INTEGER NOT NULL,
 has TEXT NOT NULL
 );
+ALTER TABLE Post_tags_Many__butane_tmp ADD FOREIGN KEY (has) REFERENCES Tag(tag);
 INSERT INTO Post_tags_Many__butane_tmp SELECT owner, has FROM Post_tags_Many;
 DROP TABLE Post_tags_Many;
 ALTER TABLE Post_tags_Many__butane_tmp RENAME TO Post_tags_Many;
 CREATE TABLE Post_tags_Many__butane_tmp (
 owner INTEGER NOT NULL,
 has TEXT NOT NULL
 );
+ALTER TABLE Post_tags_Many__butane_tmp ADD FOREIGN KEY (owner) REFERENCES Post(id);
+ALTER TABLE Post_tags_Many__butane_tmp ADD FOREIGN KEY (has) REFERENCES Tag(tag);
 INSERT INTO Post_tags_Many__butane_tmp SELECT owner, has FROM Post_tags_Many;
 DROP TABLE Post_tags_Many;
 ALTER TABLE Post_tags_Many__butane_tmp RENAME TO Post_tags_Many;