add --track flag for dolt_branch() procedure (#5312)

go/cmd/dolt/cli/arg_parser_helpers.go

@@ -129,7 +129,9 @@ var mergeAbortDetails = `Abort the current conflict resolution process, and try
 If there were uncommitted working set changes present when the merge started, {{.EmphasisLeft}}dolt merge --abort{{.EmphasisRight}} will be unable to reconstruct these changes. It is therefore recommended to always commit or stash your changes before running dolt merge.
 `
 
-// Creates the argparser shared dolt commit cli and DOLT_COMMIT.
+var branchForceFlagDesc = "Reset {{.LessThan}}branchname{{.GreaterThan}} to {{.LessThan}}startpoint{{.GreaterThan}}, even if {{.LessThan}}branchname{{.GreaterThan}} exists already. Without {{.EmphasisLeft}}-f{{.EmphasisRight}}, {{.EmphasisLeft}}dolt branch{{.EmphasisRight}} refuses to change an existing branch. In combination with {{.EmphasisLeft}}-d{{.EmphasisRight}} (or {{.EmphasisLeft}}--delete{{.EmphasisRight}}), allow deleting the branch irrespective of its merged status. In combination with -m (or {{.EmphasisLeft}}--move{{.EmphasisRight}}), allow renaming the branch even if the new branch name already exists, the same applies for {{.EmphasisLeft}}-c{{.EmphasisRight}} (or {{.EmphasisLeft}}--copy{{.EmphasisRight}})."
+
+// CreateCommitArgParser creates the argparser shared dolt commit cli and DOLT_COMMIT.
 func CreateCommitArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParser()
 	ap.SupportsString(MessageArg, "m", "msg", "Use the given {{.LessThan}}msg{{.GreaterThan}} as the commit message.")
@@ -174,7 +176,7 @@ func CreatePushArgParser() *argparser.ArgParser {
 func CreateAddArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParser()
 	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"table", "Working table(s) to add to the list tables staged to be committed. The abbreviation '.' can be used to add all tables."})
-	ap.SupportsFlag("all", "A", "Stages any and all changes (adds, deletes, and modifications).")
+	ap.SupportsFlag(AllFlag, "A", "Stages any and all changes (adds, deletes, and modifications).")
 	return ap
 }
 
@@ -255,11 +257,12 @@ func CreatePullArgParser() *argparser.ArgParser {
 
 func CreateBranchArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParser()
-	ap.SupportsFlag(ForceFlag, "f", "Ignores any foreign key warnings and proceeds with the commit.")
+	ap.SupportsFlag(ForceFlag, "f", branchForceFlagDesc)
 	ap.SupportsFlag(CopyFlag, "c", "Create a copy of a branch.")
 	ap.SupportsFlag(MoveFlag, "m", "Move/rename a branch")
 	ap.SupportsFlag(DeleteFlag, "d", "Delete a branch. The branch must be fully merged in its upstream branch.")
 	ap.SupportsFlag(DeleteForceFlag, "", "Shortcut for {{.EmphasisLeft}}--delete --force{{.EmphasisRight}}.")
+	ap.SupportsString(TrackFlag, "t", "", "When creating a new branch, set up 'upstream' configuration.")
 
 	return ap
 }

go/cmd/dolt/commands/backup.go

