Add support for primary key constraints in create_table operation

docs/operations/create_table.mdx

@@ -61,7 +61,10 @@ Each `constraint` is defined as:
 },
 ```
 
-Supported constraint types: `unique`, `check`.
+Supported constraint types: `unique`, `check`, `primary_key`.
+
+Please note that you can only configure primary keys in `columns` list or `constraints` list, but
+not in both places.
 
 ## Examples
 

internal/jsonschema/testdata/create-table-10-invalid-primary-key-constraints-extra-check.txtar

@@ -0,0 +1,38 @@
+This is an invalid 'create_table' migration.
+Primary key constraint must not constain a check expression.
+
+-- create_table.json --
+{
+  "name": "migration_name",
+  "operations": [
+    {
+      "create_table": {
+        "name": "posts",
+        "columns": [
+          {
+            "name": "title",
+            "type": "varchar(255)"
+          },
+          {
+            "name": "user_id",
+            "type": "integer",
+            "nullable": true
+          }
+        ],
+        "constraints": [
+          {
+            "name": "my_invalid_pk",
+            "type": "primary_key",
+            "columns": [
+              "title"
+            ],
+            "check": "this should not be set"
+          }
+        ]
+      }
+    }
+  ]
+}
+
+-- valid --
+false

internal/jsonschema/testdata/create-table-11-invalid-primary-key-constraints-missing-columns.txtar

@@ -0,0 +1,34 @@
+This is an invalid 'create_table' migration.
+Primary key constraint must have columns set
+
+-- create_table.json --
+{
+  "name": "migration_name",
+  "operations": [
+    {
+      "create_table": {
+        "name": "posts",
+        "columns": [
+          {
+            "name": "title",
+            "type": "varchar(255)"
+          },
+          {
+            "name": "user_id",
+            "type": "integer",
+            "nullable": true
+          }
+        ],
+        "constraints": [
+          {
+            "name": "my_invalid_pk",
+            "type": "primary_key"
+          }
+        ]
+      }
+    }
+  ]
+}
+
+-- valid --
+false

internal/jsonschema/testdata/create-table-12-valid-primary-key-constraint.txtar

@@ -0,0 +1,36 @@
+This is a valid 'create_table' migration.
+
+-- create_table.json --
+{
+  "name": "migration_name",
+  "operations": [
+    {
+      "create_table": {
+        "name": "posts",
+        "columns": [
+          {
+            "name": "title",
+            "type": "varchar(255)"
+          },
+          {
+            "name": "user_id",
+            "type": "integer"
+          }
+        ],
+        "constraints": [
+          {
+            "name": "my_pk",
+            "type": "primary_key",
+            "columns": [
+              "title",
+              "user_id"
+            ]
+          }
+        ]
+      }
+    }
+  ]
+}
+
+-- valid --
+true

pkg/migrations/errors.go

@@ -258,3 +258,11 @@ type MultiColumnConstraintsNotSupportedError struct {
 func (e MultiColumnConstraintsNotSupportedError) Error() string {
 	return fmt.Sprintf("constraint %q on table %q applies to multiple columns", e.Constraint, e.Table)
 }
+
+type PrimaryKeysAreAlreadySetError struct {
+	Table string
+}
+
+func (e PrimaryKeysAreAlreadySetError) Error() string {
+	return fmt.Sprintf("table %q already has a primary key configuration in columns list", e.Table)
+}

pkg/migrations/op_common_test.go

@@ -250,6 +250,13 @@ func NotValidatedForeignKeyMustExist(t *testing.T, db *sql.DB, schema, table, co
 	}
 }
 
+func PrimaryKeyConstraintMustExist(t *testing.T, db *sql.DB, schema, table, constraint string) {
+	t.Helper()
+	if !primaryKeyConstraintExists(t, db, schema, table, constraint) {
+		t.Fatalf("Expected constraint %q to exist", constraint)
+	}
+}
+
 func IndexMustExist(t *testing.T, db *sql.DB, schema, table, index string) {
 	t.Helper()
 	if !indexExists(t, db, schema, table, index) {
@@ -403,6 +410,26 @@ func foreignKeyExists(t *testing.T, db *sql.DB, schema, table, constraint string
 	return exists
 }
 
+func primaryKeyConstraintExists(t *testing.T, db *sql.DB, schema, table, constraint string) bool {
+	t.Helper()
+
+	var exists bool
+	err := db.QueryRow(`
+    SELECT EXISTS (
+      SELECT 1
+      FROM pg_catalog.pg_constraint
+      WHERE conrelid = $1::regclass
+      AND conname = $2
+      AND contype = 'p'
+    )`,
+		fmt.Sprintf("%s.%s", schema, table), constraint).Scan(&exists)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return exists
+}
+
 func triggerExists(t *testing.T, db *sql.DB, schema, table, trigger string) bool {
 	t.Helper()
 

pkg/migrations/op_create_table.go

@@ -79,6 +79,7 @@ func (o *OpCreateTable) Validate(ctx context.Context, s *schema.Schema) error {
 		return TableAlreadyExistsError{Name: o.Name}
 	}
 
+	hasPrimaryKeyColumns := false
 	for _, col := range o.Columns {
 		if err := ValidateIdentifierLength(col.Name); err != nil {
 			return fmt.Errorf("invalid column: %w", err)
@@ -106,6 +107,10 @@ func (o *OpCreateTable) Validate(ctx context.Context, s *schema.Schema) error {
 				}
 			}
 		}
+
+		if col.Pk {
+			hasPrimaryKeyColumns = true
+		}
 	}
 
 	for _, c := range o.Constraints {
@@ -146,6 +151,14 @@ func (o *OpCreateTable) Validate(ctx context.Context, s *schema.Schema) error {
 					Err:   fmt.Errorf("CHECK constraints cannot have NULLS NOT DISTINCT"),
 				}
 			}
+		case ConstraintTypePrimaryKey:
+			// users can only set primary keys either in columns list or in constraint list
+			if hasPrimaryKeyColumns {
+				return PrimaryKeysAreAlreadySetError{Table: o.Name}
+			}
+			if len(c.Columns) == 0 {
+				return FieldRequiredError{Name: "columns"}
+			}
 		}
 	}
 
@@ -160,12 +173,16 @@ func (o *OpCreateTable) Validate(ctx context.Context, s *schema.Schema) error {
 // the new table.
 func (o *OpCreateTable) updateSchema(s *schema.Schema) *schema.Schema {
 	columns := make(map[string]*schema.Column, len(o.Columns))
+	primaryKeys := make([]string, 0)
 	for _, col := range o.Columns {
 		columns[col.Name] = &schema.Column{
 			Name:     col.Name,
 			Unique:   col.Unique,
 			Nullable: col.Nullable,
 		}
+		if col.Pk {
+			primaryKeys = append(primaryKeys, col.Name)
+		}
 	}
 	uniqueConstraints := make(map[string]*schema.UniqueConstraint, 0)
 	checkConstraints := make(map[string]*schema.CheckConstraint, 0)
@@ -182,14 +199,8 @@ func (o *OpCreateTable) updateSchema(s *schema.Schema) *schema.Schema {
 				Columns:    c.Columns,
 				Definition: c.Check,
 			}
-		}
-	}
-
-	// Build the table's primary key from the columns that have the `Pk` flag set
-	var primaryKey []string
-	for _, col := range o.Columns {
-		if col.Pk {
-			primaryKey = append(primaryKey, col.Name)
+		case ConstraintTypePrimaryKey:
+			primaryKeys = c.Columns
 		}
 	}
 
@@ -198,7 +209,7 @@ func (o *OpCreateTable) updateSchema(s *schema.Schema) *schema.Schema {
 		Columns:           columns,
 		UniqueConstraints: uniqueConstraints,
 		CheckConstraints:  checkConstraints,
-		PrimaryKey:        primaryKey,
+		PrimaryKey:        primaryKeys,
 	})
 
 	return s
@@ -219,12 +230,14 @@ func columnsToSQL(cols []Column, tr SQLTransformer) (string, error) {
 		sql += colSQL
 
 		if col.IsPrimaryKey() {
-			primaryKeys = append(primaryKeys, pq.QuoteIdentifier(col.Name))
+			primaryKeys = append(primaryKeys, col.Name)
 		}
 	}
 
+	// Add primary key constraint if there are primary key columns.
 	if len(primaryKeys) > 0 {
-		sql += fmt.Sprintf(", PRIMARY KEY (%s)", strings.Join(primaryKeys, ", "))
+		writer := &ConstraintSQLWriter{Columns: primaryKeys}
+		sql += ", " + writer.WritePrimaryKey()
 	}
 	return sql, nil
 }
@@ -249,6 +262,8 @@ func constraintsToSQL(constraints []Constraint) (string, error) {
 			constraintsSQL[i] = writer.WriteUnique(c.NullsNotDistinct)
 		case ConstraintTypeCheck:
 			constraintsSQL[i] = writer.WriteCheck(c.Check, c.NoInherit)
+		case ConstraintTypePrimaryKey:
+			constraintsSQL[i] = writer.WritePrimaryKey()
 		}
 	}
 	if len(constraintsSQL) == 0 {
@@ -296,6 +311,17 @@ func (w *ConstraintSQLWriter) WriteCheck(check string, noInherit bool) string {
 	return constraint
 }
 
+func (w *ConstraintSQLWriter) WritePrimaryKey() string {
+	constraint := ""
+	if w.Name != "" {
+		constraint = fmt.Sprintf("CONSTRAINT %s ", pq.QuoteIdentifier(w.Name))
+	}
+	constraint += fmt.Sprintf("PRIMARY KEY (%s)", strings.Join(quoteColumnNames(w.Columns), ", "))
+	constraint += w.addIndexParameters()
+	constraint += w.addDeferrable()
+	return constraint
+}
+
 func (w *ConstraintSQLWriter) addIndexParameters() string {
 	constraint := ""
 	if len(w.IncludeColumns) != 0 {