Don't panic when attempting to update workspace table

go/libraries/doltcore/sqle/dtables/workspace_table.go

@@ -62,12 +62,17 @@ type WorkspaceTable struct {
 
 	// headSchema is the schema of the table that is being modified.
 	headSchema schema.Schema
+
+	// modifiable is true if the schemas before and after the update are identical. Used to reject updates that would
+	// be impossible to perform - such as only stage one row when the entire schema of the table is being modified.
+	modifiable bool
 }
 
 type WorkspaceTableModifier struct {
-	tableName doltdb.TableName
-	ws        *doltdb.WorkingSet
-	head      doltdb.RootValue
+	tableName          doltdb.TableName
+	workspaceTableName string
+	ws                 *doltdb.WorkingSet
+	head               doltdb.RootValue
 
 	headSch   schema.Schema
 	schemaLen int
@@ -76,6 +81,9 @@ type WorkspaceTableModifier struct {
 	tableWriter   *dsess.TableWriter
 	sessionWriter *dsess.WriteSession
 
+	// modifiable carried through from the main table.
+	modifiable bool
+
 	err *error
 }
 
@@ -151,6 +159,10 @@ func (wtu *WorkspaceTableUpdater) Update(ctx *sql.Context, old sql.Row, new sql.
 		return fmt.Errorf("Runtime error: expected non-nil inputs to WorkspaceTableUpdater.Update")
 	}
 
+	if !wtu.modifiable {
+		return errors.New(fmt.Sprintf("%s table is not modifiable due to schema change", wtu.workspaceTableName))
+	}
+
 	valid, isStaged := validateWorkspaceUpdate(old, new)
 	if !valid {
 		return errors.New("only update of column 'staged' is allowed")
@@ -205,6 +217,10 @@ func (wtd *WorkspaceTableDeleter) StatementBegin(ctx *sql.Context) {
 }
 
 func (wtd *WorkspaceTableDeleter) Delete(c *sql.Context, row sql.Row) error {
+	if !wtd.modifiable {
+		return errors.New(fmt.Sprintf("%s table is not modifiable due to schema change", wtd.workspaceTableName))
+	}
+
 	isStaged := isTrue(row[stagedColumnIdx])
 	if isStaged {
 		return fmt.Errorf("cannot delete staged rows from workspace")
@@ -313,11 +329,13 @@ func validateWorkspaceUpdate(old, new sql.Row) (valid, staged bool) {
 func (wt *WorkspaceTable) Deleter(_ *sql.Context) sql.RowDeleter {
 	cols := wt.headSchema.GetAllCols().Size()
 	modifier := WorkspaceTableModifier{
-		tableName: wt.userTblName,
-		headSch:   wt.headSchema,
-		schemaLen: cols,
-		ws:        wt.ws,
-		head:      wt.head,
+		tableName:          wt.userTblName,
+		workspaceTableName: wt.Name(),
+		headSch:            wt.headSchema,
+		schemaLen:          cols,
+		ws:                 wt.ws,
+		head:               wt.head,
+		modifiable:         wt.modifiable,
 	}
 
 	return &WorkspaceTableDeleter{
@@ -328,11 +346,13 @@ func (wt *WorkspaceTable) Deleter(_ *sql.Context) sql.RowDeleter {
 func (wt *WorkspaceTable) Updater(_ *sql.Context) sql.RowUpdater {
 	cols := wt.headSchema.GetAllCols().Size()
 	modifier := WorkspaceTableModifier{
-		tableName: wt.userTblName,
-		headSch:   wt.headSchema,
-		schemaLen: cols,
-		ws:        wt.ws,
-		head:      wt.head,
+		tableName:          wt.userTblName,
+		workspaceTableName: wt.Name(),
+		headSch:            wt.headSchema,
+		schemaLen:          cols,
+		ws:                 wt.ws,
+		head:               wt.head,
+		modifiable:         wt.modifiable,
 	}
 
 	return &WorkspaceTableUpdater{
@@ -407,6 +427,12 @@ func NewWorkspaceTable(ctx *sql.Context, workspaceTableName string, tableName do
 		fromSch = toSch
 	}
 
+	modifiable := false
+	if fromSch != nil && toSch != nil {
+		// TODO: be more intelligent about schema migrations. This is pretty strict, but it's also correct.
+		modifiable = schema.ColCollsAreEqual(fromSch.GetAllCols(), toSch.GetAllCols())
+	}
+
 	sch := sql.NewPrimaryKeySchema(GetDoltWorkspaceBaseSqlSchema())
 	baseDoltSch, err := sqlutil.ToDoltSchema(ctx, head, tableName, sch, head, sql.Collation_Default)
 	if err != nil {
@@ -430,6 +456,7 @@ func NewWorkspaceTable(ctx *sql.Context, workspaceTableName string, tableName do
 		stagedDeltas:  stgDel,
 		workingDeltas: wkDel,
 		headSchema:    fromSch,
+		modifiable:    modifiable,
 	}, nil
 }
 

go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go

@@ -269,7 +269,6 @@ var DoltWorkspaceScriptTests = []queries.ScriptTest{
 			},
 		},
 	},
-
 	{
 		Name: "dolt_workspace_* keyless table",
 		SetUpScript: []string{
@@ -1013,4 +1012,44 @@ var DoltWorkspaceScriptTests = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "dolt_workspace_* schema change partial staging unstaging",
+		SetUpScript: []string{
+			"create table tbl (x int primary key, y int)",
+			"insert into tbl values (42,42)",
+			"insert into tbl values (23,23)",
+			"call dolt_commit('-Am', 'creating table tbl')",
+			"alter table tbl add column why varchar(16)",
+			"update tbl set why = 'fourty two' where x = 42",
+			"update tbl set why = 'twenty three' where x = 23",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select * from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, false, "modified", 23, 23, "twenty three", 23, 23},
+					{1, false, "modified", 42, 42, "fourty two", 42, 42},
+				},
+			},
+			{
+				// TODO: This seems reasonable, but is difficult with current code structure. We'd need to intercept the
+				//       query earlier to determine that the user just wants to reset the table.
+				Query:          "delete from dolt_workspace_tbl",
+				ExpectedErrStr: "dolt_workspace_tbl table is not modifiable due to schema change",
+			},
+			{
+				// TODO: We could support this... see above. This is basically a `dolt add tbl`
+				Query:          "update dolt_workspace_tbl set staged = true",
+				ExpectedErrStr: "dolt_workspace_tbl table is not modifiable due to schema change",
+			},
+			{
+				Query:          "delete from dolt_workspace_tbl where id = 0",
+				ExpectedErrStr: "dolt_workspace_tbl table is not modifiable due to schema change",
+			},
+			{
+				Query:          "update dolt_workspace_tbl set staged = true where id = 0",
+				ExpectedErrStr: "dolt_workspace_tbl table is not modifiable due to schema change",
+			},
+		},
+	},
 }