Repository: ariga/atlas
Pull request #2667 (2024-04-08 13:23:10)

sql/postgres: initial support for unique constraints


Patch:
diff --git a/.golangci.yml b/.golangci.yml
index 9d9e0011..ab0a0123 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -6,6 +6,7 @@ issues:
     - EXC0012
   exclude:
     - G601
+    - G404
   exclude-rules:
     - path: _test\.go
       linters:
diff --git a/doc/md/atlas-schema/hcl.mdx b/doc/md/atlas-schema/hcl.mdx
index e43183f1..d6cf4b40 100644
--- a/doc/md/atlas-schema/hcl.mdx
+++ b/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
 ```
diff --git a/schemahcl/spec.go b/schemahcl/spec.go
index eabc9a4f..a2b0b88b 100644
--- a/schemahcl/spec.go
+++ b/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 {
diff --git a/sql/internal/sqlx/diff.go b/sql/internal/sqlx/diff.go
index 9fb308ed..edde2ac9 100644
--- a/sql/internal/sqlx/diff.go
+++ b/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.
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index 8d9a8f48..d339add6 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -189,6 +189,11 @@ func (*diff) IndexAttrChanged(from, to []schema.Attr) bool {
 	if indexNullsDistinct(to) != indexNullsDistinct(from) {
 		return true
 	}
+	_, u1 := uniqueConst(from)
+	_, u2 := uniqueConst(to)
+	if u1 != u2 {
+		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
@@ -496,6 +501,16 @@ func indexNullsDistinct(attrs []schema.Attr) bool {
 	return true
 }
 
+// uniqueConst returns the first unique constraint from the given attributes.
+func uniqueConst(attrs []schema.Attr) (*Constraint, bool) {
+	for _, a := range attrs {
+		if c, ok := a.(*Constraint); ok && c.IsUnique() {
+			return c, true
+		}
+	}
+	return nil, false
+}
+
 func trimCast(s string) string {
 	i := strings.LastIndex(s, "::")
 	if i == -1 {
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 805d4dd2..0ca9231b 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -7,9 +7,9 @@ package postgres
 import (
 	"context"
 	"database/sql"
-	"errors"
 	"fmt"
 	"hash/fnv"
+	"math/rand"
 	"net/url"
 	"strconv"
 	"time"
@@ -337,38 +337,38 @@ func (*Driver) ScanStmts(input string) ([]*migrate.Stmt, error) {
 	}).Scan(input)
 }
 
+// Use pg_try_advisory_lock to avoid deadlocks between multiple executions of Atlas (commonly tests).
+// The common case is as follows: a process (P1) of Atlas takes a lock, and another process (P2) of
+// Atlas waits for the lock. Now if P1 execute "CREATE INDEX CONCURRENTLY" (either in apply or diff),
+// the command waits all active transactions that can potentially changed the index to be finished.
+// P2 can be executed in a transaction block (opened explicitly by Atlas), or a single statement tx
+// also known as "autocommit mode". Read more: https://www.postgresql.org/docs/current/sql-begin.html.
 func acquire(ctx context.Context, conn schema.ExecQuerier, id uint32, timeout time.Duration) error {
-	switch {
-	// With timeout (context-based).
-	case timeout > 0:
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, timeout)
-		defer cancel()
-		fallthrough
-	// Infinite timeout.
-	case timeout < 0:
-		rows, err := conn.QueryContext(ctx, "SELECT pg_advisory_lock($1)", id)
-		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
-			err = schema.ErrLocked
-		}
-		if err != nil {
-			return err
-		}
-		return rows.Close()
-	// No timeout.
-	default:
+	var (
+		inter = 25
+		start = time.Now()
+	)
+	for {
 		rows, err := conn.QueryContext(ctx, "SELECT pg_try_advisory_lock($1)", id)
 		if err != nil {
 			return err
 		}
-		acquired, err := sqlx.ScanNullBool(rows)
-		if err != nil {
+		switch acquired, err := sqlx.ScanNullBool(rows); {
+		case err != nil:
 			return err
-		}
-		if !acquired.Bool {
+		case acquired.Bool:
+			return nil
+		case time.Since(start) > timeout:
 			return schema.ErrLocked
+		default:
+			if err := rows.Close(); err != nil {
+				return err
+			}
+			// 25ms~50ms, 50ms~100ms, ..., 800ms~1.6s, 1s~2s.
+			d := min(time.Duration(inter)*time.Millisecond, time.Second)
+			time.Sleep(d + time.Duration(rand.Intn(int(d))))
+			inter += inter
 		}
-		return nil
 	}
 }
 
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index 867c3f5d..08609eca 100644
--- a/sql/postgres/driver_oss.go
+++ b/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.
+}
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index f35ae6d5..7dd18a4f 100644
--- a/sql/postgres/inspect.go
+++ b/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}
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 9ab51c19..b745fb93 100644
--- a/sql/postgres/migrate.go
+++ b/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 _, ok := uniqueConst(idx.Attrs); ok {
+				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 _, ok := uniqueConst(idx.Attrs); !ok {
+			// 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 _, ok := uniqueConst(change.I.Attrs); ok {
+				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 _, ok := uniqueConst(change.I.Attrs); ok {
 				alter = append(alter, change)
 			} else {
 				dropI = append(dropI, change)
@@ -366,9 +380,21 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 					continue
 				}
 			}
+			if addU, ok := indexToUnique(change); ok {
+				alter = append(alter, addU)
+				continue
+			}
 			// Index modification requires rebuilding the index.
-			addI = append(addI, &schema.AddIndex{I: change.To})
-			dropI = append(dropI, &schema.DropIndex{I: change.From})
+			if _, ok := uniqueConst(change.From.Attrs); ok {
+				alter = append(alter, &schema.DropIndex{I: change.From})
+			} else {
+				dropI = append(dropI, &schema.DropIndex{I: change.From})
+			}
+			if _, ok := uniqueConst(change.To.Attrs); ok {
+				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 {
+	c, ok := uniqueConst(idx.Attrs)
+	if !ok {
+		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
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 6fcb61b6..c035f8c5 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -45,7 +45,7 @@ type (
 		schemahcl.DefaultExtension
 	}
 
-	// Domain holds a specification for a domain type.
+	// domain holds a specification for a domain type.
 	domain struct {
 		Name      string           `spec:",name"`
 		Qualifier string           `spec:",qualifier"`
@@ -282,6 +282,9 @@ func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, er
 	if err != nil {
 		return nil, err
 	}
+	if err := convertUniques(spec.Extra, t); err != nil {
+		return nil, err
+	}
 	if err := convertPartition(spec.Extra, t); err != nil {
 		return nil, err
 	}
@@ -310,6 +313,24 @@ func convertView(spec *sqlspec.View, parent *schema.Schema) (*schema.View, error
 	return v, nil
 }
 
+// convertUniques converts the unique constraints into indexes.
+func convertUniques(spec schemahcl.Resource, t *schema.Table) error {
+	rs := spec.Resources("unique")
+	for _, r := range rs {
+		var sx sqlspec.Index
+		if err := r.As(&sx); err != nil {
+			return fmt.Errorf("parse %s.unique constraint: %w", t.Name, err)
+		}
+		idx, err := convertIndex(&sx, t)
+		if err != nil {
+			return err
+		}
+		idx.SetUnique(true).AddAttrs(UniqueConstraint(sx.Name))
+		t.AddIndexes(idx)
+	}
+	return nil
+}
+
 // convertPartition converts and appends the partition block into the table attributes if exists.
 func convertPartition(spec schemahcl.Resource, table *schema.Table) error {
 	r, ok := spec.Resource("partition")
@@ -613,9 +634,9 @@ func schemaSpec(s *schema.Schema) (*doc, []*schema.Trigger, error) {
 }
 
 // tableSpec converts from a concrete Postgres sqlspec.Table to a schema.Table.
-func tableSpec(table *schema.Table) (*sqlspec.Table, error) {
+func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 	spec, err := specutil.FromTable(
-		table,
+		t,
 		tableColumnSpec,
 		pkSpec,
 		indexSpec,
@@ -625,7 +646,26 @@ func tableSpec(table *schema.Table) (*sqlspec.Table, error) {
 	if err != nil {
 		return nil, err
 	}
-	if p := (Partition{}); sqlx.Has(table.Attrs, &p) {
+	idxs := make([]*sqlspec.Index, 0, len(spec.Indexes))
+	for i, idx1 := range spec.Indexes {
+		if len(t.Indexes) <= i || t.Indexes[i].Name != idx1.Name {
+			return nil, fmt.Errorf("unexpected spec index %q was not found in table %q", idx1.Name, t.Name)
+		}
+		if c := (&Constraint{}); sqlx.Has(t.Indexes[i].Attrs, c) && c.IsUnique() && c.N != "" {
+			spec.Extra.Children = append(spec.Extra.Children, &schemahcl.Resource{
+				Type: "unique",
+				Name: c.N,
+				Attrs: append(
+					[]*schemahcl.Attr{schemahcl.RefsAttr("columns", idx1.Columns...)},
+					idx1.Extra.Attrs...,
+				),
+			})
+		} else {
+			idxs = append(idxs, idx1)
+		}
+	}
+	spec.Indexes = idxs
+	if p := (Partition{}); sqlx.Has(t.Attrs, &p) {
 		spec.Extra.Children = append(spec.Extra.Children, fromPartition(p))
 	}
 	return spec, nil


Test patch:
diff --git a/internal/integration/script_test.go b/internal/integration/script_test.go
index e1ccac5e..820eb56c 100644
--- a/internal/integration/script_test.go
+++ b/internal/integration/script_test.go
@@ -14,6 +14,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
+	"slices"
 	"strconv"
 	"strings"
 	"testing"
@@ -423,6 +424,11 @@ func cmdCLI(ts *testscript.TestScript, neg bool, args []string, dbURL, devURL, c
 	for i, arg := range args {
 		args[i] = r.Replace(arg)
 	}
+	// Whenever a migrate diff/apply command is executed, increase the default lock
+	// timeout to 30s, as the default (10s) for synchronizing all our tests.
+	if len(args) > 1 && args[0] == "migrate" && (args[1] == "diff" || args[1] == "apply") && !slices.Contains(args, "--lock-timeout") {
+		args = append(args, "--lock-timeout", "30s")
+	}
 	switch l := len(args); {
 	// If command was run with a unix redirect-like suffix.
 	case l > 1 && args[l-2] == ">":
diff --git a/internal/integration/testdata/postgres/index-unique-constraint.txt b/internal/integration/testdata/postgres/index-unique-constraint.txt
index 897ba132..0e7be846 100644
--- a/internal/integration/testdata/postgres/index-unique-constraint.txt
+++ b/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]
   }
 }
diff --git a/sql/postgres/driver_test.go b/sql/postgres/driver_test.go
index 9314a3c8..4a0025bd 100644
--- a/sql/postgres/driver_test.go
+++ b/sql/postgres/driver_test.go
@@ -22,20 +22,44 @@ func TestDriver_LockAcquired(t *testing.T) {
 	db, m, err := sqlmock.New()
 	require.NoError(t, err)
 	name, hash := "name", 797654004
-	m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
-		WithArgs(hash).
-		WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_lock"}).AddRow(1)).
-		RowsWillBeClosed()
-	m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
-		WithArgs(hash).
-		WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(1)).
-		RowsWillBeClosed()
 
-	d := &Driver{conn: &conn{ExecQuerier: db}}
-	unlock, err := d.Lock(context.Background(), name, 0)
-	require.NoError(t, err)
-	require.NoError(t, unlock())
-	require.NoError(t, m.ExpectationsWereMet())
+	t.Run("NoTimeout", func(t *testing.T) {
+		m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_lock"}).AddRow(1)).
+			RowsWillBeClosed()
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(1)).
+			RowsWillBeClosed()
+
+		d := &Driver{conn: &conn{ExecQuerier: db}}
+		unlock, err := d.Lock(context.Background(), name, 0)
+		require.NoError(t, err)
+		require.NoError(t, unlock())
+		require.NoError(t, m.ExpectationsWereMet())
+	})
+
+	t.Run("WithTimeout", func(t *testing.T) {
+		m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_lock"}).AddRow(0)).
+			RowsWillBeClosed()
+		m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_lock"}).AddRow(1)).
+			RowsWillBeClosed()
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(1)).
+			RowsWillBeClosed()
+
+		d := &Driver{conn: &conn{ExecQuerier: db}}
+		unlock, err := d.Lock(context.Background(), name, time.Second)
+		require.NoError(t, err)
+		require.NoError(t, unlock())
+		require.NoError(t, m.ExpectationsWereMet())
+	})
 }
 
 func TestDriver_LockError(t *testing.T) {
@@ -44,18 +68,8 @@ func TestDriver_LockError(t *testing.T) {
 	d := &Driver{conn: &conn{ExecQuerier: db}}
 	name, hash := "migrate", 979249972
 
-	t.Run("Timeout", func(t *testing.T) {
-		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_lock($1)")).
-			WithArgs(hash).
-			WillReturnError(context.DeadlineExceeded).
-			RowsWillBeClosed()
-		unlock, err := d.Lock(context.Background(), name, time.Minute)
-		require.Equal(t, schema.ErrLocked, err)
-		require.Nil(t, unlock)
-	})
-
 	t.Run("Internal", func(t *testing.T) {
-		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_lock($1)")).
+		m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
 			WithArgs(hash).
 			WillReturnError(io.EOF).
 			RowsWillBeClosed()
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index eedc3e14..afb48400 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -2326,3 +2326,92 @@ schema "s3" {
 }
 `, string(buf))
 }