@@ -202,7 +202,7 @@ func printBackups(dEnv *env.DoltEnv, apr *argparser.ArgParseResults) errhand.Ver
 	}
 
 	for _, r := range backups {
-		if apr.Contains(verboseFlag) {
+		if apr.Contains(cli.VerboseFlag) {
 			paramStr := make([]byte, 0)
 			if len(r.Params) > 0 {
 				paramStr, _ = json.Marshal(r.Params)

go/cmd/dolt/commands/branch.go

@@ -34,8 +34,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/utils/set"
 )
 
-var branchForceFlagDesc = "Reset {{.LessThan}}branchname{{.GreaterThan}} to {{.LessThan}}startpoint{{.GreaterThan}}, even if {{.LessThan}}branchname{{.GreaterThan}} exists already. Without {{.EmphasisLeft}}-f{{.EmphasisRight}}, {{.EmphasisLeft}}dolt branch{{.EmphasisRight}} refuses to change an existing branch. In combination with {{.EmphasisLeft}}-d{{.EmphasisRight}} (or {{.EmphasisLeft}}--delete{{.EmphasisRight}}), allow deleting the branch irrespective of its merged status. In combination with -m (or {{.EmphasisLeft}}--move{{.EmphasisRight}}), allow renaming the branch even if the new branch name already exists, the same applies for {{.EmphasisLeft}}-c{{.EmphasisRight}} (or {{.EmphasisLeft}}--copy{{.EmphasisRight}})."
-
 var branchDocs = cli.CommandDocumentationContent{
 	ShortDesc: `List, create, or delete branches`,
 	LongDesc: `If {{.EmphasisLeft}}--list{{.EmphasisRight}} is given, or if there are no non-option arguments, existing branches are listed. The current branch will be highlighted with an asterisk. With no options, only local branches are listed. With {{.EmphasisLeft}}-r{{.EmphasisRight}}, only remote branches are listed. With {{.EmphasisLeft}}-a{{.EmphasisRight}} both local and remote branches are listed. {{.EmphasisLeft}}-v{{.EmphasisRight}} causes the hash of the commit that the branches are at to be printed as well.
@@ -60,15 +58,7 @@ With a {{.EmphasisLeft}}-d{{.EmphasisRight}}, {{.LessThan}}branchname{{.GreaterT
 
 const (
 	listFlag        = "list"
-	forceFlag       = "force"
-	copyFlag        = "copy"
-	moveFlag        = "move"
-	deleteFlag      = "delete"
-	deleteForceFlag = "D"
-	verboseFlag     = cli.VerboseFlag
-	allFlag         = "all"
 	datasetsFlag    = "datasets"
-	remoteFlag      = "remote"
 	showCurrentFlag = "show-current"
 )
 
@@ -92,20 +82,14 @@ func (cmd BranchCmd) Docs() *cli.CommandDocumentation {
 }
 
 func (cmd BranchCmd) ArgParser() *argparser.ArgParser {
-	ap := argparser.NewArgParser()
+	ap := cli.CreateBranchArgParser()
 	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"start-point", "A commit that a new branch should point at."})
 	ap.SupportsFlag(listFlag, "", "List branches")
-	ap.SupportsFlag(forceFlag, "f", branchForceFlagDesc)
-	ap.SupportsFlag(copyFlag, "c", "Create a copy of a branch.")
-	ap.SupportsFlag(moveFlag, "m", "Move/rename a branch")
-	ap.SupportsFlag(deleteFlag, "d", "Delete a branch. The branch must be fully merged in its upstream branch.")
-	ap.SupportsFlag(deleteForceFlag, "", "Shortcut for {{.EmphasisLeft}}--delete --force{{.EmphasisRight}}.")
-	ap.SupportsFlag(verboseFlag, "v", "When in list mode, show the hash and commit subject line for each head")
-	ap.SupportsFlag(allFlag, "a", "When in list mode, shows remote tracked branches")
+	ap.SupportsFlag(cli.VerboseFlag, "v", "When in list mode, show the hash and commit subject line for each head")
+	ap.SupportsFlag(cli.AllFlag, "a", "When in list mode, shows remote tracked branches")
 	ap.SupportsFlag(datasetsFlag, "", "List all datasets in the database")
-	ap.SupportsFlag(remoteFlag, "r", "When in list mode, show only remote tracked branches. When with -d, delete a remote tracking branch.")
+	ap.SupportsFlag(cli.RemoteParam, "r", "When in list mode, show only remote tracked branches. When with -d, delete a remote tracking branch.")
 	ap.SupportsFlag(showCurrentFlag, "", "Print the name of the current branch")
-	ap.SupportsString(cli.TrackFlag, "t", "", "When creating a new branch, set up 'upstream' configuration.")
 	return ap
 }
 
@@ -121,13 +105,13 @@ func (cmd BranchCmd) Exec(ctx context.Context, commandStr string, args []string,
 	apr := cli.ParseArgsOrDie(ap, args, help)
 
 	switch {
-	case apr.Contains(moveFlag):
+	case apr.Contains(cli.MoveFlag):
 		return moveBranch(ctx, dEnv, apr, usage)
-	case apr.Contains(copyFlag):
+	case apr.Contains(cli.CopyFlag):
 		return copyBranch(ctx, dEnv, apr, usage)
-	case apr.Contains(deleteFlag):
-		return deleteBranches(ctx, dEnv, apr, usage, apr.Contains(forceFlag))
-	case apr.Contains(deleteForceFlag):
+	case apr.Contains(cli.DeleteFlag):
+		return deleteBranches(ctx, dEnv, apr, usage, apr.Contains(cli.ForceFlag))
+	case apr.Contains(cli.DeleteForceFlag):
 		return deleteBranches(ctx, dEnv, apr, usage, true)
 	case apr.Contains(listFlag):
 		return printBranches(ctx, dEnv, apr, usage)
@@ -145,9 +129,9 @@ func (cmd BranchCmd) Exec(ctx context.Context, commandStr string, args []string,
 func printBranches(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgParseResults, _ cli.UsagePrinter) int {
 	branchSet := set.NewStrSet(apr.Args)
 
-	verbose := apr.Contains(verboseFlag)
+	verbose := apr.Contains(cli.VerboseFlag)
 	printRemote := apr.Contains(cli.RemoteParam)
-	printAll := apr.Contains(allFlag)
+	printAll := apr.Contains(cli.AllFlag)
 
 	branches, err := dEnv.DoltDB.GetHeadRefs(ctx)
 
@@ -260,7 +244,7 @@ func moveBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgParseR
 		return 1
 	}
 
-	force := apr.Contains(forceFlag)
+	force := apr.Contains(cli.ForceFlag)
 	src := apr.Arg(0)
 	dest := apr.Arg(1)
 	err := actions.RenameBranch(ctx, dEnv.DbData(), src, apr.Arg(1), dEnv, force)
@@ -290,7 +274,7 @@ func copyBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgParseR
 		return 1
 	}
 
-	force := apr.Contains(forceFlag)
+	force := apr.Contains(cli.ForceFlag)
 	src := apr.Arg(0)
 	dest := apr.Arg(1)
 	err := actions.CopyBranch(ctx, dEnv, src, dest, force)
@@ -323,7 +307,7 @@ func deleteBranches(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgPa
 
 		err := actions.DeleteBranch(ctx, dEnv.DbData(), brName, actions.DeleteOptions{
 			Force:  force,
-			Remote: apr.Contains(remoteFlag),
+			Remote: apr.Contains(cli.RemoteParam),
 		}, dEnv)
 
 		if err != nil {
@@ -373,9 +357,8 @@ func createBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgPars
 		}
 
 		if apr.NArg() == 2 {
-			newBranch = apr.Arg(0)
-			startPt = apr.Arg(1)
-			remote, remoteBranch = ParseRemoteBranchName(startPt)
+			// branchName and startPt are already set
+			remote, remoteBranch = actions.ParseRemoteBranchName(startPt)
 			_, remoteOk := remotes[remote]
 			if !remoteOk {
 				return HandleVErrAndExitCode(errhand.BuildDError("'%s' is not a valid remote ref and a branch '%s' cannot be created from it", startPt, newBranch).Build(), usage)
@@ -384,12 +367,12 @@ func createBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgPars
 			// if track option is defined with no value,
 			// the track value can either be starting point name OR branch name
 			startPt = trackVal
-			remote, remoteBranch = ParseRemoteBranchName(startPt)
+			remote, remoteBranch = actions.ParseRemoteBranchName(startPt)
 			_, remoteOk := remotes[remote]
 			if !remoteOk {
 				newBranch = trackVal
 				startPt = apr.Arg(0)
-				remote, remoteBranch = ParseRemoteBranchName(startPt)
+				remote, remoteBranch = actions.ParseRemoteBranchName(startPt)
 				_, remoteOk = remotes[remote]
 				if !remoteOk {
 					return HandleVErrAndExitCode(errhand.BuildDError("'%s' is not a valid remote ref and a branch '%s' cannot be created from it", startPt, newBranch).Build(), usage)
@@ -398,7 +381,7 @@ func createBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.ArgPars
 		}
 	}
 
-	err := actions.CreateBranchWithStartPt(ctx, dEnv.DbData(), newBranch, startPt, apr.Contains(forceFlag))
+	err := actions.CreateBranchWithStartPt(ctx, dEnv.DbData(), newBranch, startPt, apr.Contains(cli.ForceFlag))
 	if err != nil {
 		return HandleVErrAndExitCode(errhand.BuildDError(err.Error()).Build(), usage)
 	}

go/cmd/dolt/commands/checkout.go

@@ -17,7 +17,6 @@ package commands
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 
@@ -154,7 +153,7 @@ func checkoutNewBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.Ar
 		} else if trackVal == "inherit" {
 			return errhand.VerboseErrorFromError(fmt.Errorf("--track='inherit' is not supported yet"))
 		}
-		remoteName, remoteBranchName = ParseRemoteBranchName(startPt)
+		remoteName, remoteBranchName = actions.ParseRemoteBranchName(startPt)
 		remotes, err := dEnv.RepoStateReader().GetRemotes()
 		if err != nil {
 			return errhand.BuildDError(err.Error()).Build()
@@ -191,7 +190,7 @@ func checkoutNewBranch(ctx context.Context, dEnv *env.DoltEnv, apr *argparser.Ar
 		if err != nil {
 			return nil
 		}
-		remoteName, remoteBranchName = ParseRemoteBranchName(startPt)
+		remoteName, remoteBranchName = actions.ParseRemoteBranchName(startPt)
 		_, remoteOk := remotes[remoteName]
 		if !remoteOk {
 			return nil
@@ -326,15 +325,7 @@ func SetRemoteUpstreamForBranchRef(dEnv *env.DoltEnv, remote, remoteBranch strin
 		return errhand.BuildDError(fmt.Errorf("%w: '%s'", err, remote).Error()).Build()
 	}
 
-	src := refSpec.SrcRef(branchRef)
-	dest := refSpec.DestRef(src)
-
-	err = dEnv.RepoStateWriter().UpdateBranch(branchRef.GetPath(), env.BranchConfig{
-		Merge: ref.MarshalableRef{
-			Ref: dest,
-		},
-		Remote: remote,
-	})
+	err = env.SetRemoteUpstreamForRefSpec(dEnv.RepoStateWriter(), refSpec, remote, branchRef)
 	if err != nil {
 		return errhand.BuildDError(err.Error()).Build()
 	}
@@ -348,12 +339,3 @@ func unreadableRootToVErr(err error) errhand.VerboseError {
 	bdr := errhand.BuildDError("error: unable to read the %s", rt.String())
 	return bdr.AddCause(doltdb.GetUnreachableRootCause(err)).Build()
 }
-
-func ParseRemoteBranchName(startPt string) (string, string) {
-	startPt = strings.TrimPrefix(startPt, "remotes/")
-	names := strings.Split(startPt, "/")
-	if len(names) < 2 {
-		return "", ""
-	}
-	return names[0], strings.Join(names[1:], "/")
-}

go/cmd/dolt/commands/filter-branch.go

@@ -84,9 +84,9 @@ func (cmd FilterBranchCmd) Docs() *cli.CommandDocumentation {
 
 func (cmd FilterBranchCmd) ArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParser()
-	ap.SupportsFlag(verboseFlag, "v", "logs more information")
+	ap.SupportsFlag(cli.VerboseFlag, "v", "logs more information")
 	ap.SupportsFlag(branchesFlag, "b", "filter all branches")
-	ap.SupportsFlag(allFlag, "a", "filter all branches and tags")
+	ap.SupportsFlag(cli.AllFlag, "a", "filter all branches and tags")
 	return ap
 }
 
@@ -112,7 +112,7 @@ func (cmd FilterBranchCmd) Exec(ctx context.Context, commandStr string, args []s
 	}
 
 	query := apr.Arg(0)
-	verbose := apr.Contains(verboseFlag)
+	verbose := apr.Contains(cli.VerboseFlag)
 	notFound := make(missingTbls)
 
 	replay := func(ctx context.Context, commit, _, _ *doltdb.Commit) (*doltdb.RootValue, error) {
@@ -160,7 +160,7 @@ func (cmd FilterBranchCmd) Exec(ctx context.Context, commandStr string, args []s
 	switch {
 	case apr.Contains(branchesFlag):
 		err = rebase.AllBranches(ctx, dEnv, replay, nerf)
-	case apr.Contains(allFlag):
+	case apr.Contains(cli.AllFlag):
 		err = rebase.AllBranchesAndTags(ctx, dEnv, replay, nerf)
 	default:
 		err = rebase.CurrentBranch(ctx, dEnv, replay, nerf)

go/cmd/dolt/commands/ls.go

@@ -66,9 +66,9 @@ func (cmd LsCmd) Docs() *cli.CommandDocumentation {
 
 func (cmd LsCmd) ArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParser()
-	ap.SupportsFlag(verboseFlag, "v", "show the hash of the table")
+	ap.SupportsFlag(cli.VerboseFlag, "v", "show the hash of the table")
 	ap.SupportsFlag(systemFlag, "s", "show system tables")
-	ap.SupportsFlag(allFlag, "a", "show system tables")
+	ap.SupportsFlag(cli.AllFlag, "a", "show system tables")
 	return ap
 }
 
@@ -83,8 +83,8 @@ func (cmd LsCmd) Exec(ctx context.Context, commandStr string, args []string, dEn
 	help, usage := cli.HelpAndUsagePrinters(cli.CommandDocsForCommandString(commandStr, lsDocs, ap))
 	apr := cli.ParseArgsOrDie(ap, args, help)
 
-	if apr.Contains(systemFlag) && apr.Contains(allFlag) {
-		verr := errhand.BuildDError("--%s and --%s are mutually exclusive", systemFlag, allFlag).SetPrintUsage().Build()
+	if apr.Contains(systemFlag) && apr.Contains(cli.AllFlag) {
+		verr := errhand.BuildDError("--%s and --%s are mutually exclusive", systemFlag, cli.AllFlag).SetPrintUsage().Build()
 		HandleVErrAndExitCode(verr, usage)
 	}
 
@@ -99,13 +99,13 @@ func (cmd LsCmd) Exec(ctx context.Context, commandStr string, args []string, dEn
 	}
 
 	if verr == nil {
-		if !apr.Contains(systemFlag) || apr.Contains(allFlag) {
-			verr = printUserTables(ctx, root, label, apr.Contains(verboseFlag))
+		if !apr.Contains(systemFlag) || apr.Contains(cli.AllFlag) {
+			verr = printUserTables(ctx, root, label, apr.Contains(cli.VerboseFlag))
 			cli.Println()
 		}
 
-		if verr == nil && (apr.Contains(systemFlag) || apr.Contains(allFlag)) {
-			verr = printSystemTables(ctx, root, dEnv.DoltDB, apr.Contains(verboseFlag))
+		if verr == nil && (apr.Contains(systemFlag) || apr.Contains(cli.AllFlag)) {
+			verr = printSystemTables(ctx, root, dEnv.DoltDB, apr.Contains(cli.VerboseFlag))
 			cli.Println()
 		}
 	}

go/cmd/dolt/commands/remote.go

@@ -85,15 +85,11 @@ func (cmd RemoteCmd) Docs() *cli.CommandDocumentation {
 }
 
 func (cmd RemoteCmd) ArgParser() *argparser.ArgParser {
-	ap := argparser.NewArgParser()
+	ap := cli.CreateRemoteArgParser()
 	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"region", "cloud provider region associated with this remote."})
 	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"creds-type", "credential type.  Valid options are role, env, and file.  See the help section for additional details."})
 	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"profile", "AWS profile to use."})
-	ap.SupportsFlag(verboseFlag, "v", "When printing the list of remotes adds additional details.")
-	ap.SupportsString(dbfactory.AWSRegionParam, "", "region", "")
-	ap.SupportsValidatedString(dbfactory.AWSCredsTypeParam, "", "creds-type", "", argparser.ValidatorFromStrList(dbfactory.AWSCredsTypeParam, dbfactory.AWSCredTypes))
-	ap.SupportsString(dbfactory.AWSCredsFileParam, "", "file", "AWS credentials file")
-	ap.SupportsString(dbfactory.AWSCredsProfile, "", "profile", "AWS profile to use")
+	ap.SupportsFlag(cli.VerboseFlag, "v", "When printing the list of remotes adds additional details.")
 	ap.SupportsString(dbfactory.OSSCredsFileParam, "", "file", "OSS credentials file")
 	ap.SupportsString(dbfactory.OSSCredsProfile, "", "profile", "OSS profile to use")
 	return ap
@@ -216,7 +212,7 @@ func printRemotes(dEnv *env.DoltEnv, apr *argparser.ArgParseResults) errhand.Ver
 	}
 
 	for _, r := range remotes {
-		if apr.Contains(verboseFlag) {
+		if apr.Contains(cli.VerboseFlag) {
 			paramStr := make([]byte, 0)
 			if len(r.Params) > 0 {
 				paramStr, _ = json.Marshal(r.Params)

go/libraries/doltcore/env/actions/remotes.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"sync"
 
 	"github.com/dolthub/dolt/go/cmd/dolt/cli"
@@ -494,3 +495,32 @@ func HandleInitRemoteStorageClientErr(name, url string, err error) error {
 	var detail = fmt.Sprintf("the remote: %s '%s' could not be accessed", name, url)
 	return fmt.Errorf("%w; %s; %s", ErrFailedToGetRemoteDb, detail, err.Error())
 }
+
+// ParseRemoteBranchName takes remote branch ref name, parses it and returns remote branch name.
+// For example, it parses the input string 'origin/john/mybranch' and returns remote name 'origin' and branch name 'john/mybranch'.
+func ParseRemoteBranchName(startPt string) (string, string) {
+	startPt = strings.TrimPrefix(startPt, "remotes/")
+	names := strings.SplitN(startPt, "/", 2)
+	if len(names) < 2 {
+		return "", ""
+	}
+	return names[0], names[1]
+}
+
+// GetRemoteBranchRef returns a remote ref with matching name for a branch for each remote.
+func GetRemoteBranchRef(ctx context.Context, ddb *doltdb.DoltDB, name string) ([]ref.RemoteRef, error) {
+	remoteRefFilter := map[ref.RefType]struct{}{ref.RemoteRefType: {}}
+	refs, err := ddb.GetRefsOfType(ctx, remoteRefFilter)
+	if err != nil {
+		return nil, err
+	}
+
+	var remoteRef []ref.RemoteRef
+	for _, rf := range refs {
+		if remRef, ok := rf.(ref.RemoteRef); ok && remRef.GetBranch() == name {
+			remoteRef = append(remoteRef, remRef)
+		}
+	}
+
+	return remoteRef, nil
+}

go/libraries/doltcore/env/actions/table.go

@@ -17,8 +17,6 @@ package actions
 import (
 	"context"
 
-	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
-
 	"github.com/dolthub/dolt/go/libraries/doltcore/diff"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/utils/set"
@@ -120,21 +118,3 @@ func RemoveDocsTable(tbls []string) []string {
 	}
 	return result
 }
-
-// GetRemoteBranchRef returns a remote ref with matching name for a branch for each remotes.
-func GetRemoteBranchRef(ctx context.Context, ddb *doltdb.DoltDB, name string) ([]ref.RemoteRef, error) {
-	remoteRefFilter := map[ref.RefType]struct{}{ref.RemoteRefType: {}}
-	refs, err := ddb.GetRefsOfType(ctx, remoteRefFilter)
-	if err != nil {
-		return nil, err
-	}
-
-	var remoteRef []ref.RemoteRef
-	for _, rf := range refs {
-		if remRef, ok := rf.(ref.RemoteRef); ok && remRef.GetBranch() == name {
-			remoteRef = append(remoteRef, remRef)
-		}
-	}
-
-	return remoteRef, nil
-}

go/libraries/doltcore/env/remotes.go

@@ -554,3 +554,17 @@ func GetDefaultBranch(dEnv *DoltEnv, branches []ref.DoltRef) string {
 
 	return branches[0].GetPath()
 }
+
+// SetRemoteUpstreamForRefSpec set upstream for given RefSpec, remote name and branch ref. It uses given RepoStateWriter
+// to persist upstream tracking branch information.
+func SetRemoteUpstreamForRefSpec(rsw RepoStateWriter, refSpec ref.RefSpec, remote string, branchRef ref.DoltRef) error {
+	src := refSpec.SrcRef(branchRef)
+	dest := refSpec.DestRef(src)
+
+	return rsw.UpdateBranch(branchRef.GetPath(), BranchConfig{
+		Merge: ref.MarshalableRef{
+			Ref: dest,
+		},
+		Remote: remote,
+	})
+}

go/libraries/doltcore/ref/remote_ref.go

@@ -57,7 +57,7 @@ func NewRemoteRef(remote, branch string) RemoteRef {
 	return RemoteRef{remote, branch}
 }
 
-// NewRemoteRefFromPathString creates a DoltRef from a string in the format origin/main, or remotes/origin/main, or
+// NewRemoteRefFromPathStr creates a DoltRef from a string in the format origin/main, or remotes/origin/main, or
 // refs/remotes/origin/main
 func NewRemoteRefFromPathStr(remoteAndPath string) (DoltRef, error) {
 	if IsRef(remoteAndPath) {

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

@@ -257,26 +257,77 @@ func loadConfig(ctx *sql.Context) *env.DoltCliConfig {
 }
 
 func createNewBranch(ctx *sql.Context, dbData env.DbData, apr *argparser.ArgParseResults) error {
-	var branchName string
+	if apr.NArg() == 0 || apr.NArg() > 2 {
+		return InvalidArgErr
+	}
+
+	var branchName = apr.Arg(0)
 	var startPt = "HEAD"
-	if apr.NArg() == 1 {
-		branchName = apr.Arg(0)
-	} else if apr.NArg() == 2 {
-		branchName = apr.Arg(0)
+	if len(branchName) == 0 {
+		return EmptyBranchNameErr
+	}
+	if apr.NArg() == 2 {
 		startPt = apr.Arg(1)
 		if len(startPt) == 0 {
 			return InvalidArgErr
 		}
 	}
 
-	if len(branchName) == 0 {
-		return EmptyBranchNameErr
+	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)
+			if err != nil {
+				return err
+			}
+		} 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
+				}
+			}
+		}
 	}
 
-	if err := branch_control.CanCreateBranch(ctx, branchName); err != nil {
+	err = branch_control.CanCreateBranch(ctx, branchName)
+	if err != nil {
 		return err
 	}
-	return actions.CreateBranchWithStartPt(ctx, dbData, branchName, startPt, apr.Contains(cli.ForceFlag))
+
+	err = actions.CreateBranchWithStartPt(ctx, dbData, branchName, startPt, apr.Contains(cli.ForceFlag))
+	if err != nil {
+		return err
+	}
+
+	if setTrackUpstream {
+		// at this point new branch is created
+		err = env.SetRemoteUpstreamForRefSpec(dbData.Rsw, refSpec, remoteName, ref.NewBranchRef(branchName))
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 func copyBranch(ctx *sql.Context, dbData env.DbData, apr *argparser.ArgParseResults) error {

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

@@ -208,15 +208,7 @@ func checkoutRemoteBranch(ctx *sql.Context, dbName string, dbData env.DbData, br
 			return errhand.BuildDError(fmt.Errorf("%w: '%s'", err, remoteRef.GetRemote()).Error()).Build()
 		}
 
-		src := refSpec.SrcRef(dbData.Rsr.CWBHeadRef())
-		dest := refSpec.DestRef(src)
-
-		return dbData.Rsw.UpdateBranch(src.GetPath(), env.BranchConfig{
-			Merge: ref.MarshalableRef{
-				Ref: dest,
-			},
-			Remote: remoteRef.GetRemote(),
-		})
+		return env.SetRemoteUpstreamForRefSpec(dbData.Rsw, refSpec, remoteRef.GetRemote(), dbData.Rsr.CWBHeadRef())
 	} else {
 		return fmt.Errorf("'%s' matched multiple (%v) remote tracking branches", branchName, len(remoteRefs))
 	}

integration-tests/bats/remotes.bats

@@ -2124,6 +2124,76 @@ SQL
     [[ "$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-$$"