Merge pull request #7122 from dolthub/macneale4/checkout-B

Add support for -B flag to dolt_checkout

https://github.com/dolthub/dolt/issues/7116

go/cmd/dolt/cli/arg_parser_helpers.go

@@ -150,6 +150,7 @@ func CreateCleanArgParser() *argparser.ArgParser {
 func CreateCheckoutArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParserWithVariableArgs("checkout")
 	ap.SupportsString(CheckoutCreateBranch, "", "branch", "Create a new branch named {{.LessThan}}new_branch{{.GreaterThan}} and start it at {{.LessThan}}start_point{{.GreaterThan}}.")
+	ap.SupportsString(CreateResetBranch, "", "branch", "Similar to '-b'. Forcibly resets the branch to {{.LessThan}}start_point{{.GreaterThan}} if it exists.")
 	ap.SupportsFlag(ForceFlag, "f", "If there is any changes in working set, the force flag will wipe out the current changes and checkout the new branch.")
 	ap.SupportsString(TrackFlag, "t", "", "When creating a new branch, set up 'upstream' configuration.")
 	return ap

go/cmd/dolt/cli/flags.go

@@ -25,6 +25,7 @@ const (
 	BranchParam          = "branch"
 	CachedFlag           = "cached"
 	CheckoutCreateBranch = "b"
+	CreateResetBranch    = "B"
 	CommitFlag           = "commit"
 	CopyFlag             = "copy"
 	DateParam            = "date"

go/cmd/dolt/commands/checkout.go

@@ -106,15 +106,20 @@ func (cmd CheckoutCmd) Exec(ctx context.Context, commandStr string, args []strin
 		return 1
 	}
 
-	branchOrTrack := apr.Contains(cli.CheckoutCreateBranch) || apr.Contains(cli.TrackFlag)
+	// Argument validation in the CLI is strictly nice to have. The stored procedure will do the same, but the errors
+	// won't be as nice.
+	branchOrTrack := apr.Contains(cli.CheckoutCreateBranch) || apr.Contains(cli.CreateResetBranch) || apr.Contains(cli.TrackFlag)
 	if (branchOrTrack && apr.NArg() > 1) || (!branchOrTrack && apr.NArg() == 0) {
 		usagePrt()
 		return 1
 	}
 
+	// Branch name retrieval here is strictly for messages. dolt_checkout procedure is the authority on logic around validation.
 	var branchName string
 	if apr.Contains(cli.CheckoutCreateBranch) {
 		branchName, _ = apr.GetValue(cli.CheckoutCreateBranch)
+	} else if apr.Contains(cli.CreateResetBranch) {
+		branchName, _ = apr.GetValue(cli.CreateResetBranch)
 	} else if apr.Contains(cli.TrackFlag) {
 		if apr.NArg() > 0 {
 			usagePrt()

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

@@ -71,9 +71,14 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 		return 1, "", err
 	}
 
-	branchOrTrack := apr.Contains(cli.CheckoutCreateBranch) || apr.Contains(cli.TrackFlag)
+	newBranch, _, err := parseBranchArgs(apr)
+	if err != nil {
+		return 1, "", err
+	}
+
+	branchOrTrack := newBranch != "" || apr.Contains(cli.TrackFlag)
 	if apr.Contains(cli.TrackFlag) && apr.NArg() > 0 {
-		return 1, "", errors.New("Improper usage.")
+		return 1, "", errors.New("Improper usage. Too many arguments provided.")
 	}
 	if (branchOrTrack && apr.NArg() > 1) || (!branchOrTrack && apr.NArg() == 0) {
 		return 1, "", errors.New("Improper usage.")
@@ -90,7 +95,7 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 	if err != nil {
 		return 1, "", err
 	}
-	if apr.Contains(cli.CheckoutCreateBranch) && readOnlyDatabase {
+	if newBranch != "" && readOnlyDatabase {
 		return 1, "", fmt.Errorf("unable to create new branch in a read-only database")
 	}
 
@@ -199,6 +204,30 @@ func doDoltCheckout(ctx *sql.Context, args []string) (statusCode int, successMes
 	return 0, successMessage, nil
 }
 
+// parseBranchArgs returns the name of the new branch and whether or not it should be created forcibly. This asserts
+// that the provided branch name may not be empty, so an empty string is returned where no -b or -B flag is provided.
+func parseBranchArgs(apr *argparser.ArgParseResults) (newBranch string, createBranchForcibly bool, err error) {
+	if apr.Contains(cli.CheckoutCreateBranch) && apr.Contains(cli.CreateResetBranch) {
+		return "", false, errors.New("Improper usage. Cannot use both -b and -B.")
+	}
+
+	if newBranch, ok := apr.GetValue(cli.CheckoutCreateBranch); ok {
+		if len(newBranch) == 0 {
+			return "", false, ErrEmptyBranchName
+		}
+		return newBranch, false, nil
+	}
+
+	if newBranch, ok := apr.GetValue(cli.CreateResetBranch); ok {
+		if len(newBranch) == 0 {
+			return "", false, ErrEmptyBranchName
+		}
+		return newBranch, true, nil
+	}
+
+	return "", false, nil
+}
+
 // isReadOnlyDatabase returns true if the named database is a read-only database. An error is returned
 // if any issues are encountered while looking up the named database.
 func isReadOnlyDatabase(ctx *sql.Context, dbName string) (bool, error) {
@@ -346,14 +375,20 @@ func checkoutNewBranch(ctx *sql.Context, dbName string, dbData env.DbData, apr *
 		newBranchName = remoteBranchName
 	}
 
-	if newBranch, ok := apr.GetValue(cli.CheckoutCreateBranch); ok {
-		if len(newBranch) == 0 {
-			return "", "", ErrEmptyBranchName
-		}
-		newBranchName = newBranch
+	// A little wonky behavior here. parseBranchArgs is actually called twice because in this procedure we pass around
+	// the parse results, but we also needed to parse the -b and -B flags in the main procedure. It ended up being
+	// a little cleaner to just call it again here than to pass the results around.
+	var createBranchForcibly bool
+	var optionBBranch string
+	optionBBranch, createBranchForcibly, err = parseBranchArgs(apr)
+	if err != nil {
+		return "", "", err
+	}
+	if optionBBranch != "" {
+		newBranchName = optionBBranch
 	}
 
-	err = actions.CreateBranchWithStartPt(ctx, dbData, newBranchName, startPt, false, rsc)
+	err = actions.CreateBranchWithStartPt(ctx, dbData, newBranchName, startPt, createBranchForcibly, rsc)
 	if err != nil {
 		return "", "", err
 	}

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

@@ -2387,6 +2387,101 @@ var DoltCheckoutScripts = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "dolt_checkout with new branch forcefully",
+		SetUpScript: []string{
+			"create table t (s varchar(5) primary key);",
+			"insert into t values ('foo');",
+			"call dolt_commit('-Am', 'commit main~2');", // will be main~2
+			"insert into t values ('bar');",
+			"call dolt_commit('-Am', 'commit main~1');", // will be main~1
+			"insert into t values ('baz');",
+			"call dolt_commit('-Am', 'commit main');", // will be main~1
+			"call dolt_branch('testbr', 'main~1');",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:            "call dolt_checkout('-B', 'testbr', 'main~2');",
+				SkipResultsCheck: true,
+			},
+			{
+				Query:    "select active_branch();",
+				Expected: []sql.Row{{"testbr"}},
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"foo"}},
+			},
+			{
+				Query:            "call dolt_checkout('main');",
+				SkipResultsCheck: true,
+			},
+			{
+				Query:    "select active_branch();",
+				Expected: []sql.Row{{"main"}},
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"bar"}, {"baz"}, {"foo"}},
+			},
+			{
+				Query:            "call dolt_checkout('-B', 'testbr', 'main~1');",
+				SkipResultsCheck: true,
+			},
+			{
+				Query:    "select active_branch();",
+				Expected: []sql.Row{{"testbr"}},
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"bar"}, {"foo"}},
+			},
+		},
+	},
+	{
+		Name: "dolt_checkout with new branch forcefully with dirty working set",
+		SetUpScript: []string{
+			"create table t (s varchar(5) primary key);",
+			"insert into t values ('foo');",
+			"call dolt_commit('-Am', 'commit main~2');", // will be main~2
+			"insert into t values ('bar');",
+			"call dolt_commit('-Am', 'commit main~1');", // will be main~1
+			"insert into t values ('baz');",
+			"call dolt_commit('-Am', 'commit main');", // will be main~1
+			"call dolt_checkout('-b', 'testbr', 'main~1');",
+			"insert into t values ('qux');",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:    "select active_branch();",
+				Expected: []sql.Row{{"testbr"}},
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"bar"}, {"foo"}, {"qux"}}, // Dirty working set
+			},
+			{
+				Query:            "call dolt_checkout('main');",
+				SkipResultsCheck: true,
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"bar"}, {"baz"}, {"foo"}},
+			},
+			{
+				Query:            "call dolt_checkout('-B', 'testbr', 'main~1');",
+				SkipResultsCheck: true,
+			},
+			{
+				Query:    "select active_branch();",
+				Expected: []sql.Row{{"testbr"}},
+			},
+			{
+				Query:    "select * from t order by s;",
+				Expected: []sql.Row{{"bar"}, {"foo"}}, // Dirty working set was forcefully overwritten
+			},
+		},
+	},
 	{
 		Name: "dolt_checkout mixed with USE statements",
 		SetUpScript: []string{
@@ -2768,6 +2863,15 @@ var DoltCheckoutReadOnlyScripts = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "dolt checkout -B returns an error for read-only databases",
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query:          "call dolt_checkout('-B', 'newBranch');",
+				ExpectedErrStr: "unable to create new branch in a read-only database",
+			},
+		},
+	},
 }
 
 var DoltInfoSchemaScripts = []queries.ScriptTest{

integration-tests/bats/checkout.bats

@@ -314,6 +314,46 @@ SQL
     [[ ! "$output" =~ "4" ]] || false
 }
 
