From de95062d6f374717bb0fa6be36a66fff2ddcf035 Mon Sep 17 00:00:00 2001
From: Neil Macneale IV <46170177+macneale4@users.noreply.github.com>
Date: Tue, 14 Jan 2025 15:43:41 -0800
Subject: [PATCH] Add --prune option to dolt_pull procedure (#8745)

Expose in CLI and test too.
---
 go/cmd/dolt/cli/arg_parser_helpers.go         |  1 +
 go/cmd/dolt/commands/pull.go                  |  3 ++
 .../doltcore/sqle/dprocedures/dolt_pull.go    | 15 ++++---
 integration-tests/bats/pull.bats              | 43 ++++++++++++++++++-
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/go/cmd/dolt/cli/arg_parser_helpers.go b/go/cmd/dolt/cli/arg_parser_helpers.go
index dba70b3c1b9..9768f8e84d1 100644
--- a/go/cmd/dolt/cli/arg_parser_helpers.go
+++ b/go/cmd/dolt/cli/arg_parser_helpers.go
@@ -211,6 +211,7 @@ func CreatePullArgParser() *argparser.ArgParser {
 	ap.SupportsFlag(NoCommitFlag, "", "Perform the merge and stop just before creating a merge commit. Note this will not prevent a fast-forward merge; use the --no-ff arg together with the --no-commit arg to prevent both fast-forwards and merge commits.")
 	ap.SupportsFlag(NoEditFlag, "", "Use an auto-generated commit message when creating a merge commit. The default for interactive CLI sessions is to open an editor.")
 	ap.SupportsString(UserFlag, "", "user", "User name to use when authenticating with the remote. Gets password from the environment variable {{.EmphasisLeft}}DOLT_REMOTE_PASSWORD{{.EmphasisRight}}.")
+	ap.SupportsFlag(PruneFlag, "p", "After fetching, remove any remote-tracking references that don't exist on the remote.")
 	ap.SupportsFlag(SilentFlag, "", "Suppress progress information.")
 	return ap
 }
diff --git a/go/cmd/dolt/commands/pull.go b/go/cmd/dolt/commands/pull.go
index cc41604a806..b5a82820623 100644
--- a/go/cmd/dolt/commands/pull.go
+++ b/go/cmd/dolt/commands/pull.go
@@ -254,6 +254,9 @@ func constructInterpolatedDoltPullQuery(apr *argparser.ArgParseResults) (string,
 	if apr.Contains(cli.NoEditFlag) {
 		args = append(args, "'--no-edit'")
 	}
+	if apr.Contains(cli.PruneFlag) {
+		args = append(args, "'--prune'")
+	}
 	if user, hasUser := apr.GetValue(cli.UserFlag); hasUser {
 		args = append(args, "'--user'")
 		args = append(args, "?")
diff --git a/go/libraries/doltcore/sqle/dprocedures/dolt_pull.go b/go/libraries/doltcore/sqle/dprocedures/dolt_pull.go
index 18b436861cc..349b751c0b7 100644
--- a/go/libraries/doltcore/sqle/dprocedures/dolt_pull.go
+++ b/go/libraries/doltcore/sqle/dprocedures/dolt_pull.go
@@ -137,12 +137,7 @@ func doDoltPull(ctx *sql.Context, args []string) (int, int, string, error) {
 		return noConflictsOrViolations, threeWayMerge, "", err
 	}
 
-	// Fetch all references
-	branchRefs, err := srcDB.GetHeadRefs(ctx)
-	if err != nil {
-		return noConflictsOrViolations, threeWayMerge, "", fmt.Errorf("%w: %s", env.ErrFailedToReadDb, err.Error())
-	}
-
+	// Assert the branch exists
 	_, hasBranch, err := srcDB.HasBranch(ctx, pullSpec.Branch.GetPath())
 	if err != nil {
 		return noConflictsOrViolations, threeWayMerge, "", err
@@ -152,7 +147,13 @@ func doDoltPull(ctx *sql.Context, args []string) (int, int, string, error) {
 			fmt.Errorf("branch %q not found on remote", pullSpec.Branch.GetPath())
 	}
 
-	mode := ref.UpdateMode{Force: true, Prune: false}
+	// Fetch all references
+	branchRefs, err := srcDB.GetHeadRefs(ctx)
+	if err != nil {
+		return noConflictsOrViolations, threeWayMerge, "", fmt.Errorf("%w: %s", env.ErrFailedToReadDb, err.Error())
+	}
+	prune := apr.Contains(cli.PruneFlag)
+	mode := ref.UpdateMode{Force: true, Prune: prune}
 	err = actions.FetchRefSpecs(ctx, dbData, srcDB, pullSpec.RefSpecs, false, &pullSpec.Remote, mode, runProgFuncs, stopProgFuncs)
 	if err != nil {
 		return noConflictsOrViolations, threeWayMerge, "", fmt.Errorf("fetch failed: %w", err)
diff --git a/integration-tests/bats/pull.bats b/integration-tests/bats/pull.bats
index 2c431c2a1fa..9229928d79c 100755
--- a/integration-tests/bats/pull.bats
+++ b/integration-tests/bats/pull.bats
@@ -21,7 +21,6 @@ setup() {
     cd $TESTDIRS
     dolt clone file://rem1 repo2
     cd $TESTDIRS/repo2
-    dolt log
     dolt branch feature
     dolt remote add test-remote file://../rem1
 
@@ -541,4 +540,44 @@ SQL
     run dolt ls
     [ "$status" -eq 0 ]
     [[ "$output" =~ "testTable" ]] || false
-}
\ No newline at end of file
+}
+
+@test "pull: clean up branches with --prune" {
+    cd repo1
+    dolt checkout -b other
+    dolt commit --allow-empty -m "new commit on other"
+    dolt push origin HEAD:new-branch
+
+    run dolt branch -a
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "origin/main" ]] || false
+    [[ "$output" =~ "origin/new-branch" ]] || false
+
+    cd ../repo2
+    dolt pull origin
+
+    run dolt branch -a
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "origin/main" ]] || false
+    [[ "$output" =~ "origin/new-branch" ]] || false
+
+    dolt merge origin/new-branch
+    dolt push origin HEAD:main
+    # Delete the remote branch.
+    dolt push origin :new-branch
+
+    cd ../repo1
+    dolt branch -d main
+    dolt checkout -b main origin/main # Ensure we are tracking the remote branch.
+
+    dolt branch
+    run dolt pull origin --prune
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "Fast-forward" ]] || false
+
+    run dolt branch -a
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "origin/main" ]] || false
+    # Verify that the remote branch was deleted.
+    [[ ! "$output" =~ "origin/new-branch" ]] || false
+}