sql/postgres: initial support for unique constraints

doc/md/atlas-schema/hcl.mdx

@@ -605,40 +605,44 @@ Indexes are child resources of a `table`, and it defines an index on the table.
 #### Example
 
 ```hcl
+# Columns only.
 index "idx_name" {
-    columns = [
-      column.name
-    ]
-    unique = true
+  unique = true
+  columns = [column.name]
 }
 
+# Columns and order.
 index "idx_name" {
-    on {
-        column = column.rank
-    }
-    on {
-        column = column.score
-        desc = true
-    }
-    unique = true
+  unique = true
+  on {
+    column = column.rank
+  }
+  on {
+    column = column.score
+    desc = true
+  }
 }
 
+# Custom index type.
 index "idx_name" {
   type = GIN
   columns = [column.data]
 }
 
+# Control storage options.
 index "idx_range" {
   type = BRIN
   columns = [column.range]
   page_per_range = 128
 }
 
+# Include non-key columns.
 index "idx_include" {
   columns = [column.range]
   include = [column.version]
 }
 
+# Define operator class.
 index "idx_operator_class" {
   type = GIN
   on {
@@ -647,12 +651,14 @@ index "idx_operator_class" {
   }
 }
 
+# Full-text index with ngram parser.
 index "index_parser" {
   type    = FULLTEXT
   columns = [column.text]
   parser  = ngram
 }
 
+# Postgres-specific NULLS [NOT] DISTINCT option.
 index "index_nulls_not_distinct" {
   unique         = true
   columns        = [column.text]
@@ -744,6 +750,33 @@ table "users" {
 }
 ```
 
