Repository: ariga/atlas
Pull request #2737 (2024-05-01 20:06:18)

sql/internal/sqlx: allow extending deps sorting


Patch:
diff --git a/sql/internal/sqlx/plan.go b/sql/internal/sqlx/plan.go
index 955f3702..3c749f32 100644
--- a/sql/internal/sqlx/plan.go
+++ b/sql/internal/sqlx/plan.go
@@ -382,8 +382,16 @@ func byKeys[T any](m map[string]T) []struct {
 	return vs
 }
 
+// SortOptions allows drivers to customize the behavior of the SortChanges function.
+type SortOptions struct {
+	// FuncDepT reports if a function depends on the given table.
+	FuncDepT func(*schema.Func, *schema.Table) bool
+	// FuncDepO reports if a function depends on the given object.
+	FuncDepO func(*schema.Func, schema.Object) bool
+}
+
 // SortChanges is a helper function to sort to level changes based on their priority.
-func SortChanges(changes []schema.Change) []schema.Change {
+func SortChanges(changes []schema.Change, opts *SortOptions) []schema.Change {
 	var views, drop, other []schema.Change
 	for _, c := range changes {
 		switch c.(type) {
@@ -404,7 +412,7 @@ func SortChanges(changes []schema.Change) []schema.Change {
 	edges := make(map[schema.Change][]schema.Change)
 	for _, c := range changes {
 		for _, c2 := range changes {
-			if c != c2 && dependsOn(c, c2) {
+			if c != c2 && dependsOn(c, c2, V(opts)) {
 				edges[c] = append(edges[c], c2)
 			}
 		}
@@ -444,7 +452,7 @@ type Depender interface {
 }
 
 // dependsOn reports if the given change depends on the other change.
-func dependsOn(c1, c2 schema.Change) bool {
+func dependsOn(c1, c2 schema.Change, opts SortOptions) bool {
 	if dependOnOf(c1, c2) {
 		return true
 	}
@@ -586,6 +594,10 @@ func dependsOn(c1, c2 schema.Change) bool {
 			if funcDep(c1.F, c2.To) {
 				return true // Relies on the new definition.
 			}
+		case *schema.AddTable:
+			if opts.FuncDepT != nil && opts.FuncDepT(c1.F, c2.T) {
+				return true
+			}
 		case *schema.AddObject:
 			t, ok := c2.O.(schema.Type)
 			if ok && (c1.F.Ret == t || slices.ContainsFunc(c1.F.Args, func(f *schema.FuncArg) bool {
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 4341f66b..c627d7ee 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -85,7 +85,7 @@ func (s *state) plan(changes []schema.Change) error {
 	if err != nil {
 		return err
 	}
-	for _, c := range sqlx.SortChanges(planned) {
+	for _, c := range sqlx.SortChanges(planned, nil) {
 		switch c := c.(type) {
 		case *schema.AddTable:
 			err = s.addTable(c)
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index cef91b67..0f7572a6 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -356,6 +356,10 @@ func convertExclude(schemahcl.Resource, *schema.Table) error {
 	return nil // unimplemented.
 }
 
+func (*state) sortChanges(changes []schema.Change) []schema.Change {
+	return sqlx.SortChanges(changes, nil)
+}
+
 func detachCycles(changes []schema.Change) ([]schema.Change, error) {
 	return sqlx.DetachCycles(changes)
 }
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 939ef799..e5f87de6 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -80,7 +80,7 @@ func (s *state) plan(changes []schema.Change) error {
 	if planned, err = detachCycles(planned); err != nil {
 		return err
 	}
-	for _, c := range sqlx.SortChanges(planned) {
+	for _, c := range s.sortChanges(planned) {
 		switch c := c.(type) {
 		case *schema.AddTable:
 			err = s.addTable(c)


Test patch:
diff --git a/sql/internal/sqlx/plan_test.go b/sql/internal/sqlx/plan_test.go
index 297e4b06..274b6a45 100644
--- a/sql/internal/sqlx/plan_test.go
+++ b/sql/internal/sqlx/plan_test.go
@@ -164,7 +164,7 @@ func TestBadTypeComparison(t *testing.T) {
 	SortChanges([]schema.Change{
 		&schema.DropObject{O: o1},
 		&schema.DropTable{T: schema.NewTable("t1").AddColumns(schema.NewColumn("c1").SetType(o1))},
-	})
+	}, nil)
 }
 
 func TestSortChanges(t *testing.T) {
@@ -186,7 +186,7 @@ func TestSortChanges(t *testing.T) {
 		&schema.AddTable{T: t1},
 		&schema.AddFunc{F: f1},
 	}
-	planned := SortChanges(changes)
+	planned := SortChanges(changes, nil)
 	require.Equal(t, []schema.Change{changes[2], changes[0], changes[3], changes[1]}, planned)
 
 	changes = []schema.Change{
@@ -194,14 +194,14 @@ func TestSortChanges(t *testing.T) {
 		&schema.DropFunc{F: f1},
 		&schema.DropTrigger{T: tr1},
 	}
-	planned = SortChanges(changes)
+	planned = SortChanges(changes, nil)
 	require.Equal(t, []schema.Change{changes[2], changes[1], changes[0]}, planned)
 
 	// No changes.
 	planned = SortChanges([]schema.Change{
 		&schema.DropFunc{F: f1},
 		&schema.DropTable{T: t1},
-	})
+	}, nil)
 	require.Equal(t, []schema.Change{&schema.DropFunc{F: f1}, &schema.DropTable{T: t1}}, planned)
 
 	// The table must be dropped before the enum type if one of its triggers depends on the enum type.
@@ -210,7 +210,7 @@ func TestSortChanges(t *testing.T) {
 		&schema.DropObject{O: e1},
 		&schema.DropTable{T: t1},
 	}
-	planned = SortChanges(changes)
+	planned = SortChanges(changes, nil)
 	require.Equal(t, []schema.Change{changes[1], changes[0]}, planned)
 
 	// The table must be dropped before the function if one of its triggers depends on the function.
@@ -218,7 +218,7 @@ func TestSortChanges(t *testing.T) {
 	planned = SortChanges([]schema.Change{
 		&schema.DropFunc{F: f1},
 		&schema.DropTable{T: t1},
-	})
+	}, nil)
 	require.Equal(t, []schema.Change{&schema.DropTable{T: t1}, &schema.DropFunc{F: f1}}, planned)
 
 	// Ignore functions that do not reside on the same schema.
@@ -227,7 +227,7 @@ func TestSortChanges(t *testing.T) {
 	planned = SortChanges([]schema.Change{
 		&schema.AddFunc{F: f1},
 		&schema.DropFunc{F: f2},
-	})
+	}, nil)
 	require.Equal(t, []schema.Change{&schema.AddFunc{F: f1}, &schema.DropFunc{F: f2}}, planned)
 
 	// Order functions that reside on the same schema.
@@ -235,7 +235,7 @@ func TestSortChanges(t *testing.T) {
 	planned = SortChanges([]schema.Change{
 		&schema.AddFunc{F: f1},
 		&schema.DropFunc{F: f2},
-	})
+	}, nil)
 	require.Equal(t, []schema.Change{&schema.DropFunc{F: f2}, &schema.AddFunc{F: f1}}, planned)
 }