Implement adding uniqueness constraints to columns

examples/14_add_reviews_table.json

@@ -0,0 +1,29 @@
+{
+  "name": "14_add_reviews_table",
+  "operations": [
+    {
+      "create_table": {
+        "name": "reviews",
+        "columns": [
+          {
+            "name": "id",
+            "type": "serial",
+            "pk": true
+          },
+          {
+            "name": "username",
+            "type": "text"
+          },
+          {
+            "name": "product",
+            "type": "text"
+          },
+          {
+            "name": "review",
+            "type": "text"
+          }
+        ]
+      }
+    }
+  ]
+}

examples/15_set_columns_unique.json

@@ -0,0 +1,15 @@
+{
+  "name": "15_set_columns_unique",
+  "operations": [
+    {
+      "set_unique": {
+        "name": "reviews_username_product_unique",
+        "table": "reviews",
+        "columns": [
+          "username",
+          "product"
+        ]
+      }
+    }
+  ]
+}

pkg/migrations/op_common.go

@@ -19,6 +19,7 @@ const (
 	OpNameCreateIndex  OpName = "create_index"
 	OpNameDropIndex    OpName = "drop_index"
 	OpNameRenameColumn OpName = "rename_column"
+	OpNameSetUnique    OpName = "set_unique"
 )
 
 func TemporaryName(name string) string {
@@ -98,6 +99,9 @@ func (v *Operations) UnmarshalJSON(data []byte) error {
 		case OpNameDropIndex:
 			item = &OpDropIndex{}
 
+		case OpNameSetUnique:
+			item = &OpSetUnique{}
+
 		default:
 			return fmt.Errorf("unknown migration type: %v", opName)
 		}
@@ -154,6 +158,9 @@ func (v Operations) MarshalJSON() ([]byte, error) {
 		case *OpDropIndex:
 			opName = OpNameDropIndex
 
+		case *OpSetUnique:
+			opName = OpNameSetUnique
+
 		default:
 			panic(fmt.Errorf("unknown operation for %T", op))
 		}

pkg/migrations/op_common_test.go

@@ -250,6 +250,13 @@ func ConstraintMustNotExist(t *testing.T, db *sql.DB, schema, table, constraint
 	}
 }
 
+func ConstraintMustExist(t *testing.T, db *sql.DB, schema, table, constraint string) {
+	t.Helper()
+	if !constraintExists(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) {

pkg/migrations/op_set_unique.go

@@ -0,0 +1,65 @@
+package migrations
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"strings"
+
+	"github.com/lib/pq"
+
+	"pg-roll/pkg/schema"
+)
+
+type OpSetUnique struct {
+	Name    string   `json:"name"`
+	Table   string   `json:"table"`
+	Columns []string `json:"columns"`
+}
+
+var _ Operation = (*OpSetUnique)(nil)
+
+func (o *OpSetUnique) Start(ctx context.Context, conn *sql.DB, schemaName string, stateSchema string, s *schema.Schema) error {
+	// create unique index concurrently
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS %s ON %s (%s)",
+		pq.QuoteIdentifier(o.Name),
+		pq.QuoteIdentifier(o.Table),
+		strings.Join(quoteColumnNames(o.Columns), ", ")))
+	return err
+}
+
+func (o *OpSetUnique) Complete(ctx context.Context, conn *sql.DB) error {
+	// create a unique constraint using the unique index
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s ADD CONSTRAINT %s UNIQUE USING INDEX %s",
+		pq.QuoteIdentifier(o.Table),
+		pq.QuoteIdentifier(o.Name),
+		pq.QuoteIdentifier(o.Name)))
+
+	return err
+}
+
+func (o *OpSetUnique) Rollback(ctx context.Context, conn *sql.DB) error {
+	// drop the index concurrently
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s", o.Name))
+
+	return err
+}
+
+func (o *OpSetUnique) Validate(ctx context.Context, s *schema.Schema) error {
+	if o.Name == "" {
+		return NameRequiredError{}
+	}
+
+	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
+
+	for _, column := range o.Columns {
+		if table.GetColumn(column) == nil {
+			return ColumnDoesNotExistError{Table: o.Table, Name: column}
+		}
+	}
+
+	return nil
+}

pkg/migrations/op_set_unique_test.go

@@ -0,0 +1,94 @@
+package migrations_test
+
+import (
+	"database/sql"
+	"testing"
+
+	"pg-roll/pkg/migrations"
+)
+
+func TestSetColumnsUnique(t *testing.T) {
+	t.Parallel()
+
+	ExecuteTests(t, TestCases{{
+		name: "set unique",
+		migrations: []migrations.Migration{
+			{
+				Name: "01_add_table",
+				Operations: migrations.Operations{
+					&migrations.OpCreateTable{
+						Name: "reviews",
+						Columns: []migrations.Column{
+							{
+								Name:       "id",
+								Type:       "serial",
+								PrimaryKey: true,
+							},
+							{
+								Name:     "username",
+								Type:     "text",
+								Nullable: false,
+							},
+							{
+								Name:     "product",
+								Type:     "text",
+								Nullable: false,
+							},
+							{
+								Name:     "review",
+								Type:     "text",
+								Nullable: false,
+							},
+						},
+					},
+				},
+			},
+			{
+				Name: "02_set_unique",
+				Operations: migrations.Operations{
+					&migrations.OpSetUnique{
+						Name:    "reviews_username_product_unique",
+						Table:   "reviews",
+						Columns: []string{"username", "product"},
+					},
+				},
+			},
+		},
+		afterStart: func(t *testing.T, db *sql.DB) {
+			// The unique index has been created on the underlying table.
+			IndexMustExist(t, db, "public", "reviews", "reviews_username_product_unique")
+
+			// Inserting values into the old schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "01_add_table", "reviews", map[string]string{
+				"username": "alice", "product": "apple", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "01_add_table", "reviews", map[string]string{
+				"username": "alice", "product": "apple", "review": "bad",
+			})
+
+			// Inserting values into the new schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "bob", "product": "orange", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "bob", "product": "orange", "review": "bad",
+			})
+		},
+		afterRollback: func(t *testing.T, db *sql.DB) {
+			// The unique index has been dropped from the the underlying table.
+			IndexMustNotExist(t, db, "public", "reviews", "reviews_username_product_unique")
+		},
+		afterComplete: func(t *testing.T, db *sql.DB) {
+			// The unique constraint has been created on the underlying table.
+			ConstraintMustExist(t, db, "public", "reviews", "reviews_username_product_unique")
+
+			// Inserting values into the new schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "carl", "product": "banana", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "carl", "product": "banana", "review": "bad",
+			})
+		},
+	}})
+}