Merge pull request #9668 from dolthub/fulghum/ws-init

Initialize a branch's working set when referenced as a branch-revision database

Author: fulghum

go/libraries/doltcore/sqle/dsess/session.go

@@ -1411,25 +1411,36 @@ func (d *DoltSession) addDB(ctx *sql.Context, db SqlDatabase) error {
 
 	if dbState.Err != nil {
 		sessionState.Err = dbState.Err
-	} else if dbState.WorkingSet != nil {
-		branchState.workingSet = dbState.WorkingSet
-
-		// TODO: this is pretty clunky, there is a silly dependency between InitialDbState and globalstate.StateProvider
-		//  that's hard to express with the current types
-		stateProvider, ok := db.(globalstate.GlobalStateProvider)
-		if !ok {
-			return fmt.Errorf("database does not contain global state store")
+	} else {
+		// If the dbState doesn't have a working set yet, try to
+		// initialize one – this will only initialize a working set
+		// if the database is a branch revision database.
+		if dbState.WorkingSet == nil {
+			if err := initializeBranchWorkingSet(ctx, db, &dbState); err != nil {
+				return err
+			}
 		}
-		sessionState.globalState = stateProvider.GetGlobalState()
 
-		tracker, err := sessionState.globalState.AutoIncrementTracker(ctx)
-		if err != nil {
-			return err
+		if dbState.WorkingSet != nil {
+			branchState.workingSet = dbState.WorkingSet
+
+			// TODO: this is pretty clunky, there is a silly dependency between InitialDbState and globalstate.StateProvider
+			//  that's hard to express with the current types
+			stateProvider, ok := db.(globalstate.GlobalStateProvider)
+			if !ok {
+				return fmt.Errorf("database does not contain global state store")
+			}
+			sessionState.globalState = stateProvider.GetGlobalState()
+
+			tracker, err := sessionState.globalState.AutoIncrementTracker(ctx)
+			if err != nil {
+				return err
+			}
+			branchState.writeSession = d.writeSessProv(nbf, branchState.WorkingSet(), tracker, editOpts)
 		}
-		branchState.writeSession = d.writeSessProv(nbf, branchState.WorkingSet(), tracker, editOpts)
 	}
 
-	// WorkingSet is nil in the case of a read only, detached head DB
+	// WorkingSet is nil in the case of a read-only, detached head DB
 	if dbState.HeadCommit != nil {
 		headRoot, err := dbState.HeadCommit.GetRootValue(ctx)
 		if err != nil {
@@ -1747,6 +1758,41 @@ func (d *DoltSession) GCSafepointController() *gcctx.GCSafepointController {
 	return d.gcSafepointController
 }
 
+// initializeBranchWorkingSet checks if |db| is a branch revision database, and if |dbState|
+// does not have a working set yet, then a new, empty working set is created and set in |dbState|.
+// If |db| is NOT a branch revision database, or |dbState| already has a working set, then this
+// function will not do anything.
+func initializeBranchWorkingSet(ctx *sql.Context, db SqlDatabase, dbState *InitialDbState) error {
+	revisionDb, isRevisionDb := db.(RevisionDatabase)
+	if !isRevisionDb || revisionDb.RevisionType() != RevisionTypeBranch || dbState.WorkingSet != nil {
+		return nil
+	}
+
+	branchRef := ref.NewBranchRef(revisionDb.Revision())
+	wsRef, err := ref.WorkingSetRefForHead(branchRef)
+	if err != nil {
+		return err
+	}
+
+	commit, err := dbState.DbData.Ddb.ResolveCommitRef(ctx, branchRef)
+	if err != nil {
+		return err
+	}
+
+	headRoot, err := commit.GetRootValue(ctx)
+	if err != nil {
+		return err
+	}
+
+	dbState.WorkingSet = doltdb.EmptyWorkingSet(wsRef).
+		WithWorkingRoot(headRoot).WithStagedRoot(headRoot)
+
+	ctx.GetLogger().Warnf("initializing empty working set for branch %s", revisionDb.Revision())
+
+	return dbState.DbData.Ddb.UpdateWorkingSet(ctx, wsRef, dbState.WorkingSet,
+		hash.Hash{}, doltdb.TodoWorkingSetMeta(), nil)
+}
+
 // validatePersistedSysVar checks whether a system variable exists and is dynamic
 func validatePersistableSysVar(name string) (sql.SystemVariable, interface{}, error) {
 	sysVar, val, ok := sql.SystemVariables.GetGlobal(name)

integration-tests/bats/sql-server-remotesrv.bats

@@ -430,6 +430,34 @@ call dolt_commit('-am', 'add one val');"
     [[ "$output" =~ "dave" ]] || false
 }
 
+# Assert that when a new branch that has been pushed to a running SQL server through the RemotesAPI,
+# it will have its working set properly initialized so that it can be written to during a write
+# operation that references the branch as a branch-revision database.
+@test "sql-server-remotesrv: can write to a branch rev db for a new branch pushed to a sql-server remote" {
+    mkdir -p db/remote
+    cd db/remote
+    dolt init
+    dolt sql-server --remotesapi-port 50051 --loglevel DEBUG &
+    srv_pid=$!
+    cd ../..
+
+    # By cloning here, we have a near-at-hand way to wait for the server to be ready.
+    dolt clone http://localhost:50051/remote cloned_remote
+    cd cloned_remote
+
+    # create a new branch in the clone and push it to the running sql-server
+    dolt checkout -b new_branch
+    dolt push origin new_branch
+
+    # Check out the new branch, do a test write, and commit
+    cd ../db/remote
+    dolt sql -q "CREATE TABLE \`remote/new_branch\`.t123 (pk int primary key);"
+    run dolt sql -q "SELECT * FROM \`remote/new_branch\`.dolt_status;"
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "t123       | 0      | new table" ]] || false
+    dolt sql -q "CALL dolt_checkout('new_branch'); CALL dolt_commit('-Am', 'add table t123');"
+}
+
 @test "sql-server-remotesrv: push to dirty workspace as super user" {
     mkdir remote
     cd remote