Add constraints option to create_table and unique constraint support

docs/operations/create_table.mdx

@@ -9,7 +9,8 @@ description: A create table operation creates a new table in the database.
 {
   "create_table": {
     "name": "name of new table",
-    "columns": [...]
+    "columns": [...],
+    "constraints": [...]
   }
 }
 ```
@@ -40,6 +41,26 @@ Each `column` is defined as:
 
 Default values are subject to the usual rules for quoting SQL expressions. In particular, string literals should be surrounded with single quotes.
 
+Each `constraint` is defined as:
+
+```json
+{
+  "name": "constraint name",
+  "type": "constraint type",
+  "columns": ["list", "of", "columns"],
+  "nulls_not_distinct": true|false,
+  "deferrable": true|false,
+  "initially_deferred": true|false,
+  "index_parameters": {
+    "tablespace": "index_tablespace",
+    "storage_parameters": "parameter=value",
+    "include_columns": ["list", "of", "columns", "included in index"]
+  },
+},
+```
+
+Supported constraint types: `unique`.
+
 ## Examples
 
 ### Create multiple tables
@@ -98,3 +119,9 @@ Create a table with a `CHECK` constraint on one column:
 Create a table with different `DEFAULT` values:
 
 <ExampleSnippet example="28_different_defaults.json" language="json" />
+
+### Create a table with table level unique constraint
+
+Create a table with table level constraints:
+
+<ExampleSnippet example="50_create_table_with_table_constraint.json" language="json" />

examples/.ledger

@@ -47,3 +47,4 @@
 47_add_table_foreign_key_constraint.json
 48_drop_tickets_check.json
 49_unset_not_null_on_indexed_column.json
+50_create_table_with_table_constraint.json

examples/50_create_table_with_table_constraint.json

@@ -0,0 +1,43 @@
+{
+  "name": "50_create_table_with_table_constraint",
+  "operations": [
+    {
+      "create_table": {
+        "name": "phonebook",
+        "columns": [
+          {
+            "name": "id",
+            "type": "serial",
+            "pk": true
+          },
+          {
+            "name": "name",
+            "type": "varchar(255)"
+          },
+          {
+            "name": "city",
+            "type": "varchar(255)"
+          },
+          {
+            "name": "phone",
+            "type": "varchar(255)"
+          }
+        ],
+        "constraints": [
+          {
+            "name": "unique_numbers",
+            "type": "unique",
+            "columns": [
+              "phone"
+            ],
+            "index_parameters": {
+              "include_columns": [
+                "name"
+              ]
+            }
+          }
+        ]
+      }
+    }
+  ]
+}

pkg/migrations/op_create_table.go

@@ -22,10 +22,16 @@ func (o *OpCreateTable) Start(ctx context.Context, conn db.DB, latestSchema stri
 		return nil, fmt.Errorf("failed to create columns SQL: %w", err)
 	}
 
+	constraintsSQL, err := constraintsToSQL(o.Constraints)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create constraints SQL: %w", err)
+	}
+
 	// Create the table
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("CREATE TABLE %s (%s)",
+	_, err = conn.ExecContext(ctx, fmt.Sprintf("CREATE TABLE %s (%s %s)",
 		pq.QuoteIdentifier(o.Name),
-		columnsSQL))
+		columnsSQL,
+		constraintsSQL))
 	if err != nil {
 		return nil, err
 	}
@@ -102,6 +108,22 @@ func (o *OpCreateTable) Validate(ctx context.Context, s *schema.Schema) error {
 		}
 	}
 
+	for _, c := range o.Constraints {
+		if c.Name == "" {
+			return FieldRequiredError{Name: "name"}
+		}
+		if err := ValidateIdentifierLength(c.Name); err != nil {
+			return fmt.Errorf("invalid constraint: %w", err)
+		}
+
+		switch c.Type { //nolint:gocritic // more cases are coming soon
+		case ConstraintTypeUnique:
+			if len(c.Columns) == 0 {
+				return FieldRequiredError{Name: "columns"}
+			}
+		}
+	}
+
 	// Update the schema to ensure that the new table is visible to validation of
 	// subsequent operations.
 	o.updateSchema(s)
@@ -118,9 +140,23 @@ func (o *OpCreateTable) updateSchema(s *schema.Schema) *schema.Schema {
 			Name: col.Name,
 		}
 	}
+	var uniqueConstraints map[string]*schema.UniqueConstraint
+	for _, c := range o.Constraints {
+		switch c.Type { //nolint:gocritic // more cases are coming soon
+		case ConstraintTypeUnique:
+			if uniqueConstraints == nil {
+				uniqueConstraints = make(map[string]*schema.UniqueConstraint)
+			}
+			uniqueConstraints[c.Name] = &schema.UniqueConstraint{
+				Name:    c.Name,
+				Columns: c.Columns,
+			}
+		}
+	}
 	s.AddTable(o.Name, &schema.Table{
-		Name:    o.Name,
-		Columns: columns,
+		Name:              o.Name,
+		Columns:           columns,
+		UniqueConstraints: uniqueConstraints,
 	})
 
 	return s
@@ -150,3 +186,88 @@ func columnsToSQL(cols []Column, tr SQLTransformer) (string, error) {
 	}
 	return sql, nil
 }
+
+func constraintsToSQL(constraints []Constraint) (string, error) {
+	constraintsSQL := make([]string, len(constraints))
+	for i, c := range constraints {
+		writer := &ConstraintSQLWriter{
+			Name:              c.Name,
+			Columns:           c.Columns,
+			InitiallyDeferred: c.InitiallyDeferred,
+			Deferrable:        c.Deferrable,
+		}
+		if c.IndexParameters != nil {
+			writer.IncludeColumns = c.IndexParameters.IncludeColumns
+			writer.StorageParameters = c.IndexParameters.StorageParameters
+			writer.Tablespace = c.IndexParameters.Tablespace
+		}
+
+		switch c.Type { //nolint:gocritic // more cases are coming soon
+		case ConstraintTypeUnique:
+			constraintsSQL[i] = writer.WriteUnique(c.NullsNotDistinct)
+		}
+	}
+	if len(constraintsSQL) == 0 {
+		return "", nil
+	}
+	return ", " + strings.Join(constraintsSQL, ", "), nil
+}
+
+type ConstraintSQLWriter struct {
+	Name              string
+	Columns           []string
+	InitiallyDeferred bool
+	Deferrable        bool
+
+	// unique, exclude, primary key constraints support the following options
+	IncludeColumns    []string
+	StorageParameters string
+	Tablespace        string
+}
+
+func (w *ConstraintSQLWriter) WriteUnique(nullsNotDistinct bool) string {
+	var constraint string
+	if w.Name != "" {
+		constraint = fmt.Sprintf("CONSTRAINT %s ", pq.QuoteIdentifier(w.Name))
+	}
+	nullsDistinct := ""
+	if nullsNotDistinct {
+		nullsDistinct = "NULLS NOT DISTINCT"
+	}
+	constraint += fmt.Sprintf("UNIQUE %s (%s)", nullsDistinct, strings.Join(quoteColumnNames(w.Columns), ", "))
+	constraint += w.addIndexParameters()
+	constraint += w.addDeferrable()
+	return constraint
+}
+
+func (w *ConstraintSQLWriter) addIndexParameters() string {
+	constraint := ""
+	if len(w.IncludeColumns) != 0 {
+		constraint += fmt.Sprintf(" INCLUDE (%s)", strings.Join(quoteColumnNames(w.IncludeColumns), ", "))
+	}
+	if w.StorageParameters != "" {
+		constraint += fmt.Sprintf(" WITH (%s)", w.StorageParameters)
+	}
+	if w.Tablespace != "" {
+		constraint += fmt.Sprintf(" USING INDEX TABLESPACE %s", w.Tablespace)
+	}
+	return constraint
+}
+
+func (w *ConstraintSQLWriter) addDeferrable() string {
+	if !w.InitiallyDeferred && !w.Deferrable {
+		return ""
+	}
+	deferrable := ""
+	if w.Deferrable {
+		deferrable += " DEFERRABLE"
+	} else {
+		deferrable += " NOT DEFERRABLE"
+	}
+	if w.InitiallyDeferred {
+		deferrable += " INITIALLY DEFERRED"
+	} else {
+		deferrable += " INITIALLY IMMEDIATE"
+	}
+	return deferrable
+}

pkg/migrations/op_create_table_test.go

@@ -483,6 +483,70 @@ func TestCreateTable(t *testing.T) {
 				ColumnMustHaveComment(t, db, schema, "users", "name", "the username")
 			},
 		},
+		{
+			name: "create table with a unique table constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "users",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name: "name",
+									Type: "text",
+								},
+							},
+							Constraints: []migrations.Constraint{
+								{
+									Name: "unique_name",
+									Type: migrations.ConstraintTypeUnique,
+									Columns: []string{
+										"name",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// The unique constraint exists on the new table.
+				UniqueConstraintMustExist(t, db, schema, "users", "unique_name")
+
+				// Inserting a row into the table succeeds when the unique constraint is satisfied.
+				MustInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"name": "alice",
+				})
+
+				// Inserting a row into the table fails when the unique constraint is not satisfied.
+				MustNotInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"name": "alice",
+				}, testutils.UniqueViolationErrorCode)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The table has been dropped, so the unique constraint is gone.
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// The check constraint exists on the new table.
+				UniqueConstraintMustExist(t, db, schema, "users", "unique_name")
+
+				// Inserting a row into the table succeeds when the unique constraint is satisfied.
+				MustInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"name": "bobby",
+				})
+
+				// Inserting a row into the table fails when the unique constraint is not satisfied.
+				MustNotInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"name": "bobby",
+				}, testutils.UniqueViolationErrorCode)
+			},
+		},
 	})
 }
 
@@ -601,6 +665,38 @@ func TestCreateTableValidation(t *testing.T) {
 			},
 			wantStartErr: migrations.InvalidIdentifierLengthError{Name: invalidName},
 		},
+		{
+			name: "missing column list in unique constraint",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "table1",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:   "name",
+									Type:   "varchar(255)",
+									Unique: true,
+								},
+							},
+							Constraints: []migrations.Constraint{
+								{
+									Name: "unique_name",
+									Type: migrations.ConstraintTypeUnique,
+								},
+							},
+						},
+					},
+				},
+			},
+			wantStartErr: migrations.FieldRequiredError{Name: "columns"},
+		},
 	})
 }