From 9d50a22527c03bf7b7ac7db5370ecf6701005601 Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Tue, 15 Aug 2023 09:11:48 +0100
Subject: [PATCH] Implement the set unique operation

---
 pkg/migrations/op_common.go          |  7 ++++
 pkg/migrations/op_set_unique.go      | 62 ++++++++++++++++++++++++++++
 pkg/migrations/op_set_unique_test.go | 10 ++---
 3 files changed, 73 insertions(+), 6 deletions(-)
 create mode 100644 pkg/migrations/op_set_unique.go

diff --git a/pkg/migrations/op_common.go b/pkg/migrations/op_common.go
index 9ca96639..0b54d8fc 100644
--- a/pkg/migrations/op_common.go
+++ b/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))
 		}
diff --git a/pkg/migrations/op_set_unique.go b/pkg/migrations/op_set_unique.go
new file mode 100644
index 00000000..3c1d06e8
--- /dev/null
+++ b/pkg/migrations/op_set_unique.go
@@ -0,0 +1,62 @@
+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 {
+	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
+}
diff --git a/pkg/migrations/op_set_unique_test.go b/pkg/migrations/op_set_unique_test.go
index d2664b99..ed8e0b5b 100644
--- a/pkg/migrations/op_set_unique_test.go
+++ b/pkg/migrations/op_set_unique_test.go
@@ -47,6 +47,7 @@ func TestSetColumnsUnique(t *testing.T) {
 				Name: "02_set_unique",
 				Operations: migrations.Operations{
 					&migrations.OpSetUnique{
+						Name:    "reviews_username_product_unique",
 						Table:   "reviews",
 						Columns: []string{"username", "product"},
 					},
@@ -55,8 +56,7 @@ func TestSetColumnsUnique(t *testing.T) {
 		},
 		afterStart: func(t *testing.T, db *sql.DB) {
 			// The unique index has been created on the underlying table.
-			idxName := migrations.IndexName("reviews", []string{"username", "product"})
-			IndexMustExist(t, db, "public", "reviews", idxName)
+			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{
@@ -76,13 +76,11 @@ func TestSetColumnsUnique(t *testing.T) {
 		},
 		afterRollback: func(t *testing.T, db *sql.DB) {
 			// The unique index has been dropped from the the underlying table.
-			idxName := migrations.IndexName("reviews", []string{"username", "product"})
-			IndexMustNotExist(t, db, "public", "reviews", idxName)
+			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.
-			constraintName := migrations.UniqueConstraintName("reviews", []string{"username", "product"})
-			ConstraintMustExist(t, db, "public", "reviews", constraintName)
+			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{