diff --git a/go/libraries/doltcore/sqle/dtables/workspace_table.go b/go/libraries/doltcore/sqle/dtables/workspace_table.go index 325ab98916c..08e1a988f72 100644 --- a/go/libraries/doltcore/sqle/dtables/workspace_table.go +++ b/go/libraries/doltcore/sqle/dtables/workspace_table.go @@ -15,6 +15,7 @@ package dtables import ( + "bytes" "context" "errors" "fmt" @@ -85,6 +86,12 @@ type WorkspaceTableModifier struct { type WorkspaceTableUpdater struct { WorkspaceTableModifier + // uniqueCols holds the ordinals of the columns covered by the table's unique indexes. A change touching none of + // them can neither free nor take a unique key, so it is applied directly. + uniqueCols []int + // deferredPuts holds the rows this statement will write, applied in [WorkspaceTableUpdater.StatementComplete] + // after every delete in the statement so that a unique key freed by one row can be taken by another. + deferredPuts []sql.Row } var _ sql.RowUpdater = (*WorkspaceTableUpdater)(nil) @@ -198,9 +205,104 @@ func (wtu *WorkspaceTableUpdater) Update(ctx *sql.Context, old sql.Row, new sql. if isDelete { return tableWriter.Delete(ctx, fromRow) - } else { + } + + if !wtu.changesUniqueKey(fromRow, toRow) { return tableWriter.Update(ctx, fromRow, toRow) } + + // Applying the delete half now and deferring the put makes every delete in the statement land before any put. + isInsert := true + for _, val := range fromRow { + if val != nil { + isInsert = false + break + } + } + if !isInsert { + if err := tableWriter.Delete(ctx, fromRow); err != nil { + return err + } + } + wtu.deferredPuts = append(wtu.deferredPuts, append(sql.Row{}, toRow...)) + return nil +} + +// uniqueColumnOrdinals returns the ordinals of the columns covered by the unique indexes of |sch|. +func uniqueColumnOrdinals(sch schema.Schema) []int { + tagToIdx := sch.GetAllCols().TagToIdx + var ords []int + seen := make(map[int]struct{}) + for _, idx := range sch.Indexes().AllIndexes() { + if !idx.IsUnique() { + continue + } + for _, tag := range idx.IndexedColumnTags() { + ord := tagToIdx[tag] + if _, ok := seen[ord]; !ok { + seen[ord] = struct{}{} + ords = append(ords, ord) + } + } + } + return ords +} + +// changesUniqueKey reports whether the change from |fromRow| to |toRow| frees or takes a unique key. Values whose +// equality cannot be determined report a change, which safely defers the write. +func (wtu *WorkspaceTableUpdater) changesUniqueKey(fromRow, toRow sql.Row) bool { + for _, ord := range wtu.uniqueCols { + from, to := fromRow[ord], toRow[ord] + if fromBytes, ok := from.([]byte); ok { + toBytes, ok := to.([]byte) + if !ok || !bytes.Equal(fromBytes, toBytes) { + return true + } + } else if from != to { + return true + } + } + return false +} + +// StatementComplete applies the puts the statement deferred and then flushes as usual. A unique key violation among +// the puts discards the partially applied statement, so a failed staging leaves the roots untouched. +func (wtu *WorkspaceTableUpdater) StatementComplete(ctx *sql.Context) error { + if wtu.err != nil { + return *wtu.err + } + if err := wtu.applyDeferredPuts(ctx); err != nil { + if wtu.tableWriter != nil { + _ = (*wtu.tableWriter).DiscardChanges(ctx, err) + } + wtu.tableWriter = nil + wtu.sessionWriter = nil + return err + } + return wtu.statementComplete(ctx) +} + +// applyDeferredPuts writes the rows this statement deferred and clears the queue. +func (wtu *WorkspaceTableUpdater) applyDeferredPuts(ctx *sql.Context) error { + defer func() { wtu.deferredPuts = nil }() + if len(wtu.deferredPuts) == 0 { + return nil + } + if wtu.tableWriter == nil { + return fmt.Errorf("Runtime error: table writer is nil") + } + for _, row := range wtu.deferredPuts { + if err := (*wtu.tableWriter).Insert(ctx, row); err != nil { + return err + } + } + return nil +} + +// DiscardChanges drops the deferred puts along with the writer's pending changes. +func (wtu *WorkspaceTableUpdater) DiscardChanges(ctx *sql.Context, errorEncountered error) error { + wtu.deferredPuts = nil + return wtu.WorkspaceTableModifier.DiscardChanges(ctx, errorEncountered) } func (wtu *WorkspaceTableUpdater) Close(c *sql.Context) error { @@ -520,7 +622,8 @@ func (wt *WorkspaceTable) Updater(ctx *sql.Context) sql.RowUpdater { } return &WorkspaceTableUpdater{ - modifier, + WorkspaceTableModifier: modifier, + uniqueCols: uniqueColumnOrdinals(wt.headSchema), } } diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go b/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go index 562faf6e7c4..dce1f465698 100644 --- a/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go +++ b/go/libraries/doltcore/sqle/enginetest/dolt_queries_workspace.go @@ -17,6 +17,8 @@ package enginetest import ( "github.com/dolthub/go-mysql-server/enginetest/queries" "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/plan" + "github.com/dolthub/go-mysql-server/sql/types" ) var DoltWorkspaceScriptTests = []queries.ScriptTest{ @@ -1179,4 +1181,171 @@ var DoltWorkspaceScriptTests = []queries.ScriptTest{ }, }, }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* stage a delete and insert that reuse a unique key", + SetUpScript: []string{ + "create table t (id int primary key, u int, unique index ix (u));", + "call dolt_commit('-Am', 'creating table t');", + "insert into t values (5, 10);", + "call dolt_commit('-am', 'insert row with u=10');", + "delete from t where id = 5;", + "insert into t values (1, 10);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_t set staged = 1 where diff_type = 'added'", + // Staging only the added row leaves both the old and the new u=10 row in staging. + ExpectedErrStr: "duplicate unique key given: [10]", + }, + { + Query: "select diff_type, staged from dolt_workspace_t order by diff_type", + // The failed staging must not stage anything. + Expected: []sql.Row{{"added", false}, {"removed", false}}, + }, + { + Query: "select id, u from t as of 'STAGED'", + Expected: []sql.Row{{5, 10}}, + }, + { + Query: "update dolt_workspace_t set staged = 1", + Expected: []sql.Row{{types.OkResult{RowsAffected: 2, Info: plan.UpdateInfo{Matched: 2, Updated: 2}}}}, + }, + { + Query: "select id, u from t as of 'STAGED'", + Expected: []sql.Row{{1, 10}}, + }, + { + Query: "select id from t as of 'STAGED' where u = 10", + // A lookup through the unique index must find the staged row, proving the index is consistent. + Expected: []sql.Row{{1}}, + }, + }, + }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* stage a delete and insert that reuse a composite unique key", + SetUpScript: []string{ + "create table foo (id varchar(36) primary key);", + "create table bar (id varchar(36) primary key, foo_id varchar(36) not null, sequence int, " + + "foreign key (foo_id) references foo(id) on delete cascade on update cascade, " + + "unique index idx_foo_id_sequence (foo_id, sequence));", + "call dolt_commit('-Am', 'creating tables foo and bar');", + "insert into foo values ('f1');", + "insert into bar values ('b1', 'f1', 10);", + "call dolt_commit('-am', 'insert foo and bar with sequence 10');", + "delete from bar where id = 'b1';", + "insert into bar values ('b2', 'f1', 10);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_bar set staged = 1", + Expected: []sql.Row{{types.OkResult{RowsAffected: 2, Info: plan.UpdateInfo{Matched: 2, Updated: 2}}}}, + }, + { + Query: "select id, sequence from bar as of 'STAGED'", + Expected: []sql.Row{{"b2", 10}}, + }, + }, + }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* stage two modifies that swap a unique key", + SetUpScript: []string{ + "create table t (id int primary key, u int, unique index ix (u));", + "call dolt_commit('-Am', 'creating table t');", + "insert into t values (1, 10), (2, 20);", + "call dolt_commit('-am', 'insert two rows');", + "update t set u = 99 where id = 1;", + "update t set u = 10 where id = 2;", + "update t set u = 20 where id = 1;", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_t set staged = 1", + Expected: []sql.Row{{types.OkResult{RowsAffected: 2, Info: plan.UpdateInfo{Matched: 2, Updated: 2}}}}, + }, + { + Query: "select id, u from t as of 'STAGED' order by id", + Expected: []sql.Row{{1, 20}, {2, 10}}, + }, + }, + }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* keyless table stage a delete and insert that reuse a unique key", + SetUpScript: []string{ + "create table t (a int, u int, unique index ix (u));", + "call dolt_commit('-Am', 'creating keyless table t');", + "insert into t values (1, 10);", + "call dolt_commit('-am', 'insert row with u=10');", + "delete from t where a = 1;", + "insert into t values (2, 10);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_t set staged = 1 where diff_type = 'added'", + // Staging only the added row leaves both the old and the new u=10 row in staging. + ExpectedErrStr: "duplicate unique key given: [10]", + }, + { + Query: "select a, u from t as of 'STAGED'", + Expected: []sql.Row{{1, 10}}, + }, + { + Query: "update dolt_workspace_t set staged = 1", + Expected: []sql.Row{{types.OkResult{RowsAffected: 2, Info: plan.UpdateInfo{Matched: 2, Updated: 2}}}}, + }, + { + Query: "select a, u from t as of 'STAGED'", + Expected: []sql.Row{{2, 10}}, + }, + }, + }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* stage a modify that leaves the unique key untouched", + SetUpScript: []string{ + "create table t (id int primary key, u int, v int, unique index ix (u));", + "call dolt_commit('-Am', 'creating table t');", + "insert into t values (1, 10, 100), (2, 20, 200);", + "call dolt_commit('-am', 'insert two rows');", + "update t set v = 101 where id = 1;", + "delete from t where id = 2;", + "insert into t values (3, 20, 300);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_t set staged = 1", + Expected: []sql.Row{{types.OkResult{RowsAffected: 3, Info: plan.UpdateInfo{Matched: 3, Updated: 3}}}}, + }, + { + Query: "select id, u, v from t as of 'STAGED' order by id", + Expected: []sql.Row{{1, 10, 101}, {3, 20, 300}}, + }, + }, + }, + { + // See https://github.com/dolthub/dolt/issues/11195. + Name: "dolt_workspace_* stage rows that share a null unique key", + SetUpScript: []string{ + "create table t (id int primary key, u int, unique index ix (u));", + "call dolt_commit('-Am', 'creating table t');", + "insert into t values (1, 10);", + "call dolt_commit('-am', 'insert row with u=10');", + "update t set u = null where id = 1;", + "insert into t values (2, null);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "update dolt_workspace_t set staged = 1", + // NULLs do not collide under a unique index. + Expected: []sql.Row{{types.OkResult{RowsAffected: 2, Info: plan.UpdateInfo{Matched: 2, Updated: 2}}}}, + }, + { + Query: "select id, u from t as of 'STAGED' order by id", + Expected: []sql.Row{{1, nil}, {2, nil}}, + }, + }, + }, }