go/libraries/doltcore/env/actions: branch.go: Fix dolt checkout so it works even if the currently checked out branch does not exist.

go/libraries/doltcore/env/actions/branch.go

@@ -314,11 +314,14 @@ func CheckoutBranch(ctx context.Context, dEnv *env.DoltEnv, brName string) error
 	}
 
 	roots, err := dEnv.Roots(ctx)
+	if errors.Is(err, doltdb.ErrBranchNotFound) {
+		roots, err = dEnv.RecoveryRoots(ctx)
+	}
 	if err != nil {
 		return err
 	}
 
-	unstagedDocs, err := GetUnstagedDocs(ctx, dEnv)
+	unstagedDocs, err := GetUnstagedDocsFromRoots(ctx, dEnv, roots)
 	if err != nil {
 		return err
 	}

go/libraries/doltcore/env/actions/docs.go

@@ -147,7 +147,10 @@ func GetUnstagedDocs(ctx context.Context, dEnv *env.DoltEnv) (doltdocs.Docs, err
 	if err != nil {
 		return nil, err
 	}
+	return GetUnstagedDocsFromRoots(ctx, dEnv, roots)
+}
 
+func GetUnstagedDocsFromRoots(ctx context.Context, dEnv *env.DoltEnv, roots doltdb.Roots) (doltdocs.Docs, error) {
 	docsOnDisk, err := dEnv.DocsReadWriter().GetDocsOnDisk()
 	if err != nil {
 		return nil, err

go/libraries/doltcore/env/environment.go

@@ -531,6 +531,33 @@ func (dEnv *DoltEnv) Roots(ctx context.Context) (doltdb.Roots, error) {
 	}, nil
 }
 
+// RecoveryRoots returns the roots for this environment in the case that the
+// currently checked out branch has been deleted or HEAD has been updated in a
+// non-principled way to point to a branch that does not exist. This is used by
+// `dolt checkout`, in particular, to go forward with a `dolt checkout` of an
+// existing branch in the degraded state where the current branch was deleted.
+func (dEnv *DoltEnv) RecoveryRoots(ctx context.Context) (doltdb.Roots, error) {
+	ws, err := dEnv.WorkingSet(ctx)
+	if err != nil {
+		return doltdb.Roots{}, err
+	}
+
+	headRoot, err := dEnv.HeadRoot(ctx)
+	if err == doltdb.ErrBranchNotFound {
+		headRoot = ws.StagedRoot()
+		err = nil
+	}
+	if err != nil {
+		return doltdb.Roots{}, err
+	}
+
+	return doltdb.Roots{
+		Head:    headRoot,
+		Working: ws.WorkingRoot(),
+		Staged:  ws.StagedRoot(),
+	}, nil
+}
+
 // UpdateRoots updates the working and staged roots for this environment
 func (dEnv *DoltEnv) UpdateRoots(ctx context.Context, roots doltdb.Roots) error {
 	ws, err := dEnv.WorkingSet(ctx)

integration-tests/bats/deleted-branches.bats

@@ -0,0 +1,16 @@
+#!/usr/bin/env bats
+load $BATS_TEST_DIRNAME/helper/common.bash
+
+setup() {
+    setup_common
+}
+
+teardown() {
+     teardown_common
+}
+
+@test "deleted-branches: can checkout existing branch after checked out branch is deleted" {
+    dolt branch -c main to_keep
+    dolt sql -q 'delete from dolt_branches where name = "main"'
+    dolt checkout to_keep
+}