+@test "checkout: -B flag will forcefully reset an existing branch" {
+    dolt sql -q 'create table test (id int primary key);'
+    dolt sql -q 'insert into test (id) values (89012);'
+    dolt commit -Am 'first change.'
+    dolt sql -q 'insert into test (id) values (76543);'
+    dolt commit -Am 'second change.'
+
+    dolt checkout -b testbr main~1
+    run dolt sql -q "select * from test;"
+    [[ "$output" =~ "89012" ]] || false
+    [[ ! "$output" =~ "76543" ]] || false
+
+    # make a change to the branch which we'll lose
+    dolt sql -q 'insert into test (id) values (19283);'
+    dolt commit -Am 'change to testbr.'
+
+    dolt checkout main
+    dolt checkout -B testbr main
+    run dolt sql -q "select * from test;"
+    [[ "$output" =~ "89012" ]] || false
+    [[ "$output" =~ "76543" ]] || false
+    [[ ! "$output" =~ "19283" ]] || false
+}
+
+@test "checkout: -B will create a branch that does not exist" {
+    dolt sql -q 'create table test (id int primary key);'
+    dolt sql -q 'insert into test (id) values (89012);'
+    dolt commit -Am 'first change.'
+    dolt sql -q 'insert into test (id) values (76543);'
+    dolt commit -Am 'second change.'
+
+    dolt checkout -B testbr main~1
+    run dolt sql -q "select * from test;"
+    [[ "$output" =~ "89012" ]] || false
+    [[ ! "$output" =~ "76543" ]] || false
+}
+
+
+
+
 @test "checkout: attempting to checkout a detached head shows a suggestion instead" {
   dolt sql -q "create table test (id int primary key);"
   dolt add .