dialect/sql/schema: without foreign keys option for atlas (#2404)

* dialect/sql: without foreign key option for atlas

* handle fks in modifytables

* add tests

* tests

* Update dialect/sql/schema/atlas.go

Co-authored-by: Ariel Mashraki <7413593+a8m@users.noreply.github.com>

* ct and tests

Co-authored-by: Ariel Mashraki <7413593+a8m@users.noreply.github.com>

dialect/sql/schema/atlas.go

@@ -23,7 +23,7 @@ import (
 type (
 	// Differ is the interface that wraps the Diff method.
 	Differ interface {
-		// Diff creates the given tables in the database.
+		// Diff returns a list of changes that construct a migration plan.
 		Diff(current, desired *schema.Schema) ([]schema.Change, error)
 	}
 
@@ -98,7 +98,7 @@ const (
 	DropCheck
 )
 
-// Is reports whether c is match the given change king.
+// Is reports whether c is match the given change kind.
 func (k ChangeKind) Is(c ChangeKind) bool {
 	return k == c || k&c != 0
 }
@@ -170,6 +170,34 @@ func filterChanges(skip ChangeKind) DiffHook {
 	}
 }
 
+func withoutForeignKeys(next Differ) Differ {
+	return DiffFunc(func(current, desired *schema.Schema) ([]schema.Change, error) {
+		changes, err := next.Diff(current, desired)
+		if err != nil {
+			return nil, err
+		}
+		for _, c := range changes {
+			switch c := c.(type) {
+			case *schema.AddTable:
+				c.T.ForeignKeys = nil
+			case *schema.ModifyTable:
+				c.T.ForeignKeys = nil
+				filtered := make([]schema.Change, 0, len(c.Changes))
+				for _, change := range c.Changes {
+					switch change.(type) {
+					case *schema.AddForeignKey, *schema.DropForeignKey, *schema.ModifyForeignKey:
+						continue
+					default:
+						filtered = append(filtered, change)
+					}
+				}
+				c.Changes = filtered
+			}
+		}
+		return changes, nil
+	})
+}
+
 type (
 	// Applier is the interface that wraps the Apply method.
 	Applier interface {
@@ -267,9 +295,6 @@ func (m *Migrate) setupAtlas() error {
 	if !m.atlas.enabled {
 		return nil
 	}
-	if !m.withForeignKeys {
-		return errors.New("sql/schema: WithForeignKeys(false) does not work in Atlas migration")
-	}
 	if m.withFixture {
 		return errors.New("sql/schema: WithFixture(true) does not work in Atlas migration")
 	}
@@ -286,6 +311,9 @@ func (m *Migrate) setupAtlas() error {
 	if skip != NoChange {
 		m.atlas.diff = append(m.atlas.diff, filterChanges(skip))
 	}
+	if !m.withForeignKeys {
+		m.atlas.diff = append(m.atlas.diff, withoutForeignKeys)
+	}
 	if m.atlas.dir != nil && m.atlas.fmt == nil {
 		m.atlas.fmt = migrate.DefaultFormatter
 	}

dialect/sql/schema/migrate_test.go

@@ -13,6 +13,8 @@ import (
 	"time"
 
 	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+
 	"entgo.io/ent/dialect"
 	"entgo.io/ent/dialect/sql"
 	"github.com/DATA-DOG/go-sqlmock"
@@ -90,3 +92,82 @@ func requireFileEqual(t *testing.T, name, contents string) {
 	require.NoError(t, err)
 	require.Equal(t, contents, string(c))
 }
+
+func TestMigrateWithoutForeignKeys(t *testing.T) {
+	tbl := &schema.Table{
+		Name: "tbl",
+		Columns: []*schema.Column{
+			{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}},
+		},
+	}
+	fk := &schema.ForeignKey{
+		Symbol:     "fk",
+		Table:      tbl,
+		Columns:    tbl.Columns[1:],
+		RefTable:   tbl,
+		RefColumns: tbl.Columns[:1],
+		OnUpdate:   schema.NoAction,
+		OnDelete:   schema.Cascade,
+	}
+	tbl.ForeignKeys = append(tbl.ForeignKeys, fk)
+	t.Run("AddTable", func(t *testing.T) {
+		mdiff := DiffFunc(func(_, _ *schema.Schema) ([]schema.Change, error) {
+			return []schema.Change{
+				&schema.AddTable{
+					T: tbl,
+				},
+			}, nil
+		})
+		df, err := withoutForeignKeys(mdiff).Diff(nil, nil)
+		require.NoError(t, err)
+		require.Len(t, df, 1)
+		actual, ok := df[0].(*schema.AddTable)
+		require.True(t, ok)
+		require.Nil(t, actual.T.ForeignKeys)
+	})
+	t.Run("ModifyTable", func(t *testing.T) {
+		mdiff := DiffFunc(func(_, _ *schema.Schema) ([]schema.Change, error) {
+			return []schema.Change{
+				&schema.ModifyTable{
+					T: tbl,
+					Changes: []schema.Change{
+						&schema.AddIndex{
+							I: &schema.Index{
+								Name: "id_key",
+								Parts: []*schema.IndexPart{
+									{C: tbl.Columns[0]},
+								},
+							},
+						},
+						&schema.DropForeignKey{
+							F: fk,
+						},
+						&schema.AddForeignKey{
+							F: fk,
+						},
+						&schema.ModifyForeignKey{
+							From:   fk,
+							To:     fk,
+							Change: schema.ChangeRefColumn,
+						},
+						&schema.AddColumn{
+							C: &schema.Column{Name: "name", Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar(255)"}}},
+						},
+					},
+				},
+			}, nil
+		})
+		df, err := withoutForeignKeys(mdiff).Diff(nil, nil)
+		require.NoError(t, err)
+		require.Len(t, df, 1)
+		actual, ok := df[0].(*schema.ModifyTable)
+		require.True(t, ok)
+		require.Len(t, actual.Changes, 2)
+		addIndex, ok := actual.Changes[0].(*schema.AddIndex)
+		require.True(t, ok)
+		require.EqualValues(t, "id_key", addIndex.I.Name)
+		addColumn, ok := actual.Changes[1].(*schema.AddColumn)
+		require.True(t, ok)
+		require.EqualValues(t, "name", addColumn.C.Name)
+	})
+}