dialect/sql/schema: without foreign keys option for atlas (ent#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>
gitlawr · Apr 13, 2022 · 815d151 · 815d151
1 parent c42c4ae
commit 815d151
Show file tree

Hide file tree

Showing 4 changed files with 114 additions and 12 deletions.
diff --git a/dialect/sql/schema/atlas.go b/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
 	}

diff --git a/dialect/sql/schema/migrate_test.go b/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)
+	})
+}
diff --git a/go.mod b/go.mod
@@ -32,7 +32,6 @@ require (
 	github.com/google/go-cmp v0.5.6 // indirect
 	github.com/hashicorp/hcl/v2 v2.10.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
-	github.com/kr/pretty v0.2.1 // indirect
 	github.com/mattn/go-runewidth v0.0.9 // indirect
 	github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
@@ -43,6 +42,5 @@ require (
 	golang.org/x/sys v0.0.0-20211205182925-97ca703d548d // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
-	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )
diff --git a/go.sum b/go.sum
@@ -1,7 +1,3 @@
-ariga.io/atlas v0.3.7-0.20220303204946-787354f533c3 h1:fjG4oFCQEfGrRi0QoxWcH2OO28CE6VYa6DkIr3yDySU=
-ariga.io/atlas v0.3.7-0.20220303204946-787354f533c3/go.mod h1:yWGf4VPiD4SW83+kAqzD624txN9VKoJC+bpVXr2pKJA=
-ariga.io/atlas v0.3.8-0.20220313134928-770640fc02bf h1:bAt5AUvr91QI8yXHME6qTsMTNM4BtfSB3M9o1cmt51E=
-ariga.io/atlas v0.3.8-0.20220313134928-770640fc02bf/go.mod h1:ipw7dUlFanAylr9nvs8lCvOUC8hFG6PGd/gtr+uJMvk=
 ariga.io/atlas v0.3.8-0.20220314111236-b2171e04c5b2 h1:qbH+CDPAMsV1FIKkHGYzy2aWP9k5QAqPbi9PYZGqz60=
 ariga.io/atlas v0.3.8-0.20220314111236-b2171e04c5b2/go.mod h1:ipw7dUlFanAylr9nvs8lCvOUC8hFG6PGd/gtr+uJMvk=
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
@@ -280,7 +276,6 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kylelemons/godebug v0.0.0-20170820004349-d65d576e9348 h1:MtvEpTB6LX3vkb4ax0b5D2DHbNAUsen0Gx5wZoq3lV4=
 github.com/kylelemons/godebug v0.0.0-20170820004349-d65d576e9348/go.mod h1:B69LEHPfb2qLo0BaaOLcbitczOKLWTsrBG9LczfCD4k=
-github.com/lib/pq v1.10.3/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lib/pq v1.10.4 h1:SO9z7FRPzA03QhHKJrH5BXA6HU1rS4V2nIVrrNC1iYk=
 github.com/lib/pq v1.10.4/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lyft/protoc-gen-star v0.5.3/go.mod h1:V0xaHgaf5oCCqmcxYcWiDfTiKsZsRc87/1qhoTACD8w=