+### Unique Constraints
+
+The `unique` block allows defining a [unique constraint](https://www.postgresql.org/docs/current/ddl-constraints.html#DDL-CONSTRAINTS-UNIQUE-CONSTRAINTS)
+supported by PostgreSQL:
+
+```hcl
+# Columns only.
+unique "name" {
+  columns = [column.name]
+}
+
+# Include non-key columns.
+unique "name_include_version" {
+  columns = [column.name]
+  include = [column.version]
+}
+```
+
+:::info <a class="sticky-anchor" id={"adding-unique-constraints-concurrently"} href={"#adding-unique-constraints-concurrently"}>Adding unique constraints concurrently</a>
+In order to add a unique constraint in non-blocking mode, the index supporting the constraint needs to be created
+concurrently first and then converted to a unique constraint. To achieve this, follow the steps below:
+1. Define a unique `index` block on the desired table.
+2. Ensure a [Diff Policy](/versioned/diff#diff-policy) is used to instruct Atlas to create the index concurrently.
+3. Apply the migration and ensure the index was created.
+4. Replace the `index` block with a `unique` block to create a new unique constraint using the existing index.
+:::
+
 ## Trigger
 
 :::info LOGIN REQUIRED
@@ -1460,8 +1493,8 @@ table "t2" {
 
 ## Extension
 
-:::info BETA FEATURE
-Extensions are currently in beta and available to logged-in users only. To use this feature, run:
+:::info LOGIN REQUIRED
+Extensions are currently available to logged-in users only. To use this feature, run:
 ```
 atlas login
 ```

internal/integration/testdata/postgres/index-unique-constraint.txt

@@ -24,12 +24,10 @@ table "users" {
     null = true
     type = text
   }
-  index "users_name_last_key" {
-    unique  = true
+  unique "users_name_last_key" {
     columns = [column.name, column.last]
   }
-  index "users_nickname_key" {
-    unique  = true
+  unique "users_nickname_key" {
     columns = [column.nickname]
   }
 }
@@ -51,8 +49,7 @@ table "users" {
     null = true
     type = text
   }
-  index "users_name_last_key" {
-    unique  = true
+  unique "users_name_last_key" {
     columns = [column.name, column.last]
   }
 }

schemahcl/spec.go

@@ -392,6 +392,20 @@ func (r *Resource) Resource(t string) (*Resource, bool) {
 	return nil, false
 }
 
+// Resources returns all child Resources by its type.
+func (r *Resource) Resources(t string) []*Resource {
+	if r == nil {
+		return nil
+	}
+	var rs []*Resource
+	for i := range r.Children {
+		if r.Children[i].Type == t {
+			rs = append(rs, r.Children[i])
+		}
+	}
+	return rs
+}
+
 // Attr returns the Attr by the provided name and reports whether it was found.
 func (r *Resource) Attr(name string) (*Attr, bool) {
 	if at, ok := attrVal(r.Attrs, name); ok {

sql/internal/sqlx/diff.go

@@ -650,6 +650,17 @@ var (
 	exprsType   = reflect.TypeOf(([]schema.Expr)(nil))
 )
 
+// AttrOr returns the first attribute of the given type,
+// or the given default value.
+func AttrOr[T schema.Attr](attrs []schema.Attr, t T) T {
+	for _, attr := range attrs {
+		if a, ok := attr.(T); ok {
+			return a
+		}
+	}
+	return t
+}
+
 // Has finds the first element in the elements list that
 // matches target, and if so, sets target to that attribute
 // value and returns true.

sql/postgres/diff.go

@@ -189,6 +189,9 @@ func (*diff) IndexAttrChanged(from, to []schema.Attr) bool {
 	if indexNullsDistinct(to) != indexNullsDistinct(from) {
 		return true
 	}
+	if sqlx.AttrOr(from, &Constraint{}).IsUnique() != sqlx.AttrOr(to, &Constraint{}).IsUnique() {
+		return true
+	}
 	var p1, p2 IndexPredicate
 	if sqlx.Has(from, &p1) != sqlx.Has(to, &p2) || (p1.P != p2.P && p1.P != sqlx.MayWrap(p2.P)) {
 		return true

sql/postgres/driver_oss.go

@@ -335,3 +335,7 @@ func convertTypes(d *doc, r *schema.Realm) error {
 	}
 	return nil
 }
+
+func indexToUnique(*schema.ModifyIndex) (*AddUniqueConstraint, bool) {
+	return nil, false // unimplemented.
+}

sql/postgres/inspect.go

@@ -1169,6 +1169,11 @@ func (o *ReferenceOption) Scan(v any) error {
 // IsUnique reports if the type is unique constraint.
 func (c Constraint) IsUnique() bool { return strings.ToLower(c.T) == "u" }
 
+// UniqueConstraint returns constraint with type "u".
+func UniqueConstraint(name string) *Constraint {
+	return &Constraint{T: "u", N: name}
+}
+
 // IntegerType returns the underlying integer type this serial type represents.
 func (s *SerialType) IntegerType() *schema.IntegerType {
 	t := &schema.IntegerType{T: TypeInteger}

sql/postgres/migrate.go

@@ -237,6 +237,14 @@ func (s *state) addTable(add *schema.AddTable) error {
 				errs = append(errs, err.Error())
 			}
 		}
+		for _, idx := range add.T.Indexes {
+			if sqlx.AttrOr(idx.Attrs, &Constraint{}).IsUnique() {
+				b.Comma().NL()
+				if err := s.unique(b, idx); err != nil {
+					errs = append(errs, err.Error())
+				}
+			}
+		}
 		if len(add.T.ForeignKeys) > 0 {
 			b.Comma()
 			s.fks(b, add.T.ForeignKeys...)
@@ -265,9 +273,11 @@ func (s *state) addTable(add *schema.AddTable) error {
 		Reverse: s.Build("DROP TABLE").Table(add.T).String(),
 	})
 	for _, idx := range add.T.Indexes {
-		// Indexes do not need to be created concurrently on new tables.
-		if err := s.addIndexes(add, add.T, &schema.AddIndex{I: idx}); err != nil {
-			return err
+		if !sqlx.AttrOr(idx.Attrs, &Constraint{}).IsUnique() {
+			// Indexes do not need to be created concurrently on new tables.
+			if err := s.addIndexes(add, add.T, &schema.AddIndex{I: idx}); err != nil {
+				return err
+			}
 		}
 	}
 	s.addComments(add, add.T)
@@ -336,11 +346,15 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 			if c := (schema.Comment{}); sqlx.Has(change.I.Attrs, &c) {
 				changes = append(changes, s.indexComment(modify, modify.T, change.I, c.Text, ""))
 			}
-			addI = append(addI, change)
+			if c := (Constraint{}); sqlx.Has(change.I.Attrs, &c) && c.IsUnique() {
+				alter = append(alter, change)
+			} else {
+				addI = append(addI, change)
+			}
 		case *schema.DropIndex:
 			// Unlike DROP INDEX statements that are executed separately,
 			// DROP CONSTRAINT are added to the ALTER TABLE statement below.
-			if isUniqueConstraint(change.I) {
+			if c := (Constraint{}); sqlx.Has(change.I.Attrs, &c) && c.IsUnique() {
 				alter = append(alter, change)
 			} else {
 				dropI = append(dropI, change)
@@ -366,9 +380,21 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 					continue
 				}
 			}
-			// Index modification requires rebuilding the index.
-			addI = append(addI, &schema.AddIndex{I: change.To})
-			dropI = append(dropI, &schema.DropIndex{I: change.From})
+			if addU, ok := indexToUnique(change); ok {
+				alter = append(alter, addU)
+				continue
+			}
+			// Index modification requires rebuilding the index or the constraint.
+			if sqlx.AttrOr(change.From.Attrs, &Constraint{}).IsUnique() {
+				alter = append(alter, &schema.DropIndex{I: change.From})
+			} else {
+				dropI = append(dropI, &schema.DropIndex{I: change.From})
+			}
+			if sqlx.AttrOr(change.To.Attrs, &Constraint{}).IsUnique() {
+				alter = append(alter, &schema.AddIndex{I: change.To})
+			} else {
+				addI = append(addI, &schema.AddIndex{I: change.To})
+			}
 		case *schema.RenameIndex:
 			changes = append(changes, &migrate.Change{
 				Source:  change,
@@ -432,6 +458,17 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 	return nil
 }
 
+type (
+	// AddUniqueConstraint to the table using the given index. Note, if the index
+	// name does not match the constraint name, PostgreSQL implicitly renames it to
+	// the constraint name.
+	AddUniqueConstraint struct {
+		schema.Change
+		Name  string        // Name of the constraint.
+		Using *schema.Index // Index to use for the constraint.
+	}
+)
+
 // alterTable modifies the given table by executing on it a list of changes in one SQL statement.
 func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 	var (
@@ -467,13 +504,22 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 			case *schema.DropColumn:
 				b.P("DROP COLUMN").Ident(change.C.Name)
 				reverse = append(reverse, &schema.AddColumn{C: change.C})
+			case *AddUniqueConstraint:
+				b.P("ADD CONSTRAINT").Ident(change.Name).P("UNIQUE USING INDEX").Ident(change.Using.Name)
+				drop := change.Using
+				if drop.Name != change.Name {
+					drop = sqlx.P(*change.Using)
+					drop.Name = change.Name
+				}
+				// Translated to the DROP CONSTRAINT below,
+				// which drops the index as well.
+				reverse = append(reverse, &schema.DropIndex{I: drop})
 			case *schema.AddIndex:
-				// Skip reversing this operation as it is the inverse of
-				// the operation above and should not be used besides this.
-				b.P("ADD CONSTRAINT").Ident(change.I.Name).P("UNIQUE")
-				if err := s.indexParts(b, change.I); err != nil {
+				b.P("ADD")
+				if err := s.unique(b, change.I); err != nil {
 					return err
 				}
+				reverse = append(reverse, &schema.DropIndex{I: change.I})
 			case *schema.DropIndex:
 				b.P("DROP CONSTRAINT").Ident(change.I.Name)
 				reverse = append(reverse, &schema.AddIndex{I: change.I})
@@ -1083,6 +1129,19 @@ func (s *state) fks(b *sqlx.Builder, fks ...*schema.ForeignKey) {
 	})
 }
 
+func (s *state) unique(b *sqlx.Builder, idx *schema.Index) error {
+	var c Constraint
+	if !sqlx.Has(idx.Attrs, &c) {
+		return fmt.Errorf("index %q is not a unique constraint", idx.Name)
+	}
+	name := c.N
+	if name == "" {
+		name = idx.Name
+	}
+	b.P("CONSTRAINT").Ident(name).P("UNIQUE")
+	return s.index(b, idx)
+}
+
 func (s *state) append(c ...*migrate.Change) {
 	s.Changes = append(s.Changes, c...)
 }
@@ -1166,41 +1225,6 @@ func check(b *sqlx.Builder, c *schema.Check) {
 	}
 }
 
-// isUniqueConstraint reports if the index is a valid UNIQUE constraint.
-func isUniqueConstraint(i *schema.Index) bool {
-	hasC := func() bool {
-		for _, a := range i.Attrs {
-			if c, ok := a.(*Constraint); ok && c.IsUnique() {
-				return true
-			}
-		}
-		return false
-	}()
-	if !hasC || !i.Unique {
-		return false
-	}
-	// UNIQUE constraint cannot use functional indexes,
-	// and all its parts must have the default sort ordering.
-	for _, p := range i.Parts {
-		if p.X != nil || p.Desc {
-			return false
-		}
-	}
-	for _, a := range i.Attrs {
-		switch a := a.(type) {
-		// UNIQUE constraints must have BTREE type indexes.
-		case *IndexType:
-			if strings.ToUpper(a.T) != IndexTypeBTree {
-				return false
-			}
-		// Partial indexes are not allowed.
-		case *IndexPredicate:
-			return false
-		}
-	}
-	return true
-}
-
 func quote(s string) string {
 	if sqlx.IsQuoted(s, '\'') {
 		return s