diff --git a/go/libraries/doltcore/sqle/dtables/workspace_table.go b/go/libraries/doltcore/sqle/dtables/workspace_table.go index f77e7d03cd5..6a7c66ea3b5 100644 --- a/go/libraries/doltcore/sqle/dtables/workspace_table.go +++ b/go/libraries/doltcore/sqle/dtables/workspace_table.go @@ -265,6 +265,14 @@ func (wtm *WorkspaceTableModifier) getWorkspaceTableWriter(ctx *sql.Context, tar setter := ds.SetWorkingRoot if targetStaging { setter = ds.SetStagingRoot + + // Ensure table exists in staging root before getting writer. + // This is necessary when staging rows from a new table that exists + // in working but not yet in staging (e.g., 'dolt add -p' on a new table). + err := wtm.ensureTableExistsInStaging(ctx) + if err != nil { + return nil, nil, err + } } gst, err := dsess.NewAutoIncrementTracker(ctx, "dolt", wtm.ws) @@ -282,6 +290,74 @@ func (wtm *WorkspaceTableModifier) getWorkspaceTableWriter(ctx *sql.Context, tar return writeSession, tableWriter, nil } +// ensureTableExistsInStaging ensures that the table exists in the staging root +// before attempting to get a table writer for it. This is necessary for the +// workspace table update workflow (used by 'dolt add -p') when staging rows +// from a table that exists in the working root but not yet in the staging root +// (i.e., a newly created table being staged for the first time). +// +// The function: +// 1. Checks if the table already exists in staging (fast path - no-op) +// 2. If not, checks if it exists in working +// 3. If yes, creates an empty table in staging with the same schema +// 4. Updates the session's staging root and wtm.ws to reflect the change +// +// Note: We create an empty table rather than copying data because 'dolt add -p' +// allows partial staging of individual rows. Each row staged will be inserted +// into the empty staging table via the workspace table UPDATE mechanism. +// +// If the table doesn't exist in either root, this function returns nil and +// lets GetTableWriter handle the error naturally. +func (wtm *WorkspaceTableModifier) ensureTableExistsInStaging(ctx *sql.Context) error { + stagedRoot := wtm.ws.StagedRoot() + + // Check if table already exists in staging - fast path + hasTable, err := stagedRoot.HasTable(ctx, wtm.tableName) + if err != nil { + return err + } + if hasTable { + return nil + } + + // Table doesn't exist in staging - check if it exists in working + workingRoot := wtm.ws.WorkingRoot() + workingTable, ok, err := workingRoot.GetTable(ctx, wtm.tableName) + if err != nil { + return err + } + if !ok { + // Table doesn't exist in either root - let GetTableWriter handle the error + return nil + } + + // Get the schema from the working table + sch, err := workingTable.GetSchema(ctx) + if err != nil { + return fmt.Errorf("failed to get schema for table %s: %w", wtm.tableName, err) + } + + // Create an empty table in staging with the same schema. + // We don't copy the data because 'dolt add -p' allows partial staging - + // each row will be inserted individually via the workspace table UPDATE. + newStaged, err := doltdb.CreateEmptyTable(ctx, stagedRoot, wtm.tableName, sch) + if err != nil { + return fmt.Errorf("failed to create table %s in staging: %w", wtm.tableName, err) + } + + // Update the session's staging root + ds := dsess.DSessFromSess(ctx.Session) + err = ds.SetStagingRoot(ctx, ctx.GetCurrentDatabase(), newStaged) + if err != nil { + return err + } + + // Update wtm.ws to reflect the new staging root + wtm.ws = wtm.ws.WithStagedRoot(newStaged) + + return nil +} + // isTrue returns true if the value is a boolean true, or an int8 value of != 0. Otherwise, it returns false. func isTrue(value interface{}) bool { switch v := value.(type) { diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go b/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go index 23e39e30aac..4ac86bb7f0d 100644 --- a/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go +++ b/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go @@ -1052,4 +1052,78 @@ var DoltWorkspaceScriptTests = []queries.ScriptTest{ }, }, }, + { + Name: "dolt_workspace_* stage individual rows from new table", + SetUpScript: []string{ + "create table new_tbl (pk int primary key, val int);", + "insert into new_tbl values (1, 100);", + "insert into new_tbl values (2, 200);", + "insert into new_tbl values (3, 300);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + // Verify initial state - all rows unstaged + Query: "select * from dolt_workspace_new_tbl", + Expected: []sql.Row{ + {0, false, "added", 1, 100, nil, nil}, + {1, false, "added", 2, 200, nil, nil}, + {2, false, "added", 3, 300, nil, nil}, + }, + }, + { + // Stage first row only - this should create empty table in staging first + Query: "UPDATE dolt_workspace_new_tbl SET staged = TRUE WHERE id = 0", + }, + { + // Verify first row is staged, others are not + Query: "select * from dolt_workspace_new_tbl", + Expected: []sql.Row{ + {0, true, "added", 1, 100, nil, nil}, + {1, false, "added", 2, 200, nil, nil}, + {2, false, "added", 3, 300, nil, nil}, + }, + }, + { + // Stage another row + Query: "UPDATE dolt_workspace_new_tbl SET staged = TRUE WHERE id = 2", + }, + { + // Verify two rows are now staged + // Note: staged rows come first, then unstaged, each sorted by their row values + Query: "select * from dolt_workspace_new_tbl", + Expected: []sql.Row{ + {0, true, "added", 1, 100, nil, nil}, + {1, true, "added", 3, 300, nil, nil}, + {2, false, "added", 2, 200, nil, nil}, + }, + }, + }, + }, + { + Name: "dolt_workspace_* stage all rows from new table at once", + SetUpScript: []string{ + "create table bulk_tbl (pk int primary key, name varchar(32));", + "insert into bulk_tbl values (1, 'alice'), (2, 'bob');", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "select * from dolt_workspace_bulk_tbl", + Expected: []sql.Row{ + {0, false, "added", 1, "alice", nil, nil}, + {1, false, "added", 2, "bob", nil, nil}, + }, + }, + { + // Stage all rows at once + Query: "UPDATE dolt_workspace_bulk_tbl SET staged = TRUE", + }, + { + Query: "select * from dolt_workspace_bulk_tbl", + Expected: []sql.Row{ + {0, true, "added", 1, "alice", nil, nil}, + {1, true, "added", 2, "bob", nil, nil}, + }, + }, + }, + }, } diff --git a/integration-tests/bats/add-patch-expect/new_table.expect b/integration-tests/bats/add-patch-expect/new_table.expect new file mode 100755 index 00000000000..53642b74f5d --- /dev/null +++ b/integration-tests/bats/add-patch-expect/new_table.expect @@ -0,0 +1,21 @@ +#!/usr/bin/expect + +set timeout 5 +set env(NO_COLOR) 1 + +source "$env(BATS_CWD)/helper/common_expect_functions.tcl" + +# This test stages only the first row from a new table that hasn't been committed yet. +spawn dolt add -p new_table + +# Stage the first row +expect_with_defaults_2 {| \+ | 1 | alice |} {Stage this row \[y,n,q,a,d,s,\?\]\? } { send "y\r"; } + +# Skip the second row +expect_with_defaults_2 {| \+ | 2 | bob |} {Stage this row \[y,n,q,a,d,s,\?\]\? } { send "n\r"; } + +# Skip the third row +expect_with_defaults_2 {| \+ | 3 | charlie |} {Stage this row \[y,n,q,a,d,s,\?\]\? } { send "n\r"; } + +expect eof +exit diff --git a/integration-tests/bats/add-patch.bats b/integration-tests/bats/add-patch.bats index 5def10235d8..84e2640f9bd 100644 --- a/integration-tests/bats/add-patch.bats +++ b/integration-tests/bats/add-patch.bats @@ -324,3 +324,53 @@ teardown() { [[ "$output" =~ "15" ]] || false } +# bats test_tags=no_lambda +@test "add-patch: new table partial staging" { + # Create a new table that doesn't exist in HEAD + dolt sql -q "CREATE TABLE new_table (pk int primary key, name varchar(32));" + dolt sql -q "INSERT INTO new_table VALUES (1, 'alice'), (2, 'bob'), (3, 'charlie');" + + # Verify table is untracked + run dolt status + [ $status -eq 0 ] + [[ "$output" =~ "new table:" ]] || false + [[ "$output" =~ "new_table" ]] || false + + # Use add -p to stage only the first row (using expect script) + run $BATS_TEST_DIRNAME/add-patch-expect/new_table.expect + [ $status -eq 0 ] + + # Verify the table is now in staging with partial changes + run dolt sql -q "SELECT * FROM new_table AS OF STAGED ORDER BY pk" + [ $status -eq 0 ] + [[ "$output" =~ "alice" ]] || false + # bob and charlie should NOT be in staged + ! [[ "$output" =~ "bob" ]] || false + ! [[ "$output" =~ "charlie" ]] || false + + # Verify workspace table shows correct state + run dolt sql -q "SELECT IF(staged, 'true', 'false'), to_pk, to_name FROM dolt_workspace_new_table ORDER BY to_pk" -r csv + [ $status -eq 0 ] + # First row should be staged (true), others should not be (false) + [[ "$output" =~ "true,1,alice" ]] || false + [[ "$output" =~ "false,2,bob" ]] || false + [[ "$output" =~ "false,3,charlie" ]] || false + + # Commit the staged changes + dolt commit -m "partial staging test" + + # Verify the table and only the first row were committed + run dolt show + [ $status -eq 0 ] + [[ "$output" =~ "added table" ]] || false + [[ "$output" =~ "| + | 1 | alice |" ]] || false + ! [[ "$output" =~ "bob" ]] || false + ! [[ "$output" =~ "charlie" ]] || false + + # Verify working still has the other rows + run dolt sql -q "SELECT * FROM dolt_workspace_new_table" + [ $status -eq 0 ] + [[ "$output" =~ "bob" ]] || false + [[ "$output" =~ "charlie" ]] || false +} +