Repository: ariga/atlas
Pull request #3163 (2024-09-29 11:29:29)

sql/internal/sqlx/plan: handle row types in planning
Approved internally. 


Patch:
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 82d7eb09..6ebfd410 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -231,6 +231,21 @@ func ObjectRef(s *schema.Schema, o SpecTypeNamer) *schemahcl.Ref {
 	return schemahcl.BuildRef([]schemahcl.PathIndex{idx})
 }
 
+// TableSpecRef returns a reference to the table in the spec. In case there is more than
+// one table with the same name, the reference will be qualified with the schema name.
+func TableSpecRef(t *schema.Table) *schemahcl.Ref {
+	typ, name := typeTable, t.Name
+	idx := schemahcl.PathIndex{T: typ, V: []string{name}}
+	if s := t.Schema; s != nil && s.Realm != nil && len(s.Realm.Schemas) > 1 && slices.ContainsFunc(s.Realm.Schemas, func(s1 *schema.Schema) bool {
+		return s1 != s && slices.ContainsFunc(s1.Tables, func(t1 *schema.Table) bool {
+			return t1.Name == t.Name
+		})
+	}) {
+		idx.V = append([]string{s.Name}, idx.V...)
+	}
+	return schemahcl.BuildRef([]schemahcl.PathIndex{idx})
+}
+
 // HCLBytesFunc returns a helper that evaluates an HCL document from a byte slice instead
 // of from an hclparse.Parser instance.
 func HCLBytesFunc(ev schemahcl.Evaluator) func(b []byte, v any, inp map[string]cty.Value) error {
diff --git a/sql/internal/sqlx/plan.go b/sql/internal/sqlx/plan.go
index 7928d950..b687c3ba 100644
--- a/sql/internal/sqlx/plan.go
+++ b/sql/internal/sqlx/plan.go
@@ -447,14 +447,21 @@ func SortChanges(changes []schema.Change, opts *SortOptions) []schema.Change {
 	return planned
 }
 
-// Depender can be implemented by an object to determine if a change to it
-// depends on other change, or if other change depends on it. For example:
-// A table creation depends on type creation, and a type deletion depends on
-// table deletion.
-type Depender interface {
-	DependsOn(change, other schema.Change) bool
-	DependencyOf(change, other schema.Change) bool
-}
+type (
+	// Depender can be implemented by an object to determine if a change to it
+	// depends on other change, or if other change depends on it. For example:
+	// A table creation depends on type creation, and a type deletion depends on
+	// table deletion.
+	Depender interface {
+		DependsOn(change, other schema.Change) bool
+		DependencyOf(change, other schema.Change) bool
+	}
+	// RowTyper can be implemented by a type to determine if its source
+	// is a regular table (e.g., row types).
+	RowTyper interface {
+		RowType() *schema.Table
+	}
+)
 
 // dependsOn reports if the given change depends on the other change.
 func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
@@ -492,6 +499,11 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 			if refTo(c1.T.ForeignKeys, c2.T) {
 				return true
 			}
+			if slices.ContainsFunc(c1.T.Columns, func(c *schema.Column) bool {
+				return c.Type != nil && typeDependsOnT(c.Type.Type, c2.T)
+			}) {
+				return true
+			}
 		case *schema.ModifyTable:
 			if (c1.T.Name != c2.T.Name || !SameSchema(c1.T.Schema, c2.T.Schema)) && refTo(c1.T.ForeignKeys, c2.T) {
 				return true
@@ -512,18 +524,43 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 		// after all resources that rely on it will be dropped.
 		switch c2 := c2.(type) {
 		case *schema.DropTable:
+			// References to this table, must be dropped first.
 			if refTo(c2.T.ForeignKeys, c1.T) {
 				return true
 			}
+			if slices.ContainsFunc(c2.T.Columns, func(c *schema.Column) bool {
+				return c.Type != nil && typeDependsOnT(c.Type.Type, c1.T)
+			}) {
+				return true
+			}
 		case *schema.ModifyTable:
-			return slices.ContainsFunc(c2.Changes, func(c schema.Change) bool {
-				fk, ok := c.(*schema.DropForeignKey)
-				return ok && refTo([]*schema.ForeignKey{fk.F}, c1.T)
-			})
+			if slices.ContainsFunc(c2.Changes, func(c schema.Change) bool {
+				switch c := c.(type) {
+				case *schema.DropForeignKey:
+					return refTo([]*schema.ForeignKey{c.F}, c1.T)
+				case *schema.DropColumn:
+					return c.C.Type != nil && typeDependsOnT(c.C.Type.Type, c1.T)
+				}
+				return false
+			}) {
+				return true
+			}
 		case *schema.DropTrigger:
 			if SameTable(c2.T.Table, c1.T) {
 				return true
 			}
+		case *schema.DropFunc:
+			if c2.F.Ret != nil && typeDependsOnT(c2.F.Ret, c1.T) || slices.ContainsFunc(c2.F.Args, func(f *schema.FuncArg) bool {
+				return typeDependsOnT(f.Type, c1.T)
+			}) {
+				return true
+			}
+		case *schema.DropProc:
+			if slices.ContainsFunc(c2.P.Args, func(f *schema.FuncArg) bool {
+				return typeDependsOnT(f.Type, c1.T)
+			}) {
+				return true
+			}
 		}
 		return depOfDrop(c1.T, c2)
 	case *schema.ModifyTable:
@@ -534,10 +571,19 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 				return true
 			}
 			// Tables need to be created before referencing them.
-			return slices.ContainsFunc(c1.Changes, func(c schema.Change) bool {
-				fk, ok := c.(*schema.AddForeignKey)
-				return ok && refTo([]*schema.ForeignKey{fk.F}, c2.T)
-			})
+			if slices.ContainsFunc(c1.Changes, func(c schema.Change) bool {
+				switch c := c.(type) {
+				case *schema.AddForeignKey:
+					return refTo([]*schema.ForeignKey{c.F}, c2.T)
+				case *schema.AddColumn:
+					return c.C.Type != nil && typeDependsOnT(c.C.Type.Type, c2.T)
+				case *schema.ModifyColumn:
+					return c.To.Type != nil && typeDependsOnT(c.To.Type.Type, c2.T)
+				}
+				return false
+			}) {
+				return true
+			}
 		case *schema.ModifyTable:
 			if c1.T != c2.T {
 				addC := make(map[*schema.Column]bool)
@@ -631,6 +677,11 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 			if opts.FuncDepT != nil && opts.FuncDepT(c1.F, c2.T) {
 				return true
 			}
+			if c1.F.Ret != nil && typeDependsOnT(c1.F.Ret, c2.T) || slices.ContainsFunc(c1.F.Args, func(f *schema.FuncArg) bool {
+				return typeDependsOnT(f.Type, c2.T)
+			}) {
+				return true
+			}
 		case *schema.AddView:
 			if opts.FuncDepV != nil && opts.FuncDepV(c1.F, c2.V) {
 				return true
@@ -679,6 +730,12 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 		switch c2 := c2.(type) {
 		case *schema.AddSchema:
 			return c1.P.Schema.Name == c2.S.Name
+		case *schema.AddTable:
+			if slices.ContainsFunc(c1.P.Args, func(f *schema.FuncArg) bool {
+				return typeDependsOnT(f.Type, c2.T)
+			}) {
+				return true
+			}
 		case *schema.DropProc:
 			return c1.P.Name == c2.P.Name && SameSchema(c1.P.Schema, c2.P.Schema) // Proc recreation.
 		case *schema.AddProc:
@@ -734,16 +791,14 @@ func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 		// Dropping a type must occur after all its usage were dropped.
 		switch c2 := c2.(type) {
 		case *schema.DropView:
-			// Dropping a view also drops its triggers
-			// and they might depend on the type.
+			// Dropping a view also drops its triggers and might depend on the type.
 			if slices.ContainsFunc(c2.V.Triggers, func(tg *schema.Trigger) bool {
 				return slices.Contains(tg.Deps, c1.O)
 			}) {
 				return true
 			}
 		case *schema.DropTable:
-			// Dropping a table also drops its triggers
-			// and they might depend on the type.
+			// Dropping a table also drops its triggers and might depend on the type.
 			if slices.ContainsFunc(c2.T.Triggers, func(tg *schema.Trigger) bool {
 				return slices.Contains(tg.Deps, c1.O)
 			}) {
@@ -912,6 +967,15 @@ func refTo(fks []*schema.ForeignKey, to *schema.Table) bool {
 	})
 }
 
+// typeDependsOnT reports if the declaration of type t1 depends on the given change.
+func typeDependsOnT(t schema.Type, tt *schema.Table) bool {
+	rt, ok := schema.UnderlyingType(t).(RowTyper)
+	if !ok {
+		return false
+	}
+	return SameTable(rt.RowType(), tt)
+}
+
 // dependsOnT reports if t1 depends on t2.
 func dependsOnT(t1, t2 schema.Type) bool {
 	// Comparing might panic due to mismatch types.
diff --git a/sql/postgres/convert.go b/sql/postgres/convert.go
index 5eb89e7c..6e2d59b5 100644
--- a/sql/postgres/convert.go
+++ b/sql/postgres/convert.go
@@ -42,6 +42,11 @@ func FormatType(t schema.Type) (string, error) {
 			return "", errors.New("postgres: missing composite type name")
 		}
 		f = t.T
+	case *RowType:
+		if t.T == nil || t.T.Name == "" {
+			return "", errors.New("postgres: missing table for composite (row) type")
+		}
+		f = t.T.Name
 	case *DomainType:
 		if t.T == "" {
 			return "", errors.New("postgres: missing domain type name")
@@ -298,12 +303,12 @@ func columnType(c *columnDesc) (schema.Type, error) {
 		if ft == "" {
 			ft = t
 		}
-		typ = &UserDefinedType{T: ft}
+		typ = &UserDefinedType{T: ft, C: c.typtype}
 	}
 	switch c.typtype {
 	case "d", "e":
 		// User-defined types supported by Atlas.
-		typ = &UserDefinedType{T: c.fmtype}
+		typ = &UserDefinedType{T: c.fmtype, C: c.typtype}
 	}
 	return typ, nil
 }
diff --git a/sql/postgres/inspect_oss.go b/sql/postgres/inspect_oss.go
index 3ed4c4c6..ed7834c1 100644
--- a/sql/postgres/inspect_oss.go
+++ b/sql/postgres/inspect_oss.go
@@ -906,6 +906,14 @@ type (
 	UserDefinedType struct {
 		schema.Type
 		T string
+		C string // Optional type class.
+
+	}
+
+	// RowType defines a composite type that represents a table row.
+	RowType struct {
+		schema.Type
+		T *schema.Table // Table that this row type represents.
 	}
 
 	// PseudoType defines a non-column pseudo-type, such as function arguments and return types.


Test patch:
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index e6340e59..da7299fa 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -1420,7 +1420,7 @@ create table atlas_types_sanity
 				},
 				{
 					Name: "tUserDefined",
-					Type: &schema.ColumnType{Type: &postgres.UserDefinedType{T: "public.address"}, Raw: "USER-DEFINED", Null: true},
+					Type: &schema.ColumnType{Type: &postgres.UserDefinedType{T: "public.address", C: "c"}, Raw: "USER-DEFINED", Null: true},
 					Default: &schema.RawExpr{
 						X: "'(ab,cd)'::public.address",
 					},
diff --git a/sql/internal/sqlx/plan_test.go b/sql/internal/sqlx/plan_test.go
index 638cc23b..7374a1eb 100644
--- a/sql/internal/sqlx/plan_test.go
+++ b/sql/internal/sqlx/plan_test.go
@@ -296,6 +296,55 @@ func TestSortChanges(t *testing.T) {
 	require.IsType(t, &schema.DropTrigger{}, planned[1])
 }
 
+// RowType is a mocked RowTyper implementation.
+type RowType struct {
+	schema.Type
+	*schema.Table
+}
+
+func (r *RowType) RowType() *schema.Table { return r.Table }
+
+func TestSortChanges_RowTyper(t *testing.T) {
+	t1, t2 := schema.NewTable("t1"), schema.NewTable("t2")
+	t2.AddColumns(schema.NewColumn("t1").SetType(&RowType{Table: t1}))
+	// Tables created before other tables that depend on them.
+	changes := []schema.Change{&schema.AddTable{T: t2}, &schema.AddTable{T: t1}}
+	planned := SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+
+	// Tables created before other columns that depend on them.
+	changes = []schema.Change{
+		&schema.ModifyTable{T: t2, Changes: schema.Changes{&schema.AddColumn{C: t2.Columns[0]}}},
+		&schema.AddTable{T: t1},
+	}
+	planned = SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+
+	// Tables dropped after other tables that depend on them.
+	changes = []schema.Change{&schema.DropTable{T: t1}, &schema.DropTable{T: t2}}
+	planned = SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+
+	// Tables dropped after other columns that depend on them.
+	changes = []schema.Change{
+		&schema.DropTable{T: t1},
+		&schema.ModifyTable{T: t2, Changes: schema.Changes{&schema.DropColumn{C: t2.Columns[0]}}},
+	}
+	planned = SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+
+	// Functions created after the tables they depend on.
+	f1 := &schema.Func{Name: "f1", Ret: &RowType{Table: t1}}
+	changes = []schema.Change{&schema.AddFunc{F: f1}, &schema.AddTable{T: t1}}
+	planned = SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+
+	// Functions dropped before the tables they depend on.
+	changes = []schema.Change{&schema.DropTable{T: t1}, &schema.DropFunc{F: f1}}
+	planned = SortChanges(changes, nil)
+	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
+}
+
 func TestSameTable(t *testing.T) {
 	t1 := schema.NewTable("t1")
 	require.True(t, SameTable(t1, t1))
diff --git a/sql/postgres/inspect_oss_test.go b/sql/postgres/inspect_oss_test.go
index 49c35a09..8ccaf014 100644
--- a/sql/postgres/inspect_oss_test.go
+++ b/sql/postgres/inspect_oss_test.go
@@ -135,7 +135,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					{Name: "c20", Type: &schema.ColumnType{Raw: "uuid", Type: &schema.UUIDType{T: "uuid"}}},
 					{Name: "c21", Type: &schema.ColumnType{Raw: "xml", Type: &XMLType{T: "xml"}}},
 					{Name: "c22", Type: &schema.ColumnType{Raw: "ARRAY", Null: true, Type: &ArrayType{Type: &schema.IntegerType{T: "integer"}, T: "integer[]"}}},
-					{Name: "c23", Type: &schema.ColumnType{Raw: "USER-DEFINED", Null: true, Type: &UserDefinedType{T: "ltree"}}},
+					{Name: "c23", Type: &schema.ColumnType{Raw: "USER-DEFINED", Null: true, Type: &UserDefinedType{T: "ltree", C: "b"}}},
 					{Name: "c24", Type: &schema.ColumnType{Raw: "state", Type: stateE}},
 					{Name: "c25", Type: &schema.ColumnType{Raw: "timestamp without time zone", Type: &schema.TimeType{T: "timestamp without time zone", Precision: sqlx.P(4)}}, Default: &schema.RawExpr{X: "now()"}},
 					{Name: "c26", Type: &schema.ColumnType{Raw: "timestamp with time zone", Type: &schema.TimeType{T: "timestamp with time zone", Precision: sqlx.P(6)}}},
@@ -147,7 +147,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					{Name: "c32", Type: &schema.ColumnType{Raw: "bigserial", Type: &SerialType{T: "bigserial", SequenceName: "t1_c32_seq"}}},
 					{Name: "c33", Type: &schema.ColumnType{Raw: `status`, Type: statusE}, Default: &schema.Literal{V: "'unknown'"}},
 					{Name: "c34", Type: &schema.ColumnType{Raw: "ARRAY", Type: &ArrayType{T: "state[]", Type: stateE}}},
-					{Name: "c35", Type: &schema.ColumnType{Raw: "domain_char", Type: &UserDefinedType{T: "domain_char"}}},
+					{Name: "c35", Type: &schema.ColumnType{Raw: "domain_char", Type: &UserDefinedType{T: "domain_char", C: "d"}}},
 					{Name: "c36", Type: &schema.ColumnType{Raw: "tsvector", Type: &TextSearchType{T: "tsvector"}}},
 					{Name: "c37", Type: &schema.ColumnType{Raw: "tsquery", Type: &TextSearchType{T: "tsquery"}}},
 					{Name: "c38", Type: &schema.ColumnType{Raw: "datemultirange", Type: &RangeType{T: "datemultirange"}}},