diff --git a/go/libraries/doltcore/sqle/dsess/session.go b/go/libraries/doltcore/sqle/dsess/session.go index 104b2917d2d..83191dd7c0c 100644 --- a/go/libraries/doltcore/sqle/dsess/session.go +++ b/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) diff --git a/integration-tests/bats/sql-server-remotesrv.bats b/integration-tests/bats/sql-server-remotesrv.bats index 73ed15a22eb..af84837f7fd 100644 --- a/integration-tests/bats/sql-server-remotesrv.bats +++ b/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