Merge pull request #646 from SeaQL/pg-alter-pk-columns

Fix Postgres alter primary key column statements
SeaQL · Jun 29, 2023 · 014b3c5 · 014b3c5
2 parents 1e83626 + e740209
commit 014b3c5
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 19 deletions.
diff --git a/src/backend/postgres/table.rs b/src/backend/postgres/table.rs
@@ -2,20 +2,10 @@ use super::*;
 
 impl TableBuilder for PostgresQueryBuilder {
     fn prepare_column_def(&self, column_def: &ColumnDef, sql: &mut dyn SqlWriter) {
-        column_def.name.prepare(sql.as_writer(), self.quote());
-
-        self.prepare_column_type_check_auto_increment(column_def, sql);
-
-        for column_spec in column_def.spec.iter() {
-            if let ColumnSpec::AutoIncrement = column_spec {
-                continue;
-            }
-            if let ColumnSpec::Comment(_) = column_spec {
-                continue;
-            }
-            write!(sql, " ").unwrap();
-            self.prepare_column_spec(column_spec, sql);
-        }
+        let f = |column_def: &ColumnDef, sql: &mut dyn SqlWriter| {
+            self.prepare_column_type_check_auto_increment(column_def, sql);
+        };
+        self.prepare_column_def_common(column_def, sql, f);
     }
 
     fn prepare_column_type(&self, column_type: &ColumnType, sql: &mut dyn SqlWriter) {
@@ -119,14 +109,20 @@ impl TableBuilder for PostgresQueryBuilder {
                     if *if_not_exists {
                         write!(sql, "IF NOT EXISTS ").unwrap();
                     }
-                    self.prepare_column_def(column, sql);
+                    let f = |column_def: &ColumnDef, sql: &mut dyn SqlWriter| {
+                        if let Some(column_type) = &column_def.types {
+                            write!(sql, " ").unwrap();
+                            self.prepare_column_type(column_type, sql);
+                        }
+                    };
+                    self.prepare_column_def_common(column, sql, f);
                 }
                 TableAlterOption::ModifyColumn(column_def) => {
-                    if column_def.types.is_some() {
+                    if let Some(column_type) = &column_def.types {
                         write!(sql, "ALTER COLUMN ").unwrap();
                         column_def.name.prepare(sql.as_writer(), self.quote());
-                        write!(sql, " TYPE").unwrap();
-                        self.prepare_column_type_check_auto_increment(column_def, sql);
+                        write!(sql, " TYPE ").unwrap();
+                        self.prepare_column_type(column_type, sql);
                     }
                     let first = column_def.types.is_none();
 
@@ -254,4 +250,24 @@ impl PostgresQueryBuilder {
             }
         }
     }
+
+    fn prepare_column_def_common<F>(&self, column_def: &ColumnDef, sql: &mut dyn SqlWriter, f: F)
+    where
+        F: Fn(&ColumnDef, &mut dyn SqlWriter),
+    {
+        column_def.name.prepare(sql.as_writer(), self.quote());
+
+        f(column_def, sql);
+
+        for column_spec in column_def.spec.iter() {
+            if let ColumnSpec::AutoIncrement = column_spec {
+                continue;
+            }
+            if let ColumnSpec::Comment(_) = column_spec {
+                continue;
+            }
+            write!(sql, " ").unwrap();
+            self.prepare_column_spec(column_spec, sql);
+        }
+    }
 }
diff --git a/tests/postgres/table.rs b/tests/postgres/table.rs
@@ -493,6 +493,7 @@ fn alter_8() {
 
 #[test]
 fn alter_9() {
+    // https://dbfiddle.uk/98Vd8pmn
     assert_eq!(
         Table::alter()
             .table(Glyph::Table)
@@ -507,7 +508,7 @@ fn alter_9() {
             .to_string(PostgresQueryBuilder),
         [
             r#"ALTER TABLE "glyph""#,
-            r#"ALTER COLUMN "aspect" TYPE serial,"#,
+            r#"ALTER COLUMN "aspect" TYPE integer,"#,
             r#"ALTER COLUMN "aspect" SET NOT NULL,"#,
             r#"ADD UNIQUE ("aspect"),"#,
             r#"ADD PRIMARY KEY ("aspect")"#,
@@ -516,6 +517,29 @@ fn alter_9() {
     );
 }
 
+#[test]
+fn alter_10() {
+    // https://dbfiddle.uk/PQksflGf
+    assert_eq!(
+        Table::alter()
+            .table(Glyph::Table)
+            .add_column(
+                ColumnDef::new(Glyph::Aspect)
+                    .integer()
+                    .auto_increment()
+                    .not_null()
+                    .unique_key()
+                    .primary_key()
+            )
+            .to_string(PostgresQueryBuilder),
+        [
+            r#"ALTER TABLE "glyph""#,
+            r#"ADD COLUMN "aspect" integer NOT NULL UNIQUE PRIMARY KEY"#,
+        ]
+        .join(" ")
+    );
+}
+
 #[test]
 fn rename_1() {
     assert_eq!(