+
+func TestMarshalSpec_UniqueConstraint(t *testing.T) {
+	s := schema.New("public").
+		AddTables(
+			schema.NewTable("t1").
+				AddColumns(
+					schema.NewIntColumn("id", "int"),
+					schema.NewIntColumn("uid", "int"),
+				).
+				AddIndexes(
+					schema.NewUniqueIndex("u1").AddColumns(schema.NewColumn("id")).AddAttrs(UniqueConstraint("u1")),
+					schema.NewUniqueIndex("u2").AddColumns(schema.NewColumn("id")).AddAttrs(UniqueConstraint("u2"), &IndexNullsDistinct{}),
+					schema.NewUniqueIndex("u3").AddColumns(schema.NewColumn("id")).AddAttrs(UniqueConstraint("u3"), &IndexInclude{Columns: []*schema.Column{schema.NewColumn("uid")}}),
+				),
+		)
+	buf, err := MarshalHCL(s)
+	require.NoError(t, err)
+	require.Equal(t, `table "t1" {
+  schema = schema.public
+  column "id" {
+    null = false
+    type = int
+  }
+  column "uid" {
+    null = false
+    type = int
+  }
+  unique "u1" {
+    columns = [column.id]
+  }
+  unique "u2" {
+    columns        = [column.id]
+    nulls_distinct = false
+  }
+  unique "u3" {
+    columns = [column.id]
+    include = [column.uid]
+  }
+}
+schema "public" {
+}
+`,
+		string(buf))
+}
+
+func TestUnmarshalSpec_UniqueConstraint(t *testing.T) {
+	var s schema.Schema
+	require.NoError(t, EvalHCLBytes([]byte(`table "t1" {
+  schema = schema.public
+  column "id" {
+    null = false
+    type = int
+  }
+  column "uid" {
+    null = false
+    type = int
+  }
+  unique "u1" {
+    columns = [column.id]
+  }
+  unique "u2" {
+    columns        = [column.id]
+    nulls_distinct = false
+  }
+  unique "u3" {
+    columns = [column.id]
+    include = [column.uid]
+  }
+}
+schema "public" {
+}
+`), &s, nil))
+	require.Len(t, s.Tables, 1)
+	require.Len(t, s.Tables[0].Indexes, 3)
+	// u1
+	u1 := s.Tables[0].Indexes[0]
+	require.Equal(t, "u1", u1.Name)
+	require.Equal(t, UniqueConstraint("u1"), u1.Attrs[0].(*Constraint))
+	// u2
+	u2 := s.Tables[0].Indexes[1]
+	require.Equal(t, "u2", u2.Name)
+	require.Equal(t, &IndexNullsDistinct{}, u2.Attrs[0])
+	require.Equal(t, UniqueConstraint("u2"), u2.Attrs[1].(*Constraint))
+	// u3
+	u3 := s.Tables[0].Indexes[2]
+	require.Equal(t, "u3", u3.Name)
+	require.Equal(t, &IndexInclude{Columns: []*schema.Column{s.Tables[0].Columns[1]}}, u3.Attrs[0])
+	require.Equal(t, UniqueConstraint("u3"), u3.Attrs[1].(*Constraint))
+}