Second pass. Downgrades from staging tested

go/libraries/doltcore/sqle/dtables/workspace.go

@@ -130,18 +130,24 @@ func (wtu *WorkspaceTableUpdater) Update(ctx *sql.Context, old sql.Row, new sql.
 		panic("row is nil")
 	}
 
-	if !validateUpdate(old, new) {
+	valid, isStaged := validateWorkspaceUpdate(old, new)
+	if !valid {
 		return errors.New("only update of column 'staged' is allowed")
 	}
-	// old and new are the same. Just use one.
-	new = nil
 
 	// We could do this up front once. NM4. Also - not always the same schema??
 	schemaLen := wtu.headSch.GetAllCols().Size()
 
-	// loop over toRow, and if it's all nil, it's a delete. NM4 - is there a better way to pass through the diff type?
+	// old and new are the same. Just use one.
+	new = nil
+
 	toRow := old[3 : 3+schemaLen]
 	fromRow := old[3+schemaLen : len(old)]
+	if !isStaged {
+		toRow, fromRow = fromRow, toRow
+	}
+
+	// loop over toRow, and if it's all nil, it's a delete. NM4 - is there a better way to pass through the diff type?
 	isDelete := true
 	for _, val := range toRow {
 		if val != nil {
@@ -159,29 +165,46 @@ func (wtu *WorkspaceTableUpdater) Update(ctx *sql.Context, old sql.Row, new sql.
 	}
 }
 
-// validateUpdate returns true IFF old and new row are identical - except the "staged" column. Updating that
-// column to TRUE or FALSE is the only update allowed, and any other update will result in returning an false.
-func validateUpdate(old, new sql.Row) bool {
+// NM4 - is there a better way?
+func isTrue(value interface{}) bool {
+	switch v := value.(type) {
+	case bool:
+		return v
+	case int8:
+		return v != 0
+	default:
+		return false
+	}
+}
+
+// validateWorkspaceUpdate returns true IFF old and new row are identical - except the "staged" column. Updating that
+// column to TRUE or FALSE is the only update allowed, and any other update will result in 'valid' being false. If
+// valid is true, then 'staged' will be the value in the "staged" column of the new row.
+func validateWorkspaceUpdate(old, new sql.Row) (valid, staged bool) {
 	// NM4 - I think it's impossible to have equal rows, but we should rule that out.
 	if old == nil {
-		return false
+		return false, false
 	}
 
 	if len(old) != len(new) {
-		return false
+		return false, false
 	}
 
-	for i := range old {
+	isStaged := false
+
+	for i := range new {
 		if i == 1 {
+			// NM4 - not required in the iterator, right?
+			isStaged = isTrue(new[i])
 			// skip the "staged" column. NM4 - is there a way to not use a constant index here?
 			continue
 		}
 
 		if old[i] != new[i] {
-			return false
+			return false, false
 		}
 	}
-	return true
+	return true, isStaged
 }
 
 func (wtu *WorkspaceTableUpdater) Close(c *sql.Context) error {

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

@@ -530,4 +530,207 @@ var DoltWorkspaceScriptTests = []queries.ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "dolt_workspace_* modifies downgrade from staging",
+		SetUpScript: []string{
+			"create table tbl (pk int primary key, v char(36));",
+			"insert into tbl values (42, UUID());",
+			"insert into tbl values (23, UUID());",
+			"call dolt_commit('-Am', 'creating table tbl');",
+			"update tbl set v = UUID();",
+			"call dolt_add('tbl')",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "modified", 23},
+					{1, true, "modified", 42},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = false where id = 0",
+			},
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "modified", 42},
+					{1, false, "modified", 23},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = 0 where staged = TRUE",
+			},
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, false, "modified", 23},
+					{1, false, "modified", 42},
+				},
+			},
+		},
+	},
+	{
+		Name: "dolt_workspace_* inserts downgrade from staging",
+		SetUpScript: []string{
+			"create table tbl (pk int primary key, v char(36));",
+			"call dolt_commit('-Am', 'creating table tbl');",
+			"insert into tbl values (42, UUID());",
+			"insert into tbl values (23, UUID());",
+			"call dolt_add('tbl')",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "added", 23},
+					{1, true, "added", 42},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = false where id = 0",
+			},
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "added", 42},
+					{1, false, "added", 23},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = 0 where staged = TRUE",
+			},
+			{
+				Query: "select id, staged, diff_type, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, false, "added", 23},
+					{1, false, "added", 42},
+				},
+			},
+		},
+	},
+	{
+		Name: "dolt_workspace_* deletes downgrade from staging",
+		SetUpScript: []string{
+			"create table tbl (pk int primary key, v char(36));",
+			"insert into tbl values (42, UUID());",
+			"insert into tbl values (23, UUID());",
+			"call dolt_commit('-Am', 'creating table tbl');",
+			"delete from tbl",
+			"call dolt_add('tbl')",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "removed", 23, nil},
+					{1, true, "removed", 42, nil},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = false where from_pk = 23",
+			},
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "removed", 42, nil},
+					{1, false, "removed", 23, nil},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = FALSE", // Unstage everything.
+			},
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, false, "removed", 23, nil},
+					{1, false, "removed", 42, nil},
+				},
+			},
+		},
+	},
+	{
+		Name: "dolt_workspace_* complicated mixed updates",
+		SetUpScript: []string{
+			"create table tbl (pk int primary key, v char(36));",
+			"insert into tbl values (42, UUID());",
+			"insert into tbl values (23, UUID());",
+			"call dolt_commit('-Am', 'creating table tbl');",
+			"update tbl set v = UUID() where pk = 42;",
+			"delete from tbl where pk = 23;",
+			"call dolt_add('tbl')",
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "removed", 23, nil},
+					{1, true, "modified", 42, 42},
+				},
+			},
+			{
+				Query: "update tbl set v = UUID() where pk = 42",
+			},
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "removed", 23, nil},
+					{1, true, "modified", 42, 42},
+					{2, false, "modified", 42, 42},
+				},
+			},
+			{
+				Query: "insert into tbl values (23, UUID())",
+			},
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "removed", 23, nil},
+					{1, true, "modified", 42, 42},
+					{2, false, "added", nil, 23},
+					{3, false, "modified", 42, 42},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = 1 where id = 2",
+			},
+			{
+				Query: "select id, staged, diff_type, from_pk, to_pk from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "modified", 23, 23},
+					{1, true, "modified", 42, 42},
+					{2, false, "modified", 42, 42},
+				},
+			},
+		},
+	},
+	{
+		Name: "dolt_workspace_* downgrades keep working changes",
+		SetUpScript: []string{
+			`create table tbl (pk int primary key, v varchar(20))`,
+			`insert into tbl values (42, "inserted")`,
+			`call dolt_commit('-Am', 'creating table tbl')`,
+			`update tbl set v = "staged"`,
+			`call dolt_add('tbl')`,
+			`update tbl set v = "working"`,
+		},
+		Assertions: []queries.ScriptTestAssertion{
+			{
+				Query: "select * from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, true, "modified", 42, "staged", 42, "inserted"},
+					{1, false, "modified", 42, "working", 42, "staged"},
+				},
+			},
+			{
+				Query: "update dolt_workspace_tbl set staged = false where id = 0",
+			},
+			{
+				Query: "select * from dolt_workspace_tbl",
+				Expected: []sql.Row{
+					{0, false, "modified", 42, "working", 42, "inserted"},
+				},
+			},
+		},
+	},
 }