Merge pull request #9433 from dolthub/nathan/upstream

--set-upstream and improved --track for dolt branch

go/cmd/dolt/cli/arg_parser_helpers.go

@@ -244,14 +244,8 @@ func createTracklessBranchArgParser() *argparser.ArgParser {
 
 func CreateBranchArgParser() *argparser.ArgParser {
 	ap := createTracklessBranchArgParser()
-	ap.SupportsString(TrackFlag, "t", "", "When creating a new branch, set up 'upstream' configuration.")
-
-	return ap
-}
-
-func CreateBranchArgParserWithNoTrackValue() *argparser.ArgParser {
-	ap := createTracklessBranchArgParser()
-	ap.SupportsFlag(TrackFlag, "t", "When creating a new branch, set up 'upstream' configuration.")
+	ap.SupportsFlag(TrackFlag, "t", "Set up upstream configuration for a branch. Uses current branch as default")
+	ap.SupportsString(SetUpstreamToFlag, "u", "", "Set upstream configuration for a branch.")
 
 	return ap
 }

go/cmd/dolt/cli/flags.go

@@ -68,6 +68,7 @@ const (
 	QuietFlag            = "quiet"
 	RemoteParam          = "remote"
 	SetUpstreamFlag      = "set-upstream"
+	SetUpstreamToFlag    = "set-upstream-to"
 	ShallowFlag          = "shallow"
 	ShowIgnoredFlag      = "ignored"
 	ShowSignatureFlag    = "show-signature"

go/cmd/dolt/commands/branch.go

@@ -140,6 +140,8 @@ func (cmd BranchCmd) Exec(ctx context.Context, commandStr string, args []string,
 		return printCurrentBranch(sqlCtx, queryEngine)
 	case apr.Contains(datasetsFlag):
 		return printAllDatasets(sqlCtx, dEnv)
+	case apr.ContainsAny(cli.SetUpstreamToFlag, cli.TrackFlag):
+		return updateUpstream(sqlCtx, queryEngine, apr, args)
 	case apr.NArg() > 0:
 		return createBranch(sqlCtx, queryEngine, apr, args, usage)
 	default:
@@ -317,35 +319,57 @@ func generateBranchSql(args []string) (string, error) {
 	return dbr.InterpolateForDialect(buffer.String(), queryValues, dialect.MySQL)
 }
 
-func createBranch(sqlCtx *sql.Context, queryEngine cli.Queryist, apr *argparser.ArgParseResults, args []string, usage cli.UsagePrinter) int {
-
-	trackVal, setTrackUpstream := apr.GetValue(cli.TrackFlag)
-	if setTrackUpstream {
-		if trackVal == "inherit" {
-			return HandleVErrAndExitCode(errhand.BuildDError("--track='inherit' is not supported yet").Build(), usage)
+func updateUpstream(sqlCtx *sql.Context, queryEngine cli.Queryist, apr *argparser.ArgParseResults, args []string) int {
+	var branchName, upstreamBranch string
+	var err error
+	if apr.NArg() == 0 {
+		if apr.Contains(cli.TrackFlag) {
+			cli.PrintErrln("error: must specify branch to be created when using --track")
 		}
-
-		if trackVal == "direct" {
-			if apr.NArg() != 2 {
-				return HandleVErrAndExitCode(errhand.BuildDError("invalid arguments").Build(), usage)
-			}
-		} else {
-			// --track did not have an associated parameter; we parsed a positional arg as its value.
-			// There is no way to determine what position that arg was supposed to be in.
-
-			// --track accepts a parameter but can also be passed in on its own.
-			// We can determine which based on the number of arguments.
-			// We initially parsed args assuming that --track accepted a parameter,
-			// but now we have to parse the args again with it as a flag instead.
-
-			var err error
-			apr, err = cli.CreateBranchArgParserWithNoTrackValue().Parse(args)
-			if err != nil {
-				return HandleVErrAndExitCode(errhand.VerboseErrorFromError(err), usage)
-			}
+		branchName, err = getActiveBranchName(sqlCtx, queryEngine)
+		if err != nil {
+			cli.PrintErrln("error: failed to get current branch from database")
+			return 1
 		}
+	} else {
+		branchName = apr.Arg(0)
 	}
 
+	if apr.Contains(cli.TrackFlag) && apr.Contains(cli.SetUpstreamToFlag) {
+		cli.PrintErrln(fmt.Sprintf("error: --%s and --%s are mutually exclusive options.", cli.SetUpstreamToFlag, cli.TrackFlag))
+		return 1
+	} else if apr.Contains(cli.TrackFlag) {
+		if apr.NArg() == 1 { // If we're only given the branch, we'll presume the current branch will be the upstream
+			upstreamBranch, err = getActiveBranchName(sqlCtx, queryEngine)
+			if err != nil {
+				cli.PrintErrln("error: If --track is used, you must specify the branch to modify, or a new branch.")
+				return 1
+			}
+		} else if apr.NArg() == 2 {
+			upstreamBranch = apr.Arg(1)
+		} else {
+			cli.PrintErrln("error: If --track is used, you must specify the branch to modify, or a new branch.")
+			return 1
+		}
+	} else if apr.Contains(cli.SetUpstreamToFlag) {
+		if apr.NArg() > 2 {
+			cli.PrintErrln("error: If --set-upstream-to is used, you must specify the branch to be modified, then upstream branch.")
+			return 1
+		}
+		upstreamBranch, _ = apr.GetValue(cli.SetUpstreamToFlag)
+	}
+
+	res := callStoredProcedure(sqlCtx, queryEngine, args)
+	if res != 0 {
+		return res
+	}
+	cli.Printf("branch '%s' set up to track '%s'\n", branchName, upstreamBranch)
+	return 0
+}
+
+func createBranch(sqlCtx *sql.Context, queryEngine cli.Queryist, apr *argparser.ArgParseResults, args []string, usage cli.UsagePrinter) int {
+	remoteName, useUpstream := apr.GetValue(cli.SetUpstreamToFlag)
+
 	if apr.NArg() != 1 && apr.NArg() != 2 {
 		usage()
 		return 1
@@ -378,8 +402,8 @@ func createBranch(sqlCtx *sql.Context, queryEngine cli.Queryist, apr *argparser.
 		return result
 	}
 
-	if apr.Contains(cli.TrackFlag) {
-		cli.Printf("branch '%s' set up to track '%s'\n", apr.Arg(0), apr.Arg(1))
+	if useUpstream {
+		cli.Printf("branch '%s' set up to track '%s'\n", apr.Arg(0), remoteName)
 	}
 
 	return 0

go/cmd/dolt/commands/status.go

@@ -178,7 +178,7 @@ func createPrintData(err error, queryist cli.Queryist, sqlCtx *sql.Context, show
 		return nil, err
 	}
 
-	ahead, behind, err := getRemoteInfo(queryist, sqlCtx, branchName, remoteName, remoteBranchName, currentBranchCommit)
+	ahead, behind, err := getUpstreamInfo(queryist, sqlCtx, branchName, remoteName, remoteBranchName, currentBranchCommit)
 	if err != nil {
 		return nil, err
 	}
@@ -334,31 +334,38 @@ func getConflictedTables(statusRows []sql.Row) map[string]bool {
 	return conflictedTables
 }
 
-func getRemoteInfo(queryist cli.Queryist, sqlCtx *sql.Context, branchName string, remoteName string, remoteBranchName string, currentBranchCommit string) (ahead int64, behind int64, err error) {
+func getUpstreamInfo(queryist cli.Queryist, sqlCtx *sql.Context, branchName string, remoteName string, upstreamBranchName string, currentBranchCommit string) (ahead int64, behind int64, err error) {
 	ahead = 0
 	behind = 0
-	if len(remoteName) > 0 {
+	if len(remoteName) > 0 || len(upstreamBranchName) > 0 {
 		// get remote branch
-		remoteBranchRef := fmt.Sprintf("remotes/%s/%s", remoteName, remoteBranchName)
-		q := fmt.Sprintf("select name, hash from dolt_remote_branches where name = '%s';", remoteBranchRef)
-		remoteBranches, err := GetRowsForSql(queryist, sqlCtx, q)
+		var q string
+		var upstreamBranchRef string
+		if remoteName != "" {
+			upstreamBranchRef = fmt.Sprintf("remotes/%s/%s", remoteName, upstreamBranchName)
+			q = fmt.Sprintf("select name, hash from dolt_remote_branches where name = '%s';", upstreamBranchRef)
+		} else if upstreamBranchName != "" {
+			q = fmt.Sprintf("select name, hash from dolt_branches where name = '%s'", upstreamBranchName)
+		}
+
+		upstreamBranches, err := GetRowsForSql(queryist, sqlCtx, q)
 		if err != nil {
 			return ahead, behind, err
 		}
-		if len(remoteBranches) > 1 {
-			return ahead, behind, fmt.Errorf("runtime error: too many results returned for remote branch %s", remoteBranchRef)
-		} else if len(remoteBranches) == 0 {
+		if len(upstreamBranches) > 1 {
+			return ahead, behind, fmt.Errorf("runtime error: too many results returned for upstream branch %s", upstreamBranchRef)
+		} else if len(upstreamBranches) == 0 {
 			return ahead, behind, nil
 		}
-		remoteBranchCommit := remoteBranches[0][1].(string)
+		upstreamBranchCommit := upstreamBranches[0][1].(string)
 
-		q = fmt.Sprintf("call dolt_count_commits('--from', '%s', '--to', '%s')", currentBranchCommit, remoteBranchCommit)
+		q = fmt.Sprintf("call dolt_count_commits('--from', '%s', '--to', '%s')", currentBranchCommit, upstreamBranchCommit)
 		rows, err := GetRowsForSql(queryist, sqlCtx, q)
 		if err != nil {
 			return ahead, behind, err
 		}
 		if len(rows) != 1 {
-			return ahead, behind, fmt.Errorf("could not count commits between %s and %s", currentBranchCommit, remoteBranchCommit)
+			return ahead, behind, fmt.Errorf("could not count commits between %s and %s", currentBranchCommit, upstreamBranchCommit)
 		}
 		aheadDb := rows[0][0]
 		behindDb := rows[0][1]
@@ -496,11 +503,14 @@ func printEverything(data *printData) error {
 	cli.Printf(branchHeader, data.branchName)
 
 	// remote info
-	if data.remoteName != "" {
+	if data.remoteName != "" || data.remoteBranchName != "" {
 		ahead := data.ahead
 		behind := data.behind
-		remoteBranchRef := fmt.Sprintf("%s/%s", data.remoteName, data.remoteBranchName)
-
+		var remoteBranchRef string
+		if data.remoteName != "" {
+			remoteBranchRef = fmt.Sprintf("%s/", data.remoteName)
+		}
+		remoteBranchRef += data.remoteBranchName
 		if ahead > 0 && behind > 0 {
 			cli.Printf(`Your branch and '%s' have diverged,
 and have %v and %v different commits each, respectively.

go/libraries/doltcore/sqle/dprocedures/dolt_branch.go

@@ -79,6 +79,8 @@ func doDoltBranch(ctx *sql.Context, args []string) (int, error) {
 		err = renameBranch(ctx, dbData, apr, dSess, dbName, &rsc)
 	case apr.Contains(cli.DeleteFlag), apr.Contains(cli.DeleteForceFlag):
 		err = deleteBranches(ctx, dbData, apr, dSess, dbName, &rsc)
+	case apr.Contains(cli.SetUpstreamToFlag):
+		err = setBranchUpstream(ctx, dbData, apr, &rsc)
 	default:
 		err = createNewBranch(ctx, dbData, apr, &rsc)
 	}
@@ -301,7 +303,7 @@ func shouldAllowDefaultBranchDeletion(ctx *sql.Context) bool {
 	return userVar != nil
 }
 
-// validateBranchNotActiveInAnySessions returns an error if the specified branch is currently
+// validateBranchNotActiveInAnySession returns an error if the specified branch is currently
 // selected as the active branch for any active server sessions.
 func validateBranchNotActiveInAnySession(ctx *sql.Context, branchName string) error {
 	currentDbName := ctx.GetCurrentDatabase()
@@ -362,6 +364,50 @@ func loadConfig(ctx *sql.Context) *env.DoltCliConfig {
 	return dEnv.Config
 }
 
+func setBranchUpstream(ctx *sql.Context, dbData env.DbData[*sql.Context], apr *argparser.ArgParseResults, rsc *doltdb.ReplicationStatusController) error {
+	var branchName, fullRemote string
+	var err error
+
+	if apr.NArg() == 0 {
+		branchName, err = currentBranch(ctx)
+		if err != nil {
+			return err
+		}
+	} else {
+		branchName = apr.Arg(0)
+		ok, err := actions.IsBranch(ctx, dbData.Ddb, branchName)
+		if err != nil {
+			return err
+		}
+		if !ok {
+			return createNewBranch(ctx, dbData, apr, rsc)
+		}
+	}
+
+	if apr.NArg() > 2 {
+		return InvalidArgErr
+	}
+	fullRemote, ok := apr.GetValue(cli.SetUpstreamToFlag)
+	if !ok {
+		return fmt.Errorf("could not parse upstream value for dolt branch")
+	}
+
+	remoteName, remoteBranch, err := validateTracking(ctx, dbData, fullRemote, branchName)
+	if err != nil {
+		return err
+	}
+
+	refSpec, err := ref.ParseRefSpecForRemote(remoteName, remoteBranch)
+	if err != nil {
+		return err
+	}
+	err = env.SetRemoteUpstreamForRefSpec(dbData.Rsw, refSpec, remoteName, ref.NewBranchRef(branchName))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func createNewBranch(ctx *sql.Context, dbData env.DbData[*sql.Context], apr *argparser.ArgParseResults, rsc *doltdb.ReplicationStatusController) error {
 	if apr.NArg() == 0 || apr.NArg() > 2 {
 		return InvalidArgErr
@@ -372,46 +418,45 @@ func createNewBranch(ctx *sql.Context, dbData env.DbData[*sql.Context], apr *arg
 	if len(branchName) == 0 {
 		return EmptyBranchNameErr
 	}
-	if apr.NArg() == 2 {
-		startPt = apr.Arg(1)
-		if len(startPt) == 0 {
-			return InvalidArgErr
-		}
-	}
 
 	var remoteName, remoteBranch string
 	var refSpec ref.RefSpec
 	var err error
-	trackVal, setTrackUpstream := apr.GetValue(cli.TrackFlag)
-	if setTrackUpstream {
-		if trackVal == "inherit" {
-			return fmt.Errorf("--track='inherit' is not supported yet")
-		} else if trackVal == "direct" && apr.NArg() != 2 {
-			return InvalidArgErr
-		}
-
-		if apr.NArg() == 2 {
-			// branchName and startPt are already set
-			remoteName, remoteBranch = actions.ParseRemoteBranchName(startPt)
-			refSpec, err = ref.ParseRefSpecForRemote(remoteName, remoteBranch)
+	var trackVal string
+	var setTrackUpstream bool
+	if apr.Contains(cli.SetUpstreamToFlag) && apr.Contains(cli.TrackFlag) {
+		return fmt.Errorf("error: --%s and --%s are mutually exclusive options.", cli.SetUpstreamToFlag, cli.TrackFlag)
+	} else if apr.Contains(cli.SetUpstreamToFlag) {
+		trackVal, setTrackUpstream = apr.GetValue(cli.SetUpstreamToFlag)
+	} else if apr.Contains(cli.TrackFlag) {
+		if apr.NArg() == 1 {
+			trackVal, err = currentBranch(ctx)
 			if err != nil {
 				return err
 			}
+		} else if apr.NArg() == 2 {
+			trackVal = apr.Arg(1)
 		} else {
-			// if track option is defined with no value,
-			// the track value can either be starting point name OR branch name
-			startPt = trackVal
-			remoteName, remoteBranch = actions.ParseRemoteBranchName(startPt)
-			refSpec, err = ref.ParseRefSpecForRemote(remoteName, remoteBranch)
-			if err != nil {
-				branchName = trackVal
-				startPt = apr.Arg(0)
-				remoteName, remoteBranch = actions.ParseRemoteBranchName(startPt)
-				refSpec, err = ref.ParseRefSpecForRemote(remoteName, remoteBranch)
-				if err != nil {
-					return err
-				}
-			}
+			return InvalidArgErr
+		}
+		setTrackUpstream = true
+	}
+
+	if apr.NArg() == 2 && !apr.Contains(cli.TrackFlag) {
+		startPt = apr.Arg(1)
+	}
+	if len(startPt) == 0 {
+		return InvalidArgErr
+	}
+
+	if setTrackUpstream {
+		remoteName, remoteBranch, err = validateTracking(ctx, dbData, trackVal, branchName)
+		if err != nil {
+			return err
+		}
+		refSpec, err = ref.ParseRefSpecForRemote(remoteName, remoteBranch)
+		if err != nil {
+			return err
 		}
 	}
 
@@ -419,7 +464,6 @@ func createNewBranch(ctx *sql.Context, dbData env.DbData[*sql.Context], apr *arg
 	if err != nil {
 		return err
 	}
-
 	err = actions.CreateBranchWithStartPt(ctx, dbData, branchName, startPt, apr.Contains(cli.ForceFlag), rsc)
 	if err != nil {
 		return err
@@ -485,3 +529,42 @@ func copyABranch(ctx *sql.Context, dbData env.DbData[*sql.Context], srcBr string
 
 	return nil
 }
+
+// validateTracking takes in a full remote path, like `origin/main`, or a branch like 'main' and verifies that it's a valid upstream.
+// It errors out if it can't find the remote, or if it can't find the given branch. It returns the remote name and upstream branch.
+func validateTracking(ctx *sql.Context, dbData env.DbData[*sql.Context], maybeUpstream string, selectedBranch string) (string, string, error) {
+
+	//First we check if it's a remote
+	headRef, err := dbData.Rsr.CWBHeadRef(ctx)
+	if err != nil {
+		return "", "", err
+	}
+	cs, err := doltdb.NewCommitSpec(maybeUpstream)
+	if err != nil {
+		return "", "", err
+	}
+	_, err = dbData.Ddb.Resolve(ctx, cs, headRef)
+
+	if err == nil { // Valid remote
+		remoteName, remoteBranch := actions.ParseRemoteBranchName(maybeUpstream)
+		if remoteName != "" && remoteBranch != "" {
+			return remoteName, remoteBranch, nil
+		}
+	}
+
+	// It's not a remote, so now we check if it's a valid branch
+	ok, err := actions.IsBranch(ctx, dbData.Ddb, maybeUpstream)
+	if err != nil {
+		return "", "", err
+	}
+	if !ok {
+		return "", "", fmt.Errorf("branch not found: '%s'", maybeUpstream)
+	}
+
+	if maybeUpstream == selectedBranch {
+		return "", "", fmt.Errorf("	not setting '%s' as its own upstream", selectedBranch)
+	}
+
+	// In this case we use the local branch for upstream, so the remote name is empty
+	return "", maybeUpstream, nil
+}

integration-tests/bats/branch.bats

@@ -27,13 +27,12 @@ teardown() {
 }
 
 @test "branch: deleting a branch deletes its working set" {
-    dolt checkout -b to_delete
+    dolt branch to_delete
 
-    run dolt branch --datasets
+    run dolt --branch to_delete branch --datasets
     [[ "$output" =~ "workingSets/heads/main" ]] || false
     [[ "$output" =~ "workingSets/heads/to_delete" ]] || false
 
-    dolt checkout main
     dolt branch -d -f to_delete
 
     run dolt branch --datasets
@@ -42,6 +41,9 @@ teardown() {
 }
 
 @test "branch: moving current working branch takes its working set" {
+    if [ "$SQL_ENGINE" = "remote-engine" ]; then
+          skip "moves main branch which is not allowed with remote server"
+    fi
     dolt sql -q 'create table test (id int primary key);'
     dolt branch -m main new_main
     run dolt branch --show-current
@@ -64,20 +66,17 @@ teardown() {
     dolt push --set-upstream origin b3
 
     # b1 is one commit ahead of the remote
-    dolt checkout b1
-    dolt sql -q "create table t2 (id int primary key);"
-    dolt commit -Am "new table"
+    dolt --branch b1 sql -q "create table t2 (id int primary key);"
+    dolt --branch b1 commit -Am "new table"
 
     # b2 is even with the remote
 
     # b3 is one commit behind the remote
-    dolt checkout b3
-    dolt sql -q "create table t2 (id int primary key);"
-    dolt commit -Am "new table"
-    dolt push origin b3
-    dolt reset --hard HEAD~
-    
-    dolt checkout main
+    dolt --branch b3 sql -q "create table t2 (id int primary key);"
+    dolt --branch b3 commit -Am "new table"
+    dolt --branch b3 push origin b3
+    dolt --branch b3 reset --hard HEAD~
+
     run dolt branch -d b1
     [ "$status" -ne 0 ]
     [[ "$output" =~ "branch 'b1' is not fully merged" ]] || false
@@ -104,20 +103,17 @@ teardown() {
     dolt branch b3
     
     # b1 is one commit ahead of main
-    dolt checkout b1
-    dolt sql -q "create table t3 (id int primary key);"
-    dolt commit -Am "new table"
+    dolt --branch b1 sql -q "create table t3 (id int primary key);"
+    dolt --branch b1 commit -Am "new table"
     # two additional copies
-    dolt branch b1-1
-    dolt branch b1-2
+    dolt --branch b1 branch b1-1
+    dolt --branch b1 branch b1-2
 
     # b2 is even with main
 
     # b3 is one commit behind main
-    dolt checkout b3
-    dolt reset --hard HEAD~
-    
-    dolt checkout main
+    dolt --branch b3 reset --hard HEAD~
+
     run dolt branch -d b1
     [ "$status" -ne 0 ]
     [[ "$output" =~ "branch 'b1' is not fully merged" ]] || false
@@ -125,11 +121,9 @@ teardown() {
 
     dolt branch -D b1
 
-    dolt checkout b1-1
     # this works because it's even with the checked out branch (but not with main)
-    dolt branch -d b1-2
+    dolt --branch b1-1 branch -d b1-2
 
-    dolt checkout main
     dolt branch -D b1-1
     dolt branch -d b2
     dolt branch -d b3
@@ -142,6 +136,9 @@ teardown() {
 }
 
 @test "branch: attempting to delete the currently checked out branch results in an error" {
+    if [ "$SQL_ENGINE" = "remote-engine" ]; then
+          skip "deletes main branch which is not allowed within  remote server"
+    fi
     run dolt branch -D main
     [ "$status" -ne 0 ]
     [[ "$output" =~ "Cannot delete checked out branch 'main'" ]] || false
@@ -238,7 +235,7 @@ teardown() {
     [[ $output =~ "0" ]] || false
 }
 
-@test "branch: print nothing on successfull create" {
+@test "branch: print nothing on successful create" {
     run dolt branch newbranch1 HEAD
     [ $status -eq "0" ]
     [[ $output == "" ]] || false
@@ -284,24 +281,25 @@ teardown() {
     [ $status -eq "1" ]
     [[ "$output" =~ "is an invalid branch name" ]] || false
 
-    dolt checkout altBranch
-
-    run dolt branch -m HEAD
+    run dolt --branch altBranch branch -m HEAD
     [ $status -eq "1" ]
     [[ "$output" == "HEAD is an invalid branch name" ]] || false
-    run dolt branch -m $hash
+    run dolt --branch altBranch branch -m $hash
     [ $status -eq "1" ]
     [[ "$output" =~ "is an invalid branch name" ]] || false
 
-    run dolt branch -c HEAD
+    run dolt --branch altBranch branch -c HEAD
     [ $status -eq "1" ]
     [[ "$output" == "HEAD is an invalid branch name" ]] || false
-    run dolt branch -c $hash
+    run dolt --branch altBranch branch -c $hash
     [ $status -eq "1" ]
     [[ "$output" =~ "is an invalid branch name" ]] || false
 }
 
 @test "branch: renaming default branch should update init.defaultbranch config" {
+    if [ "$SQL_ENGINE" = "remote-engine" ]; then
+          skip "renames main branch which is not allowed with remote server"
+    fi
     # Set up initial default branch config
     dolt config --local --add init.defaultbranch main
     
@@ -354,3 +352,201 @@ teardown() {
     [ $status -eq 0 ]
     [[ "$output" =~ "main" ]] || false
 }
+
+@test "branch: dolt branch set upstream flag sets upstream" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt push --set-upstream origin main
+
+    run dolt branch testUpstream --set-upstream-to origin/main
+    [ $status -eq 0 ]
+    [[ "$output" =~ "branch 'testUpstream' set up to track 'origin/main'" ]] || false
+
+    run dolt sql -q "select remote, branch from dolt_branches where name = 'testUpstream'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "origin,main" ]] || false
+}
+
+@test "branch: can change upstream of existing branch with --set-upstream-to" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt push --set-upstream origin main
+    dolt branch br1 --set-upstream-to origin/main
+    dolt branch other
+    dolt --branch other push --set-upstream origin other
+
+    run dolt sql -q "select remote, branch from dolt_branches where name = 'br1'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "origin,main" ]] || false
+
+    dolt branch br1 --set-upstream-to origin/other
+
+    run dolt sql -q "select remote, branch from dolt_branches where name = 'br1'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "origin,other" ]] || false
+}
+
+@test "branch: can change upstream of existing branch with --set-upstream-to and current branch is assumed" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt push --set-upstream origin main
+    dolt branch other
+    dolt --branch other push --set-upstream origin other
+
+    dolt branch --set-upstream-to origin/other
+
+    run dolt sql -q "select remote, branch from dolt_branches where name = 'main'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "origin,other" ]] || false
+}
+
+@test "branch: cannot set upstream of branches with invalid remote" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    run dolt branch br1 --track origin/invalid
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "error: branch not found: 'origin/invalid'" ]] || false
+
+    run dolt branch main --set-upstream-to origin/invalid
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "error: branch not found: 'origin/invalid'" ]] || false
+}
+
+@test "branch: cannot use both --track and --set-upstream-to" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt push --set-upstream origin main
+
+    run dolt branch br1 --set-upstream-to origin/main --track origin/main
+    [ $status -eq 1 ]
+    [[ "$output" =~ "error: --set-upstream-to and --track are mutually exclusive options" ]] || false
+}
+
+@test "branch: --track sets upstream" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt push --set-upstream origin main
+
+    run dolt branch br1 --track origin/main
+    [ $status -eq 0 ]
+    [[ "$output" =~ "branch 'br1' set up to track 'origin/main'" ]] || false
+}
+
+@test "branch: can specify local branch with --track" {
+    run dolt branch br1 --track main
+    [ $status -eq 0 ]
+    [[ "$output" =~ "branch 'br1' set up to track 'main'" ]] || false
+}
+
+@test "branch: --track presumes current branch without argument" {
+    run dolt branch br1 --track
+    [ $status -eq 0 ]
+    [[ "$output" =~ "branch 'br1' set up to track 'main'" ]] || false
+}
+
+@test "branch: --set-upstream-to works with starting point" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt sql -q "create table t (i int)"
+    dolt commit -Am "Created a table"
+    dolt push --set-upstream origin main
+
+    # get the second to last commit hash
+    hash=`dolt sql -q "select commit_hash from dolt_log where message = 'Initialize data repository'" -r csv | sed -n '2p'`
+    dolt branch br1 --set-upstream-to origin/main "$hash"
+
+    run dolt --branch br1 ls
+    [ $status -eq 0 ]
+    [[ "$output" =~ "No tables in working set" ]] || false
+}
+
+@test "branch: --set-upstream-to and --track presume HEAD starting point" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+
+    dolt branch br1
+    dolt --branch br1 commit --allow-empty -m "A new commit"
+    dolt --branch br1 push --set-upstream origin main
+
+    dolt branch setUpstream --set-upstream-to origin/main
+    run dolt sql -q "select latest_commit_message from dolt_branches where name = 'setUpstream'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "Initialize data repository" ]] || false
+
+    dolt branch trackUpstream --track origin/main
+    run dolt sql -q "select latest_commit_message from dolt_branches where name = 'trackUpstream'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "Initialize data repository" ]] || false
+}
+
+@test "branch: --set-upstream-to and --track can set HEAD starting point" {
+    mkdir remote
+    mkdir repo1
+
+    cd repo1
+    dolt init
+    dolt remote add origin file://../remote
+    dolt branch br1
+    dolt --branch br1 commit --allow-empty -m "A new commit"
+    dolt --branch br1 push --set-upstream origin main
+
+    dolt branch setUpstream --set-upstream-to origin/main HEAD
+    run dolt sql -q "select latest_commit_message from dolt_branches where name = 'setUpstream'" -r csv
+    [ $status -eq 0 ]
+    [[ "$output" =~ "Initialize data repository" ]] || false
+    ! [[ "$output" =~ "A new commit" ]] || false
+}
+
+@test "branch: --set-upstream-to and --track cannot set branch as its own upstream" {
+    run dolt branch --set-upstream-to main
+    [ $status -eq 1 ]
+    [[ "$output" =~ "not setting 'main' as its own upstream" ]] || false
+
+    run dolt branch br1 --track br1
+    [ $status -eq 1 ]
+    [[ "$output" =~ "branch not found: 'br1'" ]] || false
+}
+
+@test "branch: --set-upstream-to and --track cannot set relative commit as upstream" {
+    dolt commit --allow-empty -m "Empty commit 1"
+
+    run dolt branch br1 --track HEAD~1
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "branch not found: 'HEAD~1'" ]] || false
+
+    run dolt branch br1 --set-upstream-to HEAD~1
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "branch not found: 'HEAD~1'" ]] || false
+}

integration-tests/bats/helper/local-remote.bash

@@ -114,7 +114,6 @@ SKIP_SERVER_TESTS=$(cat <<-EOM
 ~sql-fetch.bats~
 ~foreign-keys-invert-pk.bats~
 ~merge-base.bats~
-~branch.bats~
 ~auto_increment.bats~
 ~creds.bats~
 ~schema-conflicts.bats~

integration-tests/bats/remotes.bats

@@ -1684,139 +1684,6 @@ SQL
     [[ "$output" =~ "Everything up-to-date" ]] || false
 }
 
-@test "remotes: dolt branch track flag sets upstream" {
-    mkdir remote
-    mkdir repo1
-
-    cd repo1
-    dolt init
-    dolt remote add origin file://../remote
-    dolt sql -q "CREATE TABLE a (pk int)"
-    dolt add .
-    dolt commit -am "add table a"
-    dolt push --set-upstream origin main
-    dolt checkout -b other
-    dolt push --set-upstream origin other
-
-    cd ..
-    dolt clone file://./remote repo2
-
-    cd repo2
-    dolt branch
-    [[ ! "$output" =~ "other" ]] || false
-
-    run dolt branch --track other origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'other' set up to track 'origin/other'" ]] || false
-
-    run dolt status
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "On branch main" ]] || false
-
-    dolt checkout other
-    run dolt pull
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "Everything up-to-date" ]] || false
-
-    # NOTE: this command fails with git, requiring `--track=direct`, when both branch name and starting point name are defined, but Dolt allows both formats.
-    run dolt branch feature --track direct origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'feature' set up to track 'origin/other'" ]] || false
-
-    run dolt status
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "On branch other" ]] || false
-
-    dolt commit --allow-empty -m "new commit to other"
-    dolt push
-    dolt checkout feature
-    run dolt pull
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "Fast-forward" ]] || false
-
-    run dolt branch feature1 --track origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'feature1' set up to track 'origin/other'" ]] || false
-
-    run dolt branch --track direct feature2 origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'feature2' set up to track 'origin/other'" ]] || false
-
-    run dolt branch --track=direct feature3 origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'feature3' set up to track 'origin/other'" ]] || false
-}
-
-@test "remotes: call dolt_branch track flag sets upstream" {
-    mkdir remote
-    mkdir repo1
-
-    cd repo1
-    dolt init
-    dolt remote add origin file://../remote
-    dolt sql -q "CREATE TABLE a (pk int)"
-    dolt commit -Am "add table a"
-    dolt push --set-upstream origin main
-    dolt checkout -b other
-    dolt push --set-upstream origin other
-
-    cd ..
-    dolt clone file://./remote repo2
-
-    cd repo2
-    dolt branch
-    [[ ! "$output" =~ "other" ]] || false
-
-    dolt sql -q "CALL DOLT_BRANCH('--track','other','origin/other');"
-    run dolt status
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "On branch main" ]] || false
-
-    run dolt checkout other
-    [ "$status" -eq 0 ]
-
-    run dolt status
-    [[ "$output" =~ "Your branch is up to date with 'origin/other'." ]] || false
-
-    run dolt pull
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "Everything up-to-date" ]] || false
-
-    # NOTE: this command fails with git, requiring `--track=direct`, when both branch name and starting point name are defined, but Dolt allows both formats.
-    dolt sql -q "CALL DOLT_BRANCH('feature','--track','direct','origin/other');"
-    run dolt status
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "On branch other" ]] || false
-
-    dolt commit --allow-empty -m "new commit to other"
-    dolt push
-
-    run dolt checkout feature
-    [ "$status" -eq 0 ]
-
-    run dolt pull
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "Fast-forward" ]] || false
-
-    run dolt branch feature1 --track origin/other
-    [ "$status" -eq 0 ]
-    [[ "$output" =~ "branch 'feature1' set up to track 'origin/other'" ]] || false
-
-    dolt sql -q "CALL DOLT_BRANCH('--track','direct','feature2','origin/other');"
-    run dolt checkout feature2
-    [ "$status" -eq 0 ]
-
-    run dolt status
-    [[ "$output" =~ "Your branch is up to date with 'origin/other'." ]] || false
-
-    dolt sql -q "CALL DOLT_BRANCH('--track=direct','feature3','origin/other');"
-    run dolt checkout feature3
-    [ "$status" -eq 0 ]
-
-    run dolt status
-    [[ "$output" =~ "Your branch is up to date with 'origin/other'." ]] || false
-}
-
 @test "remotes: dolt_clone failure cleanup" {
     repoDir="$BATS_TMPDIR/dolt-repo-$$"
 

integration-tests/bats/status.bats

@@ -552,6 +552,22 @@ SQL
     [[ "${lines[2]}" = "  (use \"dolt push\" to publish your local commits)" ]] || false
 }
 
+@test "status: with local upstream" {
+    dolt branch br1 --track
+
+    run dolt --branch br1 status
+    [ "$status" -eq 0 ]
+    [[ "${lines[0]}" = "On branch br1" ]] || false
+    [[ "${lines[1]}" = "Your branch is up to date with 'main'." ]] || false
+    [[ "${lines[2]}" = "nothing to commit, working tree clean" ]] || false
+
+    dolt commit --allow-empty -m "Empty commit"
+    run dolt --branch br1 status
+    [ "$status" -eq 0 ]
+    [[ "${lines[0]}" = "On branch br1" ]] || false
+    [[ "${lines[1]}" = "Your branch is behind 'main' by 1 commit, and can be fast-forwarded." ]] || false
+}
+
 @test "status: tables with no observable changes don't show in status but can be staged" {
 dolt sql <<SQL
 create table t1 (id int primary key);