fix: preserve constraint INCLUDE columns (#673)

internal/diff/constraint.go

@@ -26,7 +26,7 @@ func generateConstraintSQL(constraint *ir.Constraint, targetSchema string, quali
 		if constraint.IsTemporal && len(cols) > 0 {
 			cols[len(cols)-1] = cols[len(cols)-1] + " WITHOUT OVERLAPS"
 		}
-		return fmt.Sprintf("CONSTRAINT %s PRIMARY KEY (%s)%s", ir.QuoteIdentifier(constraint.Name), strings.Join(cols, ", "), deferrableClause(constraint))
+		return fmt.Sprintf("CONSTRAINT %s PRIMARY KEY (%s)%s%s", ir.QuoteIdentifier(constraint.Name), strings.Join(cols, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 	case ir.ConstraintTypeUnique:
 		// Always include CONSTRAINT name to be explicit and consistent
 		cols := getColumnNames(constraint.Columns)
@@ -37,7 +37,7 @@ func generateConstraintSQL(constraint *ir.Constraint, targetSchema string, quali
 		if constraint.NullsNotDistinct {
 			modifier = " NULLS NOT DISTINCT"
 		}
-		return fmt.Sprintf("CONSTRAINT %s UNIQUE%s (%s)%s", ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(cols, ", "), deferrableClause(constraint))
+		return fmt.Sprintf("CONSTRAINT %s UNIQUE%s (%s)%s%s", ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(cols, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 	case ir.ConstraintTypeForeignKey:
 		// Always include CONSTRAINT name to preserve explicit FK names
 		// Use QualifyEntityNameWithQuotes to add schema qualifier when referencing tables in other schemas
@@ -91,6 +91,18 @@ func generateConstraintSQL(constraint *ir.Constraint, targetSchema string, quali
 	}
 }
 
+// constraintIncludeClause renders the ordered non-key columns of a PK or UNIQUE.
+func constraintIncludeClause(constraint *ir.Constraint) string {
+	if len(constraint.IncludeColumns) == 0 {
+		return ""
+	}
+	columns := make([]string, len(constraint.IncludeColumns))
+	for i, column := range constraint.IncludeColumns {
+		columns[i] = ir.QuoteIdentifier(column)
+	}
+	return " INCLUDE (" + strings.Join(columns, ", ") + ")"
+}
+
 // deferrableClause returns the " DEFERRABLE [INITIALLY DEFERRED]" suffix for a
 // constraint, or "" when it is not deferrable. PRIMARY KEY, UNIQUE, and FOREIGN
 // KEY constraints all support this; EXCLUDE constraints carry the clause inside
@@ -234,6 +246,10 @@ func constraintsEqual(old, new *ir.Constraint) bool {
 		}
 	}
 
+	if !slices.Equal(old.IncludeColumns, new.IncludeColumns) {
+		return false
+	}
+
 	// Compare referenced columns
 	if len(old.ReferencedColumns) != len(new.ReferencedColumns) {
 		return false

internal/diff/constraint_include_test.go

@@ -0,0 +1,184 @@
+package diff
+
+import (
+	"context"
+	"database/sql"
+	"testing"
+
+	"github.com/pgplex/pgschema/ir"
+	"github.com/pgplex/pgschema/ir/queries"
+	"github.com/pgplex/pgschema/testutil"
+	"github.com/stretchr/testify/require"
+)
+
+// Compare native catalog definitions rather than generating the expected SQL
+// with the same renderer that produces the migration.
+func constraintIncludeCatalog(t *testing.T, db *sql.DB) []string {
+	t.Helper()
+	rows, err := db.Query(`SELECT t.relname || ':' || c.conname || ':' || pg_get_constraintdef(c.oid) || ':' || pg_get_indexdef(c.conindid)
+ FROM pg_constraint c JOIN pg_class t ON t.oid = c.conrelid
+ WHERE t.relnamespace = 'public'::regnamespace AND c.contype IN ('p', 'u') ORDER BY t.relname, c.conname`)
+	require.NoError(t, err)
+	defer rows.Close()
+	var definitions []string
+	for rows.Next() {
+		var definition string
+		require.NoError(t, rows.Scan(&definition))
+		definitions = append(definitions, definition)
+	}
+	require.NoError(t, rows.Err())
+	return definitions
+}
+
+func TestConstraintIncludeRoundTrip(t *testing.T) {
+	const source = `
+ CREATE TABLE ordinary (id int, a int, payload text, "select" text,
+ CONSTRAINT ordinary_pk PRIMARY KEY (a, id) INCLUDE ("select", payload) DEFERRABLE INITIALLY DEFERRED,
+ CONSTRAINT ordinary_unique UNIQUE (id, a) INCLUDE (payload, "select") DEFERRABLE);
+ CREATE TABLE events (id int, a int, payload text, "select" text,
+ CONSTRAINT root_unique UNIQUE (id, a) INCLUDE (payload, "select"),
+ CONSTRAINT root_pk PRIMARY KEY (a, id) INCLUDE ("select", payload)) PARTITION BY RANGE (id);
+ CREATE TABLE events_low PARTITION OF events FOR VALUES FROM (0) TO (10) PARTITION BY RANGE (a);
+ CREATE TABLE events_leaf PARTITION OF events_low FOR VALUES FROM (0) TO (10);
+ DO $$ DECLARE old_name text; BEGIN
+ SELECT conname INTO old_name FROM pg_constraint WHERE conrelid = 'events_low'::regclass AND contype = 'u';
+ EXECUTE format('ALTER TABLE events_low RENAME CONSTRAINT %I TO "renamed middle unique"', old_name);
+ SELECT conname INTO old_name FROM pg_constraint WHERE conrelid = 'events_leaf'::regclass AND contype = 'p';
+ EXECUTE format('ALTER TABLE events_leaf RENAME CONSTRAINT %I TO "renamed leaf pk"', old_name);
+ END $$;
+ `
+	db, _, _, _, _, _ := testutil.ConnectToPostgres(t, sharedTestPostgres)
+	defer db.Close()
+	original := testutil.ParseSQLToIR(t, sharedTestPostgres, source, "public")
+	rows, err := queries.New(db).GetConstraints(context.Background())
+	require.NoError(t, err)
+	found := false
+	for _, row := range rows {
+		if row.ConstraintName == "ordinary_pk" {
+			require.Equal(t, []string{"select", "payload"}, row.IncludeColumns)
+			found = true
+		}
+	}
+	require.True(t, found)
+	require.Equal(t, constraintColumns("a", "id"), original.Schemas["public"].Tables["ordinary"].Constraints["ordinary_pk"].Columns)
+	expected := constraintIncludeCatalog(t, db)
+	for _, mode := range []string{"dump", "initial migration"} {
+		t.Run(mode, func(t *testing.T) {
+			var statements []Diff
+			if mode == "dump" {
+				var err error
+				statements, err = GenerateDump(original, "public", false)
+				require.NoError(t, err)
+			} else {
+				statements = GenerateMigration(&ir.IR{}, original, "public")
+			}
+			sqlText := buildSQLFromSteps(statements)
+			require.Contains(t, sqlText, `UNIQUE (id, a) INCLUDE (payload, "select")`)
+			replayed := testutil.ParseSQLToIR(t, sharedTestPostgres, sqlText, "public")
+			require.Equal(t, expected, constraintIncludeCatalog(t, db))
+			require.Empty(t, GenerateMigration(replayed, original, "public"))
+			// Included values must not enlarge the uniqueness key.
+			_, err := db.Exec(`INSERT INTO events VALUES (1, 1, 'first', 'x'); INSERT INTO events VALUES (1, 1, 'second', 'y')`)
+			require.Error(t, err)
+		})
+	}
+}
+
+func TestConstraintIncludeMigration(t *testing.T) {
+	for _, topology := range []string{"ordinary", "partitioned"} {
+		for _, kind := range []string{"PRIMARY KEY", "UNIQUE"} {
+			for _, transition := range []struct{ name, old, new string }{
+				{"add", "", ` INCLUDE (payload, "select")`},
+				{"reorder", ` INCLUDE (payload, "select")`, ` INCLUDE ("select", payload)`},
+				{"change", ` INCLUDE (payload)`, ` INCLUDE ("select")`},
+				{"remove", ` INCLUDE (payload)`, ""},
+			} {
+				t.Run(topology+"/"+kind+"/"+transition.name, func(t *testing.T) {
+					schema := func(include string) string {
+						partition := ""
+						if topology == "partitioned" {
+							partition = ` PARTITION BY RANGE (id); CREATE TABLE leaf PARTITION OF parent FOR VALUES FROM (0) TO (10)`
+						}
+						return `CREATE TABLE parent (id int NOT NULL, payload text, "select" text, CONSTRAINT parent_key ` + kind + ` (id)` + include + `)` + partition + `; CREATE TABLE child (id int REFERENCES parent(id));`
+					}
+					desired := testutil.ParseSQLToIR(t, sharedTestPostgres, schema(transition.new), "public")
+					db, _, _, _, _, _ := testutil.ConnectToPostgres(t, sharedTestPostgres)
+					defer db.Close()
+					expected := constraintIncludeCatalog(t, db)
+					old := testutil.ParseSQLToIR(t, sharedTestPostgres, schema(transition.old), "public")
+					steps := GenerateMigration(old, desired, "public")
+					require.NotEmpty(t, steps)
+					_, err := db.Exec(buildSQLFromSteps(steps))
+					require.NoError(t, err)
+					require.Equal(t, expected, constraintIncludeCatalog(t, db))
+					actual, err := ir.NewInspector(db, nil).BuildIR(context.Background(), "public")
+					require.NoError(t, err)
+					require.Empty(t, GenerateMigration(actual, desired, "public"))
+				})
+			}
+		}
+	}
+}
+
+func TestConstraintIncludeColumnChanges(t *testing.T) {
+	for _, kind := range []string{"PRIMARY KEY", "UNIQUE"} {
+		for _, change := range []struct{ name, old, new string }{
+			{"add constraint", `CREATE TABLE parent (id int NOT NULL, payload text);`, `CREATE TABLE parent (id int NOT NULL, payload text, CONSTRAINT parent_key ` + kind + ` (id) INCLUDE (payload));`},
+			{"add columns", `CREATE TABLE parent (existing int);`, `CREATE TABLE parent (existing int, id int, payload text, CONSTRAINT parent_key ` + kind + ` (id) INCLUDE (payload));`},
+			{"drop included column", `CREATE TABLE parent (id int NOT NULL, payload text, CONSTRAINT parent_key ` + kind + ` (id) INCLUDE (payload));`, `CREATE TABLE parent (id int NOT NULL, CONSTRAINT parent_key ` + kind + ` (id));`},
+		} {
+			t.Run(kind+"/"+change.name, func(t *testing.T) {
+				desired := testutil.ParseSQLToIR(t, sharedTestPostgres, change.new, "public")
+				db, _, _, _, _, _ := testutil.ConnectToPostgres(t, sharedTestPostgres)
+				defer db.Close()
+				expected := constraintIncludeCatalog(t, db)
+				old := testutil.ParseSQLToIR(t, sharedTestPostgres, change.old, "public")
+				_, err := db.Exec(buildSQLFromSteps(GenerateMigration(old, desired, "public")))
+				require.NoError(t, err)
+				require.Equal(t, expected, constraintIncludeCatalog(t, db))
+				actual, err := ir.NewInspector(db, nil).BuildIR(context.Background(), "public")
+				require.NoError(t, err)
+				require.Empty(t, GenerateMigration(actual, desired, "public"))
+			})
+		}
+	}
+}
+
+func TestConstraintIncludeAttachDetach(t *testing.T) {
+	const attached = `CREATE TABLE parent (id int, payload text, "select" text,
+ CONSTRAINT parent_pk PRIMARY KEY (id) INCLUDE ("select", payload),
+ CONSTRAINT parent_unique UNIQUE (id) INCLUDE (payload, "select")) PARTITION BY RANGE (id);
+ CREATE TABLE leaf PARTITION OF parent FOR VALUES FROM (0) TO (10);`
+	const detached = attached + ` ALTER TABLE parent DETACH PARTITION leaf;`
+	for _, change := range []struct{ name, old, new string }{
+		{"detach", attached, detached}, {"attach", detached, attached},
+		{"attach different INCLUDE", detached + ` ALTER TABLE leaf DROP CONSTRAINT leaf_pkey; ALTER TABLE leaf ADD CONSTRAINT leaf_pkey PRIMARY KEY (id) INCLUDE (payload);`, attached},
+	} {
+		t.Run(change.name, func(t *testing.T) {
+			desired := testutil.ParseSQLToIR(t, sharedTestPostgres, change.new, "public")
+			db, _, _, _, _, _ := testutil.ConnectToPostgres(t, sharedTestPostgres)
+			defer db.Close()
+			expected := constraintIncludeCatalog(t, db)
+			old := testutil.ParseSQLToIR(t, sharedTestPostgres, change.old, "public")
+			_, err := db.Exec(buildSQLFromSteps(GenerateMigration(old, desired, "public")))
+			require.NoError(t, err)
+			require.Equal(t, expected, constraintIncludeCatalog(t, db))
+			actual, err := ir.NewInspector(db, nil).BuildIR(context.Background(), "public")
+			require.NoError(t, err)
+			require.Empty(t, GenerateMigration(actual, desired, "public"))
+		})
+	}
+}
+
+func TestConstraintIncludeComparison(t *testing.T) {
+	for _, kind := range []ir.ConstraintType{ir.ConstraintTypePrimaryKey, ir.ConstraintTypeUnique} {
+		original := &ir.Constraint{Name: "key", Type: kind, Columns: constraintColumns("id"), IncludeColumns: []string{"payload", "select"}}
+		for _, included := range [][]string{nil, {"payload"}, {"select", "payload"}, {"payload", "select"}} {
+			changed := *original
+			changed.IncludeColumns = included
+			equal := len(included) == 2 && included[0] == "payload"
+			require.Equal(t, equal, constraintsEqual(original, &changed))
+			require.Equal(t, equal, parentSuppliesConstraint([]*ir.Constraint{original}, &changed))
+		}
+	}
+}

internal/diff/partition.go

@@ -228,7 +228,8 @@ func parentSuppliesConstraint(supplied []*ir.Constraint, constraint *ir.Constrai
 				candidate.MatchFull != constraint.MatchFull {
 				continue
 			}
-		} else if candidate.NullsNotDistinct != constraint.NullsNotDistinct {
+		} else if candidate.NullsNotDistinct != constraint.NullsNotDistinct ||
+			!slices.Equal(candidate.IncludeColumns, constraint.IncludeColumns) {
 			continue
 		}
 		return true

internal/diff/partition_test.go

@@ -1,6 +1,7 @@
 package diff
 
 import (
+	"slices"
 	"testing"
 
 	"github.com/pgplex/pgschema/ir"
@@ -30,6 +31,8 @@ func TestParentSuppliesConstraint(t *testing.T) {
 	childPK.Name = "events_2026_01_pkey"
 	deferrablePK := childPK
 	deferrablePK.Deferrable = true
+	includedPK := childPK
+	includedPK.IncludeColumns = []string{"payload"}
 	childFK := *parentFK
 	childFK.Name = "events_2026_01_account_id_fkey"
 	setNullAllFK := childFK
@@ -43,6 +46,7 @@ func TestParentSuppliesConstraint(t *testing.T) {
 	}{
 		{"primary key copy", &childPK, true},
 		{"primary key with different deferrability", &deferrablePK, false},
+		{"primary key with different included columns", &includedPK, false},
 		{"foreign key copy", &childFK, true},
 		{"foreign key with different SET NULL column list", &setNullAllFK, false},
 		{"check constraint", checkCopy, false},
@@ -60,7 +64,7 @@ func TestParentSuppliesConstraint(t *testing.T) {
 // constraints under the names the catalog gave them, so the old state used for
 // the diff must list them; the inspected table itself is left untouched. (#639)
 func TestWithConstraintsRetainedByDetach(t *testing.T) {
-	parentPK := &ir.Constraint{Schema: "public", Table: "events", Name: "events_pkey", Type: ir.ConstraintTypePrimaryKey, Columns: constraintColumns("id", "at"), Comment: "parent comment"}
+	parentPK := &ir.Constraint{Schema: "public", Table: "events", Name: "events_pkey", Type: ir.ConstraintTypePrimaryKey, Columns: constraintColumns("id", "at"), IncludeColumns: []string{"payload", "select"}, Comment: "parent comment"}
 	parentFK := &ir.Constraint{Schema: "public", Table: "events", Name: "events_account_id_fkey", Type: ir.ConstraintTypeForeignKey, Columns: constraintColumns("account_id"), ReferencedTable: "accounts", ReferencedColumns: constraintColumns("id")}
 	parent := &ir.Table{Schema: "public", Name: "events", Constraints: map[string]*ir.Constraint{parentPK.Name: parentPK, parentFK.Name: parentFK}}
 	check := &ir.Constraint{Schema: "public", Table: "events_2026_01", Name: "events_amount_check", Type: ir.ConstraintTypeCheck}
@@ -80,7 +84,7 @@ func TestWithConstraintsRetainedByDetach(t *testing.T) {
 		t.Fatalf("got %d constraints, want 3", len(got.Constraints))
 	}
 	pk := got.Constraints["events_2026_01_pkey"]
-	if pk == nil || pk.Type != ir.ConstraintTypePrimaryKey || pk.Table != "events_2026_01" || pk.Comment != "" {
+	if pk == nil || pk.Type != ir.ConstraintTypePrimaryKey || pk.Table != "events_2026_01" || pk.Comment != "" || !slices.Equal(pk.IncludeColumns, parentPK.IncludeColumns) {
 		t.Errorf("retained primary key = %+v", pk)
 	}
 	if fk := got.Constraints["events_account_id_fkey"]; fk == nil || fk.ReferencedTable != "accounts" || fk.Table != "events_2026_01" {

internal/diff/table.go

@@ -1199,6 +1199,11 @@ func constraintDroppedWithColumns(constraint *ir.Constraint, droppedColumnSet ma
 			return true
 		}
 	}
+	for _, column := range constraint.IncludeColumns {
+		if droppedColumnSet[column] {
+			return true
+		}
+	}
 
 	return false
 }
@@ -1402,6 +1407,10 @@ func (td *tableDiff) generateAlterTableStatements(targetSchema string, collector
 			if deferredFKSet[constraintPathKey(constraint)] {
 				continue
 			}
+			// INCLUDE needs table-level syntax after all new columns exist.
+			if len(constraint.IncludeColumns) > 0 {
+				continue
+			}
 			// Only add inline for single-column constraints
 			if len(constraint.Columns) == 1 && constraint.Columns[0].Name == column.Name {
 				switch constraint.Type {
@@ -1522,8 +1531,8 @@ func (td *tableDiff) generateAlterTableStatements(targetSchema string, collector
 			if constraint.NullsNotDistinct {
 				modifier = " NULLS NOT DISTINCT"
 			}
-			sql := fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s UNIQUE%s (%s)%s;",
-				tableName, ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(columnNames, ", "), deferrableClause(constraint))
+			sql := fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s UNIQUE%s (%s)%s%s;",
+				tableName, ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(columnNames, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 
 			context := &diffContext{
 				Type:                DiffTypeTableConstraint,
@@ -1591,8 +1600,8 @@ func (td *tableDiff) generateAlterTableStatements(targetSchema string, collector
 				columnNames[len(columnNames)-1] = columnNames[len(columnNames)-1] + " WITHOUT OVERLAPS"
 			}
 			tableName := getTableNameWithSchema(td.Table.Schema, td.Table.Name, targetSchema)
-			sql := fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s PRIMARY KEY (%s)%s;",
-				tableName, ir.QuoteIdentifier(constraint.Name), strings.Join(columnNames, ", "), deferrableClause(constraint))
+			sql := fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s PRIMARY KEY (%s)%s%s;",
+				tableName, ir.QuoteIdentifier(constraint.Name), strings.Join(columnNames, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 
 			context := &diffContext{
 				Type:                DiffTypeTableConstraint,
@@ -1660,8 +1669,8 @@ func (td *tableDiff) generateAlterTableStatements(targetSchema string, collector
 			if constraint.NullsNotDistinct {
 				modifier = " NULLS NOT DISTINCT"
 			}
-			addSQL = fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s UNIQUE%s (%s)%s;",
-				tableName, ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(columnNames, ", "), deferrableClause(constraint))
+			addSQL = fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s UNIQUE%s (%s)%s%s;",
+				tableName, ir.QuoteIdentifier(constraint.Name), modifier, strings.Join(columnNames, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 
 		case ir.ConstraintTypeCheck:
 			// Add CHECK constraint with ensured outer parentheses
@@ -1698,8 +1707,8 @@ func (td *tableDiff) generateAlterTableStatements(targetSchema string, collector
 			if constraint.IsTemporal && len(columnNames) > 0 {
 				columnNames[len(columnNames)-1] = columnNames[len(columnNames)-1] + " WITHOUT OVERLAPS"
 			}
-			addSQL = fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s PRIMARY KEY (%s)%s;",
-				tableName, ir.QuoteIdentifier(constraint.Name), strings.Join(columnNames, ", "), deferrableClause(constraint))
+			addSQL = fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s PRIMARY KEY (%s)%s%s;",
+				tableName, ir.QuoteIdentifier(constraint.Name), strings.Join(columnNames, ", "), constraintIncludeClause(constraint), deferrableClause(constraint))
 
 		case ir.ConstraintTypeExclusion:
 			addSQL = fmt.Sprintf("ALTER TABLE %s\nADD CONSTRAINT %s %s;",

ir/inspector.go

@@ -603,6 +603,7 @@ func (i *Inspector) buildConstraints(ctx context.Context, schema *IR, targetSche
 				Name:             constraintName,
 				Type:             cType,
 				Columns:          []*ConstraintColumn{},
+				IncludeColumns:   constraint.IncludeColumns,
 				NoInherit:        constraint.NoInherit,
 				IsTemporal:       constraint.IsPeriod.Bool,         // PG18 temporal constraint (WITHOUT OVERLAPS / PERIOD)
 				NullsNotDistinct: constraint.NullsNotDistinct.Bool, // PG15+ UNIQUE NULLS NOT DISTINCT

ir/ir.go

@@ -262,6 +262,7 @@ type Constraint struct {
 	Name                string              `json:"name"`
 	Type                ConstraintType      `json:"type"`
 	Columns             []*ConstraintColumn `json:"columns"`
+	IncludeColumns      []string            `json:"include_columns,omitempty"` // Ordered non-key columns of the backing index
 	ReferencedSchema    string              `json:"referenced_schema,omitempty"`
 	ReferencedTable     string              `json:"referenced_table,omitempty"`
 	ReferencedColumns   []*ConstraintColumn `json:"referenced_columns,omitempty"`