Support staging new tables with dolt add -p

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) {

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},
+				},
+			},
+		},
+	},
 }

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

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
+}
+