Merge pull request #8446 from dolthub/zachmu/checkout

Implemented checking out a table or tables from another commit

go/cmd/dolt/commands/checkout.go

@@ -44,13 +44,17 @@ dolt checkout {{.LessThan}}branch{{.GreaterThan}}
    To prepare for working on {{.LessThan}}branch{{.GreaterThan}}, switch to it by updating the index and the tables in the working tree, and by pointing HEAD at the branch. Local modifications to the tables in the working
    tree are kept, so that they can be committed to the {{.LessThan}}branch{{.GreaterThan}}.
 
+dolt checkout {{.LessThan}}commit{{.GreaterThan}} [--] {{.LessThan}}table{{.GreaterThan}}...
+	 Specifying table names after a commit reference (branch, commit hash, tag, etc.) updates the working set to match that commit for one or more tables, but keeps the current branch. Local modifications to the tables named will be overwritten by their versions in the commit named.
+
 dolt checkout -b {{.LessThan}}new_branch{{.GreaterThan}} [{{.LessThan}}start_point{{.GreaterThan}}]
    Specifying -b causes a new branch to be created as if dolt branch were called and then checked out.
 
 dolt checkout {{.LessThan}}table{{.GreaterThan}}...
   To update table(s) with their values in HEAD `,
 	Synopsis: []string{
 		`{{.LessThan}}branch{{.GreaterThan}}`,
+		`{{.LessThan}}commit{{.GreaterThan}} [--] {{.LessThan}}table{{.GreaterThan}}...`,
 		`{{.LessThan}}table{{.GreaterThan}}...`,
 		`-b {{.LessThan}}new-branch{{.GreaterThan}} [{{.LessThan}}start-point{{.GreaterThan}}]`,
 		`--track {{.LessThan}}remote{{.GreaterThan}}/{{.LessThan}}branch{{.GreaterThan}}`,

go/libraries/doltcore/sqle/dprocedures/dolt_checkout.go

@@ -28,6 +28,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/env/actions"
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/resolve"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
 	"github.com/dolthub/dolt/go/store/hash"
 )
@@ -80,7 +81,7 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 	if apr.Contains(cli.TrackFlag) && apr.NArg() > 0 {
 		return 1, "", errors.New("Improper usage. Too many arguments provided.")
 	}
-	if (branchOrTrack && apr.NArg() > 1) || (!branchOrTrack && apr.NArg() == 0) {
+	if !branchOrTrack && apr.NArg() == 0 {
 		return 1, "", errors.New("Improper usage.")
 	}
 
@@ -122,10 +123,25 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 	if err != nil {
 		return 1, "", err
 	}
-	if !isModification {
+	if !isModification && apr.NArg() == 1 {
 		return 0, fmt.Sprintf("Already on branch '%s'", branchName), nil
 	}
 
+	// Handle dolt_checkout HEAD~3 -- table1 table2 table3
+	if apr.NArg() > 1 {
+		database := ctx.GetCurrentDatabase()
+		if database == "" {
+			return 1, "", sql.ErrNoDatabaseSelected.New()
+		}
+		err = checkoutTablesFromCommit(ctx, database, branchName, apr.Args[1:])
+		if err != nil {
+			return 0, "", err
+		}
+
+		dsess.WaitForReplicationController(ctx, rsc)
+		return 0, "", nil
+	}
+
 	// Check if user wants to checkout branch.
 	if isBranch, err := actions.IsBranch(ctx, dbData.Ddb, branchName); err != nil {
 		return 1, "", err
@@ -190,7 +206,7 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 		return 0, "", err
 	}
 
-	err = checkoutTables(ctx, roots, currentDbName, apr.Args)
+	err = checkoutTablesFromHead(ctx, roots, currentDbName, apr.Args)
 	if err != nil && apr.NArg() == 1 {
 		upstream, err := checkoutRemoteBranch(ctx, dSess, currentDbName, dbData, branchName, apr, &rsc)
 		if err != nil {
@@ -479,6 +495,82 @@ func checkoutExistingBranch(ctx *sql.Context, dbName string, branchName string,
 	return nil
 }
 
+// checkoutTablesFromCommit checks out the tables named from the branch named and overwrites those tables in the
+// staged and working roots.
+func checkoutTablesFromCommit(
+	ctx *sql.Context,
+	databaseName string,
+	commitRef string,
+	tables []string,
+) error {
+	dSess := dsess.DSessFromSess(ctx.Session)
+	dbData, ok := dSess.GetDbData(ctx, databaseName)
+	if !ok {
+		return fmt.Errorf("Could not load database %s", ctx.GetCurrentDatabase())
+	}
+
+	currentBranch, err := dSess.GetBranch()
+	if err != nil {
+		return err
+	}
+	if currentBranch == "" {
+		return fmt.Errorf("error: no current branch")
+	}
+	currentBranchRef := ref.NewBranchRef(currentBranch)
+
+	cs, err := doltdb.NewCommitSpec(commitRef)
+	if err != nil {
+		return err
+	}
+
+	headCommit, err := dbData.Ddb.Resolve(ctx, cs, currentBranchRef)
+	if err != nil {
+		return err
+	}
+	cm, ok := headCommit.ToCommit()
+	if !ok {
+		return fmt.Errorf("HEAD is not a commit")
+	}
+
+	headRoot, err := cm.GetRootValue(ctx)
+	if err != nil {
+		return err
+	}
+
+	ws, err := dSess.WorkingSet(ctx, databaseName)
+	if err != nil {
+		return err
+	}
+
+	var tableNames []doltdb.TableName
+	if len(tables) == 1 && tables[0] == "." {
+		tableNames, err = doltdb.UnionTableNames(ctx, ws.WorkingRoot())
+		if err != nil {
+			return err
+		}
+	} else {
+		tableNames = make([]doltdb.TableName, len(tables))
+		for i, table := range tables {
+			// TODO: we should allow schema-qualified table names here as well
+			name, _, tableExistsInHead, err := resolve.Table(ctx, headRoot, table)
+			if err != nil {
+				return err
+			}
+			if !tableExistsInHead {
+				return fmt.Errorf("table %s does not exist in %s", table, commitRef)
+			}
+			tableNames[i] = name
+		}
+	}
+
+	newRoot, err := actions.MoveTablesBetweenRoots(ctx, tableNames, headRoot, ws.WorkingRoot())
+	if err != nil {
+		return err
+	}
+
+	return dSess.SetWorkingSet(ctx, databaseName, ws.WithStagedRoot(newRoot).WithWorkingRoot(newRoot))
+}
+
 // doGlobalCheckout implements the behavior of the `dolt checkout` command line, moving the working set into
 // the new branch and persisting the checked-out branch into future sessions
 func doGlobalCheckout(ctx *sql.Context, branchName string, isForce bool, isNewBranch bool) error {
@@ -490,7 +582,9 @@ func doGlobalCheckout(ctx *sql.Context, branchName string, isForce bool, isNewBr
 	return nil
 }
 
-func checkoutTables(ctx *sql.Context, roots doltdb.Roots, name string, tables []string) error {
+// checkoutTablesFromHead checks out the tables named from the current head and overwrites those tables in the
+// working root. The working root is then set as the new staged root.
+func checkoutTablesFromHead(ctx *sql.Context, roots doltdb.Roots, name string, tables []string) error {
 	tableNames := make([]doltdb.TableName, len(tables))
 
 	for i, table := range tables {

go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go

@@ -106,39 +106,7 @@ func TestSchemaOverrides(t *testing.T) {
 func TestSingleScript(t *testing.T) {
 	t.Skip()
 
-	var scripts = []queries.ScriptTest{
-		{
-			Name: "create database in a transaction",
-			SetUpScript: []string{
-				"START TRANSACTION",
-				"CREATE DATABASE test",
-			},
-			Assertions: []queries.ScriptTestAssertion{
-				{
-					Query:            "USE test",
-					SkipResultsCheck: true,
-				},
-				{
-					Query:            "CREATE TABLE foo (bar INT)",
-					SkipResultsCheck: true,
-				},
-				{
-					Query:            "USE mydb",
-					SkipResultsCheck: true,
-				},
-				{
-					Query:            "INSERT INTO test.foo VALUES (1)",
-					SkipResultsCheck: true,
-				},
-				{
-					Query: "SELECT * FROM test.foo",
-					Expected: []sql.Row{
-						{1},
-					},
-				},
-			},
-		},
-	}
+	var scripts = []queries.ScriptTest{}
 
 	for _, script := range scripts {
 		harness := newDoltHarness(t)

go/libraries/doltcore/sqle/enginetest/dolt_queries.go

@@ -3174,6 +3174,169 @@ var DoltCheckoutScripts = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "Checkout tables from commit",
+		SetUpScript: []string{
+			"create table t1 (a int primary key, b int);",
+			"create table t2 (a int primary key, b int);",
+			"call dolt_commit('-Am', 'creating tables');",
+			"call dolt_tag('tag1');",
+			"insert into t1 values (1, 1);",
+			"insert into t2 values (2, 2);",
+			"call dolt_commit('-Am', 'one row in each table');",
+			"call dolt_branch('b1');",
+			"insert into t1 values (3, 3);",
+			"insert into t2 values (4, 4);",
+			"call dolt_commit('-Am', 'two rows in each table');",
+			"insert into t1 values (5, 5);",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "call dolt_checkout('HEAD~', '--', 't1')",
+				Expected: []sql.Row{
+					{0, ""},
+				},
+			},
+			{
+				Query: "select * from t1 order by 1",
+				Expected: []sql.Row{
+					{1, 1},
+				},
+			},
+			{
+				Query: "select * from t2 order by 1",
+				Expected: []sql.Row{
+					{2, 2},
+					{4, 4},
+				},
+			},
+			{
+				Query: "select * from dolt_status",
+				Expected: []sql.Row{
+					{"t1", true, "modified"},
+				},
+			},
+			{
+				Query: "call dolt_reset('--hard')",
+				Expected: []sql.Row{
+					{0},
+				},
+			},
+			{
+				Query: "call dolt_checkout('HEAD~', '--', 't2')",
+				Expected: []sql.Row{
+					{0, ""},
+				},
+			},
+			{
+				Query: "select * from t1 order by 1",
+				Expected: []sql.Row{
+					{1, 1},
+					{3, 3},
+				},
+			},
+			{
+				Query: "select * from t2 order by 1",
+				Expected: []sql.Row{
+					{2, 2},
+				},
+			},
+			{
+				Query: "select * from dolt_status",
+				Expected: []sql.Row{
+					{"t2", true, "modified"},
+				},
+			},
+			{
+				Query: "call dolt_reset('--hard')",
+				Expected: []sql.Row{
+					{0},
+				},
+			},
+			{
+				Query: "call dolt_checkout('b1', 't2', 't1')",
+				Expected: []sql.Row{
+					{0, ""},
+				},
+			},
+			{
+				Query: "select * from t1 order by 1",
+				Expected: []sql.Row{
+					{1, 1},
+				},
+			},
+			{
+				Query: "select * from t2 order by 1",
+				Expected: []sql.Row{
+					{2, 2},
+				},
+			},
+			{
+				Query: "call dolt_reset('--hard')",
+				Expected: []sql.Row{
+					{0},
+				},
+			},
+			{
+				Query: "call dolt_checkout('tag1', '.')",
+				Expected: []sql.Row{
+					{0, ""},
+				},
+			},
+			{
+				Query:    "select * from t1 order by 1",
+				Expected: []sql.Row{},
+			},
+			{
+				Query:    "select * from t2 order by 1",
+				Expected: []sql.Row{},
+			},
+			{
+				Query: "select * from dolt_status",
+				Expected: []sql.Row{
+					{"t1", true, "modified"},
+					{"t2", true, "modified"},
+				},
+			},
+			{
+				Query: "call dolt_reset('--hard')",
+				Expected: []sql.Row{
+					{0},
+				},
+			},
+			{
+				Query:    "SET @commit1 = (select commit_hash from dolt_log order by date desc limit 1);",
+				Expected: []sql.Row{{}},
+			},
+			{
+				Query: "call dolt_checkout(@commit1, 't1')",
+				Expected: []sql.Row{
+					{0, ""},
+				},
+			},
+			{
+				Query: "select * from t1 order by 1",
+				Expected: []sql.Row{
+					{1, 1},
+					{3, 3},
+				},
+			},
+			{
+				Query: "call dolt_reset('--hard')",
+				Expected: []sql.Row{
+					{0},
+				},
+			},
+			{
+				Query:          "call dolt_checkout('nosuchbranch', 't1')",
+				ExpectedErrStr: "branch not found: nosuchbranch",
+			},
+			{
+				Query:          "call dolt_checkout('HEAD', 't3')",
+				ExpectedErrStr: "table t3 does not exist in HEAD",
+			},
+		},
+	},
 }
 
 var DoltCheckoutReadOnlyScripts = []queries.ScriptTest{

go/libraries/utils/argparser/parser_test.go

@@ -150,6 +150,13 @@ func TestArgParser(t *testing.T) {
 			map[string]string{"param": "value", "optional": ""},
 			[]string{},
 		},
+		{
+			NewArgParserWithVariableArgs("test").SupportsString("param", "p", "", ""),
+			[]string{"--param", "value", "arg1", "--", "table1", "table2"},
+			nil,
+			map[string]string{"param": "value"},
+			[]string{"arg1", "table1", "table2"},
+		},
 	}
 
 	for _, test := range tests {