refactored checkout to retain feature version data

go/libraries/doltcore/doltdb/feature_version_test.go

@@ -19,7 +19,6 @@ import (
 	"github.com/dolthub/dolt/go/cmd/dolt/commands"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/stretchr/testify/assert"
-	"math"
 	"testing"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
@@ -54,10 +53,10 @@ type fvCommand struct {
 	args args
 }
 
-func (cmd fvCommand) exec(t *testing.T, ctx context.Context, dEnv *env.DoltEnv) int {
+func (cmd fvCommand) exec(ctx context.Context, dEnv *env.DoltEnv) int {
 	// execute the command using |cmd.user|'s Feature Version
-	doltdb.DoltFeatureVersion, cmd.user.vers = cmd.user.vers, doltdb.DoltFeatureVersion
-	defer func() { doltdb.DoltFeatureVersion = cmd.user.vers }()
+	doltdb.DoltFeatureVersion = cmd.user.vers
+	defer func() { doltdb.DoltFeatureVersion = DoltFeatureVersionCopy }()
 
 	return cmd.cmd.Exec(ctx, cmd.cmd.Name(), cmd.args, dEnv)
 }
@@ -153,6 +152,17 @@ func TestFeatureVersion(t *testing.T) {
 			},
 			expVer: oldVersion,
 		},
+		{
+			name: "Checkout does not update feature version",
+			setup: []fvCommand{
+				{OldClient, commands.SqlCmd{}, args{"-q", "CREATE TABLE test (pk int PRIMARY KEY);"}},
+				{OldClient, commands.CommitCmd{}, args{"-a", "-m", "created table"}},
+				{OldClient, commands.BranchCmd{}, args{"other"}},
+				{NewClient, commands.CheckoutCmd{}, args{"other"}},
+				{NewClient, commands.CheckoutCmd{}, args{"master"}},
+			},
+			expVer: oldVersion,
+		},
 		{
 			name: "new client writes to table, locking out old client",
 			setup: []fvCommand{
@@ -173,22 +183,22 @@ func TestFeatureVersion(t *testing.T) {
 	ctx := context.Background()
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+
+			doltdb.DoltFeatureVersion = oldVersion
 			dEnv := dtestutils.CreateTestEnv()
+			doltdb.DoltFeatureVersion = DoltFeatureVersionCopy
 
 			for _, cmd := range test.setup {
-				code := cmd.exec(t, ctx, dEnv)
+				code := cmd.exec(ctx, dEnv)
 				require.Equal(t, 0, code)
 			}
 			for _, cmd := range test.errCmds {
-				code := cmd.exec(t, ctx, dEnv)
+				code := cmd.exec(ctx, dEnv)
 				require.NotEqual(t, 0, code)
 			}
 
-			// ensure |doltdb.DoltFeatureVersion| was restored
-			assert.Equal(t, DoltFeatureVersionCopy, doltdb.DoltFeatureVersion)
-
-			// execute assertions with max feature version
-			doltdb.DoltFeatureVersion = math.MaxInt64
+			// execute assertions with newVersion to avoid OutOfDate errors
+			doltdb.DoltFeatureVersion = newVersion
 			defer func() { doltdb.DoltFeatureVersion = DoltFeatureVersionCopy }()
 
 			root, err := dEnv.WorkingRoot(ctx)
@@ -196,8 +206,11 @@ func TestFeatureVersion(t *testing.T) {
 
 			act, ok, err := root.GetFeatureVersion(ctx)
 			require.NoError(t, err)
-			assert.True(t, ok)
+			require.True(t, ok)
 			assert.Equal(t, test.expVer, act)
 		})
+
+		// ensure |doltdb.DoltFeatureVersion| was restored
+		assert.Equal(t, DoltFeatureVersionCopy, doltdb.DoltFeatureVersion)
 	}
 }

go/libraries/doltcore/doltdb/root_val.go

@@ -49,75 +49,24 @@ type RootValue struct {
 	fkc     *ForeignKeyCollection // cache the first load
 }
 
-func NewRootValue(ctx context.Context, vrw types.ValueReadWriter, tables map[string]hash.Hash, ssMap types.Map, fkMap types.Map) (*RootValue, error) {
-	values := make([]types.Value, 2*len(tables))
-
-	index := 0
-	for k, v := range tables {
-		values[index] = types.String(k)
-		valForHash, err := vrw.ReadValue(ctx, v)
-
-		if err != nil {
-			return nil, err
-		}
-
-		if valForHash == nil {
-			return nil, ErrHashNotFound
-		}
-
-		values[index+1], err = types.NewRef(valForHash, vrw.Format())
-
-		if err != nil {
-			return nil, err
-		}
-
-		index += 2
-	}
-
-	tblMap, err := types.NewMap(ctx, vrw, values...)
-
-	if err != nil {
-		return nil, err
-	}
-
-	return newRootFromMaps(vrw, tblMap, ssMap, fkMap)
-}
-
 func newRootValue(vrw types.ValueReadWriter, st types.Struct) *RootValue {
 	return &RootValue{vrw, st, nil}
 }
 
 func emptyRootValue(ctx context.Context, vrw types.ValueReadWriter) (*RootValue, error) {
-	m, err := types.NewMap(ctx, vrw)
-
+	empty, err := types.NewMap(ctx, vrw)
 	if err != nil {
 		return nil, err
 	}
 
-	mm, err := types.NewMap(ctx, vrw)
-
-	if err != nil {
-		return nil, err
-	}
-
-	mmm, err := types.NewMap(ctx, vrw)
-
-	if err != nil {
-		return nil, err
-	}
-
-	return newRootFromMaps(vrw, m, mm, mmm)
-}
-
-func newRootFromMaps(vrw types.ValueReadWriter, tblMap types.Map, ssMap types.Map, fkMap types.Map) (*RootValue, error) {
 	sd := types.StructData{
-		tablesKey:       tblMap,
-		superSchemasKey: ssMap,
-		foreignKeyKey:   fkMap,
+		tablesKey:       empty,
+		superSchemasKey: empty,
+		foreignKeyKey:   empty,
+		featureVersKey:  types.Int(DoltFeatureVersion),
 	}
 
 	st, err := types.NewStruct(vrw.Format(), ddbRootStructName, sd)
-
 	if err != nil {
 		return nil, err
 	}
@@ -759,79 +708,6 @@ func (root *RootValue) HashOf() (hash.Hash, error) {
 	return root.valueSt.Hash(root.vrw.Format())
 }
 
-// UpdateTablesFromOther takes the tables from the given root and applies them to the calling root, along with any
-// foreign keys and other table-related data.
-func (root *RootValue) UpdateTablesFromOther(ctx context.Context, tblNames []string, other *RootValue) (*RootValue, error) {
-	tableMap, err := root.getTableMap()
-	if err != nil {
-		return nil, err
-	}
-	fkCollection, err := root.GetForeignKeyCollection(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	otherMap, err := other.getTableMap()
-	if err != nil {
-		return nil, err
-	}
-	otherFkCollection, err := other.GetForeignKeyCollection(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	var fksToAdd []ForeignKey
-	var fksToRemove []ForeignKey
-
-	me := tableMap.Edit()
-	for _, tblName := range tblNames {
-		key := types.String(tblName)
-		if val, ok, err := otherMap.MaybeGet(ctx, key); err != nil {
-			return nil, err
-		} else if ok {
-			me = me.Set(key, val)
-			newFks, _ := otherFkCollection.KeysForTable(tblName)
-			fksToAdd = append(fksToAdd, newFks...)
-			// must remove deleted fks too
-			currentFks, _ := fkCollection.KeysForTable(tblName)
-			newFksSet := make(map[string]struct{})
-			for _, newFk := range newFks {
-				newFksSet[newFk.Name] = struct{}{}
-			}
-			for _, currentFk := range currentFks {
-				_, ok := newFksSet[currentFk.Name]
-				if !ok {
-					fksToRemove = append(fksToRemove, currentFk)
-				}
-			}
-		} else if _, ok, err := tableMap.MaybeGet(ctx, key); err != nil {
-			return nil, err
-		} else if ok {
-			me = me.Remove(key)
-			fks, _ := fkCollection.KeysForTable(tblName)
-			fksToRemove = append(fksToRemove, fks...)
-		}
-	}
-
-	m, err := me.Map(ctx)
-	if err != nil {
-		return nil, err
-	}
-	rootValSt, err := root.valueSt.Set(tablesKey, m)
-	if err != nil {
-		return nil, err
-	}
-
-	newRoot := newRootValue(root.vrw, rootValSt)
-	fkCollection.Stage(ctx, fksToAdd, fksToRemove)
-	newRoot, err = newRoot.PutForeignKeyCollection(ctx, fkCollection)
-	if err != nil {
-		return nil, err
-	}
-
-	return newRoot, nil
-}
-
 // UpdateSuperSchemasFromOther updates SuperSchemas of tblNames using SuperSchemas from other.
 func (root *RootValue) UpdateSuperSchemasFromOther(ctx context.Context, tblNames []string, other *RootValue) (*RootValue, error) {
 	newRoot := root

go/libraries/doltcore/env/actions/branch.go

@@ -24,7 +24,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/utils/set"
 	"github.com/dolthub/dolt/go/store/hash"
-	"github.com/dolthub/dolt/go/store/types"
 )
 
 var ErrAlreadyExists = errors.New("already exists")
@@ -215,72 +214,51 @@ func updateRootsForBranch(ctx context.Context, dbData env.DbData, dref ref.DoltR
 	if !hasRef {
 		return hash.Hash{}, hash.Hash{}, doltdb.ErrBranchNotFound
 	}
-
 	if ref.Equals(dbData.Rsr.CWBHeadRef(), dref) {
 		return hash.Hash{}, hash.Hash{}, doltdb.ErrAlreadyOnBranch
 	}
 
 	currRoots, err := getRoots(ctx, dbData.Ddb, dbData.Rsr, HeadRoot, WorkingRoot, StagedRoot)
-
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
 
 	cs, err := doltdb.NewCommitSpec(brName)
-
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, RootValueUnreadable{HeadRoot, err}
 	}
 
 	cm, err := dbData.Ddb.Resolve(ctx, cs, nil)
-
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, RootValueUnreadable{HeadRoot, err}
 	}
 
 	newRoot, err := cm.GetRootValue()
-
-	if err != nil {
-		return hash.Hash{}, hash.Hash{}, err
-	}
-
-	ssMap, err := newRoot.GetSuperSchemaMap(ctx)
-
-	if err != nil {
-		return hash.Hash{}, hash.Hash{}, err
-	}
-
-	fkMap, err := newRoot.GetForeignKeyCollectionMap(ctx)
-
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
 
 	conflicts := set.NewStrSet([]string{})
-	wrkTblHashes, err := tblHashesForCO(ctx, currRoots[HeadRoot], newRoot, currRoots[WorkingRoot], conflicts)
 
+	wrkTblHashes, err := moveModifiedTables(ctx, currRoots[HeadRoot], newRoot, currRoots[WorkingRoot], conflicts)
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
 
-	stgTblHashes, err := tblHashesForCO(ctx, currRoots[HeadRoot], newRoot, currRoots[StagedRoot], conflicts)
-
+	stgTblHashes, err := moveModifiedTables(ctx, currRoots[HeadRoot], newRoot, currRoots[StagedRoot], conflicts)
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
-
 	if conflicts.Size() > 0 {
 		return hash.Hash{}, hash.Hash{}, CheckoutWouldOverwrite{conflicts.AsSlice()}
 	}
 
-	wrkHash, err = writeRoot(ctx, dbData.Ddb, wrkTblHashes, ssMap, fkMap)
-
+	wrkHash, err = writeRoot(ctx, dbData.Ddb, newRoot, wrkTblHashes)
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
 
-	stgHash, err = writeRoot(ctx, dbData.Ddb, stgTblHashes, ssMap, fkMap)
-
+	stgHash, err = writeRoot(ctx, dbData.Ddb, newRoot, stgTblHashes)
 	if err != nil {
 		return hash.Hash{}, hash.Hash{}, err
 	}
@@ -345,29 +323,29 @@ func CheckoutBranchWithoutDocs(ctx context.Context, dbData env.DbData, brName st
 
 var emptyHash = hash.Hash{}
 
-func tblHashesForCO(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.RootValue, conflicts *set.StrSet) (map[string]hash.Hash, error) {
+// moveModifiedTables handles working set changes during a branch change.
+// When moving between branches, changes in the working set should travel with you.
+// Working set changes cannot be moved if the table differs between the old and new head,
+// in this case, we throw a conflict and error (as per Git).
+func moveModifiedTables(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.RootValue, conflicts *set.StrSet) (map[string]hash.Hash, error) {
 	resultMap := make(map[string]hash.Hash)
 	tblNames, err := newRoot.GetTableNames(ctx)
-
 	if err != nil {
 		return nil, err
 	}
 
 	for _, tblName := range tblNames {
 		oldHash, _, err := oldRoot.GetTableHash(ctx, tblName)
-
 		if err != nil {
 			return nil, err
 		}
 
 		newHash, _, err := newRoot.GetTableHash(ctx, tblName)
-
 		if err != nil {
 			return nil, err
 		}
 
 		changedHash, _, err := changedRoot.GetTableHash(ctx, tblName)
-
 		if err != nil {
 			return nil, err
 		}
@@ -384,7 +362,6 @@ func tblHashesForCO(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.R
 	}
 
 	tblNames, err = changedRoot.GetTableNames(ctx)
-
 	if err != nil {
 		return nil, err
 	}
@@ -392,13 +369,11 @@ func tblHashesForCO(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.R
 	for _, tblName := range tblNames {
 		if _, exists := resultMap[tblName]; !exists {
 			oldHash, _, err := oldRoot.GetTableHash(ctx, tblName)
-
 			if err != nil {
 				return nil, err
 			}
 
 			changedHash, _, err := changedRoot.GetTableHash(ctx, tblName)
-
 			if err != nil {
 				return nil, err
 			}
@@ -414,43 +389,36 @@ func tblHashesForCO(ctx context.Context, oldRoot, newRoot, changedRoot *doltdb.R
 	return resultMap, nil
 }
 
-func writeRoot(ctx context.Context, ddb *doltdb.DoltDB, tblHashes map[string]hash.Hash, ssMap types.Map, fkMap types.Map) (hash.Hash, error) {
+func writeRoot(ctx context.Context, ddb *doltdb.DoltDB, head *doltdb.RootValue, tblHashes map[string]hash.Hash) (hash.Hash, error) {
+	names, err := head.GetTableNames(ctx)
+	if err != nil {
+		return hash.Hash{}, err
+	}
+
+	var toDrop []string
+	for _, name := range names {
+		if _, ok := tblHashes[name]; !ok {
+			toDrop = append(toDrop, name)
+		}
+	}
+
+	head, err = head.RemoveTables(ctx, toDrop...)
+	if err != nil {
+		return hash.Hash{}, err
+	}
+
 	for k, v := range tblHashes {
 		if v == emptyHash {
-			delete(tblHashes, k)
+			continue
+		}
+
+		head, err = head.SetTableHash(ctx, k, v)
+		if err != nil {
+			return hash.Hash{}, err
 		}
 	}
 
-	root, err := doltdb.NewRootValue(ctx, ddb.ValueReadWriter(), tblHashes, ssMap, fkMap)
-	if err != nil {
-		if err == doltdb.ErrHashNotFound {
-			return emptyHash, errors.New("corrupted database? Can't find hash of current table")
-		}
-		return emptyHash, doltdb.ErrNomsIO
-	}
-
-	return ddb.WriteRootValue(ctx, root)
-
-}
-
-func RootsWithTable(ctx context.Context, dEnv *env.DoltEnv, table string) (RootTypeSet, error) {
-	roots, err := getRoots(ctx, dEnv.DoltDB, dEnv.RepoStateReader(), ActiveRoots...)
-
-	if err != nil {
-		return nil, err
-	}
-
-	rootsWithTable := make([]RootType, 0, len(roots))
-
-	for rt, root := range roots {
-		if has, err := root.HasTable(ctx, table); err != nil {
-			return nil, err
-		} else if has {
-			rootsWithTable = append(rootsWithTable, rt)
-		}
-	}
-
-	return NewRootTypeSet(rootsWithTable...), nil
+	return ddb.WriteRootValue(ctx, head)
 }
 
 func IsBranch(ctx context.Context, ddb *doltdb.DoltDB, str string) (bool, error) {