diff --git a/x/sqlbuilder/upserter/execer.go b/x/sqlbuilder/upserter/execer.go index bf085f3..2c4bad8 100644 --- a/x/sqlbuilder/upserter/execer.go +++ b/x/sqlbuilder/upserter/execer.go @@ -36,99 +36,185 @@ type execer struct { mode Mode } +// MultiExecer executes the statement for several rows at once, all or none of +// them, each row holding the values of all the markers of the statement as Exec +// does. The result holds the sum of the rows affected. The rows must not be +// empty and the query values unique. +// +// The key constrained Insert and Upsert modes do it with a single INSERT, which +// is bound by the parameter limit of the database as the rows are not chunked. +// Without the key constraint, the statements of every row run one after the +// other in a single transaction, which saves the transaction of each row but +// not its round trips. +// +// The key constrained Update mode, the key constrained Upsert mode with a +// PredicateClause and the statements with Returning do not implement it, so +// that the capability can be discovered with a type assertion. +type MultiExecer interface { + sqlbuilder.Execer + + MultiExec(context.Context, []map[string]any) (sql.Result, error) +} + +type multiExecer struct { + sqlbuilder.Execer + + ie *sqlbuilder.InsertExecer +} + +func (me *multiExecer) MultiExec(ctx context.Context, vvs []map[string]any) (sql.Result, error) { + res, err := me.ie.MultiExec(ctx, vvs, nil) + + return res, errors.WithStack(err) +} + +func insertFields(stmt Statement) []sqlbuilder.Marker { + return append( + append( + append([]sqlbuilder.Marker{}, stmt.QueryValues...), + stmt.InsertValues..., + ), + stmt.SetValues..., + ) +} + +func excludedUpdate(ms []sqlbuilder.Marker) sqlbuilder.Update { + res := make(sqlbuilder.Update, len(ms)) + + for i, m := range ms { + res[i] = sqlbuilder.Excluded(m) + } + + return res +} + func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { if len(stmt.QueryValues) == 0 { return errExecer{errNoQueryValues} } if stmt.QueryConstrained { - qb := sqlbuilder.QueryBuilder{Queryer: te.queryer()} + return newQueryConstrainedExecer(te.queryer(), stmt) + } - switch stmt.mode() { - case Upsert: - var action sqlbuilder.OnConflictAction = sqlbuilder.Nothing + return newTxExecer(te, stmt) +} - if len(stmt.SetValues) > 0 { - action = sqlbuilder.Update(stmt.SetValues) - } +// newQueryConstrainedExecer relies on the key constraint of the query values +// to write the row with a single statement. +func newQueryConstrainedExecer(q sql.Queryer, stmt Statement) sqlbuilder.Execer { + qb := sqlbuilder.QueryBuilder{Queryer: q} - e := qb.PrepareInsert( - sqlbuilder.InsertStatement{ - Table: stmt.Table, - Fields: append( - append( - append([]sqlbuilder.Marker{}, stmt.QueryValues...), - stmt.InsertValues..., - ), - stmt.SetValues..., - ), - OnConfict: &sqlbuilder.OnConflictClause{ - Target: &sqlbuilder.OnConflictTarget{ - Fields: stmt.QueryValues, - }, - Action: action, - Where: stmt.PredicateClause, + switch stmt.mode() { + case Upsert: + var action sqlbuilder.OnConflictAction = sqlbuilder.Nothing + + if len(stmt.SetValues) > 0 { + action = sqlbuilder.Update(stmt.SetValues) + } + + e := qb.PrepareInsert( + sqlbuilder.InsertStatement{ + Table: stmt.Table, + Fields: insertFields(stmt), + OnConfict: &sqlbuilder.OnConflictClause{ + Target: &sqlbuilder.OnConflictTarget{ + Fields: stmt.QueryValues, }, - Returning: stmt.Returning, + Action: action, + Where: stmt.PredicateClause, }, - ) + Returning: stmt.Returning, + }, + ) - return &sqlbuilder.RetryExecer{ - Execer: e, - ShouldRetry: func(err error) bool { - var ce sql.ConstraintError + var re sqlbuilder.Execer = &sqlbuilder.RetryExecer{ + Execer: e, + ShouldRetry: func(err error) bool { + var ce sql.ConstraintError - if !errors.As(err, &ce) { - return false - } + if !errors.As(err, &ce) { + return false + } - return ce.Type == sql.PrimaryKey || ce.Type == sql.Unique - }, - RetryCount: -1, - } - case Insert: - return qb.PrepareInsert( + return ce.Type == sql.PrimaryKey || ce.Type == sql.Unique + }, + RetryCount: -1, + } + + // the predicate is bound to the values of a single row, which the + // ON CONFLICT clause of a multi rows INSERT cannot hold + if stmt.Returning != nil || stmt.PredicateClause != nil { + return re + } + + var multiAction sqlbuilder.OnConflictAction = sqlbuilder.Nothing + + if len(stmt.SetValues) > 0 { + multiAction = excludedUpdate(stmt.SetValues) + } + + return &multiExecer{ + Execer: re, + ie: qb.PrepareInsert( sqlbuilder.InsertStatement{ - Table: stmt.Table, - Fields: append( - append( - append([]sqlbuilder.Marker{}, stmt.QueryValues...), - stmt.InsertValues..., - ), - stmt.SetValues..., - ), + Table: stmt.Table, + Fields: insertFields(stmt), OnConfict: &sqlbuilder.OnConflictClause{ Target: &sqlbuilder.OnConflictTarget{ Fields: stmt.QueryValues, }, - Action: sqlbuilder.Nothing, + Action: multiAction, }, - Returning: stmt.Returning, }, - ) - case Update: - clauses := make([]sqlbuilder.PredicateClause, len(stmt.QueryValues)) + ), + } + case Insert: + ie := qb.PrepareInsert( + sqlbuilder.InsertStatement{ + Table: stmt.Table, + Fields: insertFields(stmt), + OnConfict: &sqlbuilder.OnConflictClause{ + Target: &sqlbuilder.OnConflictTarget{ + Fields: stmt.QueryValues, + }, + Action: sqlbuilder.Nothing, + }, + Returning: stmt.Returning, + }, + ) - for i, v := range stmt.QueryValues { - clauses[i] = sqlbuilder.Eq(v) - } + if stmt.Returning != nil { + return ie + } - if stmt.PredicateClause != nil { - clauses = append(clauses, stmt.PredicateClause) - } + return &multiExecer{Execer: ie, ie: ie} + case Update: + clauses := make([]sqlbuilder.PredicateClause, len(stmt.QueryValues)) - return qb.PrepareUpdate( - sqlbuilder.UpdateStatement{ - Table: stmt.Table, - Fields: stmt.SetValues, - WhereClause: sqlbuilder.And(clauses...), - }, - ) - default: - return errExecer{fmt.Errorf("invalid mode: %d", stmt.Mode)} + for i, v := range stmt.QueryValues { + clauses[i] = sqlbuilder.Eq(v) } + + if stmt.PredicateClause != nil { + clauses = append(clauses, stmt.PredicateClause) + } + + return qb.PrepareUpdate( + sqlbuilder.UpdateStatement{ + Table: stmt.Table, + Fields: stmt.SetValues, + WhereClause: sqlbuilder.And(clauses...), + }, + ) + default: + return errExecer{fmt.Errorf("invalid mode: %d", stmt.Mode)} } +} +// newTxExecer looks the row up and writes it accordingly within a +// transaction, as the table holds no key constraint on the query values. +func newTxExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { var selectClauses = []sqlbuilder.Marker{oneMarker} for _, m := range stmt.SetValues { @@ -207,6 +293,10 @@ func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { e.us.WhereClause = sqlbuilder.And(clauses...) + if stmt.Returning != nil { + return singleExecer{Execer: &e} + } + return &e } @@ -225,9 +315,16 @@ func (am *assertMarker) WriteTo(qw sqlbuilder.QueryWriter, vs map[string]any) er return err } -func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error) { +// rowArgs holds the arguments of the statements of a row, the variables they +// are scanned into included. +type rowArgs struct { + qvs map[string]any + existing map[string]any + lastID *int64 +} + +func (e *execer) prepareRow(vs map[string]any) (rowArgs, error) { var ( - res sql.Result lastID int64 one int64 @@ -239,7 +336,7 @@ func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error v, ok := vs[f] if !ok { - return nil, sqlbuilder.ErrMissingKey{Key: f} + return rowArgs{}, sqlbuilder.ErrMissingKey{Key: f} } qvs[f] = v @@ -249,7 +346,7 @@ func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error v, ok := vs[f] if !ok { - return nil, sqlbuilder.ErrMissingKey{Key: f} + return rowArgs{}, sqlbuilder.ErrMissingKey{Key: f} } var val sql.NullBool @@ -263,54 +360,128 @@ func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error existing[m.Binding()] = &lastID } - return res, e.te.executeTx(ctx, func(q sql.Queryer) error { - var ( - err error + return rowArgs{qvs: qvs, existing: existing, lastID: &lastID}, nil +} - qb = sqlbuilder.QueryBuilder{Queryer: q} - ) +// execRow runs the statements of a row. It returns sql.ErrRollback when +// nothing has to be written. +func (e *execer) execRow(ctx context.Context, q sql.Queryer, r rowArgs, vs map[string]any) (sql.Result, error) { + var ( + res sql.Result + err error - switch err = qb.PrepareSelect(e.ss).QueryRow(ctx, qvs).Scan(existing); err { - case nil: - pristine := true + qb = sqlbuilder.QueryBuilder{Queryer: q} + ) - for _, sf := range e.sfs { - if val := existing[sf].(*sql.NullBool); !val.Bool { - pristine = false + err = qb.PrepareSelect(e.ss).QueryRow(ctx, r.qvs).Scan(r.existing) - break - } - } + switch { + case err == nil: + pristine := true - if pristine || e.mode&Update == 0 { - res = driver.RowsAffected(0) + for _, sf := range e.sfs { + if val := r.existing[sf].(*sql.NullBool); !val.Bool { + pristine = false - if e.returningMarker != nil { - res = lastIDResult{Result: res, id: lastID} - } + break + } + } + + if pristine || e.mode&Update == 0 { + res = driver.RowsAffected(0) - return sql.ErrRollback + if e.returningMarker != nil { + res = lastIDResult{Result: res, id: *r.lastID} } - res, err = qb.PrepareUpdate(e.us).Exec(ctx, vs) + return res, sql.ErrRollback + } + + res, err = qb.PrepareUpdate(e.us).Exec(ctx, vs) + + if err == nil && e.returningMarker != nil { + res = lastIDResult{Result: res, id: *r.lastID} + } + case errors.Is(err, sql.ErrNoRows): + if e.mode&Insert == 0 { + return driver.RowsAffected(0), sql.ErrRollback + } - if err == nil && e.returningMarker != nil { - res = lastIDResult{Result: res, id: lastID} + res, err = qb.PrepareInsert(e.is).Exec(ctx, vs) + default: + } + + // the error is forwarded as it is, as Exec always did + return res, err //nolint:wrapcheck +} + +func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error) { + r, err := e.prepareRow(vs) + + if err != nil { + return nil, err + } + + var res sql.Result + + err = e.te.executeTx(ctx, func(q sql.Queryer) error { + var err error + + res, err = e.execRow(ctx, q, r, vs) + + return err + }) + + return res, err +} + +// MultiExec runs the statements of every row in a single transaction, one row +// after the other as each of them depends on the content of the table. The +// result holds the sum of the rows affected. +func (e *execer) MultiExec(ctx context.Context, vvs []map[string]any) (sql.Result, error) { + rs := make([]rowArgs, len(vvs)) + + for i, vs := range vvs { + r, err := e.prepareRow(vs) + + if err != nil { + return nil, err + } + + rs[i] = r + } + + var total int64 + + err := e.te.executeTx(ctx, func(q sql.Queryer) error { + total = 0 + + for i, r := range rs { + res, err := e.execRow(ctx, q, r, vvs[i]) + + // a row with nothing to write must not roll back the others + if err != nil && !errors.Is(err, sql.ErrRollback) { + return err } - case sql.ErrNoRows: - if e.mode&Insert == 0 { - res = driver.RowsAffected(0) - return sql.ErrRollback + + n, err := res.RowsAffected() + + if err != nil { + return errors.Wrap(err, "cant read the affected rows") } - res, err = qb.PrepareInsert(e.is).Exec(ctx, vs) - default: + total += n } - return err + return nil }) + + return driver.RowsAffected(total), err } +// singleExecer hides the MultiExec method of the executor it wraps. +type singleExecer struct{ sqlbuilder.Execer } + type lastIDResult struct { sql.Result diff --git a/x/sqlbuilder/upserter/upserter_test.go b/x/sqlbuilder/upserter/upserter_test.go index a12103d..7b9cd6b 100644 --- a/x/sqlbuilder/upserter/upserter_test.go +++ b/x/sqlbuilder/upserter/upserter_test.go @@ -8,6 +8,7 @@ import ( _ "github.com/mattn/go-sqlite3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/upfluence/errors" "github.com/upfluence/log" "github.com/upfluence/sql" @@ -628,3 +629,315 @@ func TestInTxUpserterPristine(t *testing.T) { assertResultAffected(t, res, 0) }) } + +func readFoo(t *testing.T, db sql.DB) map[string][2]string { + t.Helper() + + cur, err := db.Query(context.Background(), "SELECT x, y, z FROM foo") + + require.NoError(t, err) + + res := make(map[string][2]string) + + require.NoError( + t, + sql.ScrollCursor( + cur, + func(sc sql.Scanner) error { + var x, y, z string + + if err := sc.Scan(&x, &y, &z); err != nil { + return errors.WithStack(err) + } + + res[x] = [2]string{y, z} + + return nil + }, + ), + ) + + return res +} + +func TestUpserterMultiExec(t *testing.T) { + var ( + x = sqlbuilder.Column("x") + y = sqlbuilder.Column("y") + z = sqlbuilder.Column("z") + ) + + sqltest.NewTestCase( + sqltest.WithMigratorFunc( + buildMigrator( + map[string]string{ + "1_initial.up.sql": "CREATE TABLE foo (x TEXT, y TEXT, z TEXT)", + "1_initial.down.sql": "DROP TABLE foo", + "2_initial.up.sql": "CREATE UNIQUE INDEX foo_pk ON foo(x)", + "2_initial.down.sql": "DROP INDEX foo_pk", + }, + ), + ), + ).Run(t, func(t *testing.T, db sql.DB) { + ctx := context.Background() + u := Upserter{DB: db} + + for _, tt := range []struct { + name string + + haveSetup string + haveStmt Statement + haveRows [][]map[string]any + + wantAffected []int64 + wantRows map[string][2]string + }{ + { + name: "upsert updates every row with its own values", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + SetValues: []sqlbuilder.Marker{y}, + InsertValues: []sqlbuilder.Marker{z}, + QueryConstrained: true, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "1", "z": "1"}, + {"x": "b", "y": "2", "z": "2"}, + }, + { + {"x": "b", "y": "3", "z": "9"}, + {"x": "c", "y": "4", "z": "4"}, + }, + }, + wantAffected: []int64{2, 2}, + wantRows: map[string][2]string{ + "a": {"1", "1"}, + "b": {"3", "2"}, + "c": {"4", "4"}, + }, + }, + { + name: "upsert without set values only inserts", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + InsertValues: []sqlbuilder.Marker{y, z}, + QueryConstrained: true, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "1", "z": "1"}, + {"x": "b", "y": "2", "z": "2"}, + }, + { + {"x": "b", "y": "3", "z": "3"}, + {"x": "c", "y": "4", "z": "4"}, + }, + }, + wantAffected: []int64{2, 1}, + wantRows: map[string][2]string{ + "a": {"1", "1"}, + "b": {"2", "2"}, + "c": {"4", "4"}, + }, + }, + { + name: "without the key constraint upsert writes only what changes", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + SetValues: []sqlbuilder.Marker{y}, + InsertValues: []sqlbuilder.Marker{z}, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "1", "z": "1"}, + {"x": "b", "y": "2", "z": "2"}, + }, + { + {"x": "b", "y": "3", "z": "9"}, + {"x": "c", "y": "4", "z": "4"}, + {"x": "a", "y": "1", "z": "1"}, + }, + }, + wantAffected: []int64{2, 2}, + wantRows: map[string][2]string{ + "a": {"1", "1"}, + "b": {"3", "2"}, + "c": {"4", "4"}, + }, + }, + { + name: "without the key constraint insert skips the existing rows", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + SetValues: []sqlbuilder.Marker{y, z}, + Mode: Insert, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "1", "z": "1"}, + {"x": "b", "y": "2", "z": "2"}, + }, + { + {"x": "b", "y": "3", "z": "3"}, + {"x": "c", "y": "4", "z": "4"}, + }, + }, + wantAffected: []int64{2, 1}, + wantRows: map[string][2]string{ + "a": {"1", "1"}, + "b": {"2", "2"}, + "c": {"4", "4"}, + }, + }, + { + name: "without the key constraint update skips the missing rows", + haveSetup: "INSERT INTO foo(x, y, z) VALUES ('a', '1', '1'), ('b', '2', '2')", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + SetValues: []sqlbuilder.Marker{y}, + Mode: Update, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "5"}, + {"x": "c", "y": "6"}, + }, + }, + wantAffected: []int64{1}, + wantRows: map[string][2]string{ + "a": {"5", "1"}, + "b": {"2", "2"}, + }, + }, + { + name: "insert skips the existing rows", + haveStmt: Statement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{x}, + SetValues: []sqlbuilder.Marker{y, z}, + Mode: Insert, + QueryConstrained: true, + }, + haveRows: [][]map[string]any{ + { + {"x": "a", "y": "1", "z": "1"}, + {"x": "b", "y": "2", "z": "2"}, + }, + { + {"x": "b", "y": "3", "z": "3"}, + {"x": "c", "y": "4", "z": "4"}, + }, + }, + wantAffected: []int64{2, 1}, + wantRows: map[string][2]string{ + "a": {"1", "1"}, + "b": {"2", "2"}, + "c": {"4", "4"}, + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := db.Exec(ctx, "DELETE FROM foo") + + require.NoError(t, err) + + if tt.haveSetup != "" { + _, err := db.Exec(ctx, tt.haveSetup) + + require.NoError(t, err) + } + + me, ok := u.PrepareUpsert(tt.haveStmt).(MultiExecer) + + require.True(t, ok) + + for i, rows := range tt.haveRows { + res, err := me.MultiExec(ctx, rows) + + require.NoError(t, err) + assertResultAffected(t, res, tt.wantAffected[i]) + } + + assert.Equal(t, tt.wantRows, readFoo(t, db)) + }) + } + }) +} + +func TestUpserterMultiExecNotSupported(t *testing.T) { + var ( + u = Upserter{} + + x = []sqlbuilder.Marker{sqlbuilder.Column("x")} + y = []sqlbuilder.Marker{sqlbuilder.Column("y")} + ) + + for _, tt := range []struct { + name string + + haveStmt Statement + }{ + { + name: "update", + haveStmt: Statement{ + Table: "foo", + QueryValues: x, + SetValues: y, + Mode: Update, + QueryConstrained: true, + }, + }, + { + name: "not query constrained with returning", + haveStmt: Statement{ + Table: "foo", + QueryValues: x, + SetValues: y, + Returning: &sql.Returning{Field: "id"}, + }, + }, + { + name: "upsert with returning", + haveStmt: Statement{ + Table: "foo", + QueryValues: x, + SetValues: y, + Returning: &sql.Returning{Field: "id"}, + QueryConstrained: true, + }, + }, + { + name: "upsert with predicate", + haveStmt: Statement{ + Table: "foo", + QueryValues: x, + SetValues: y, + PredicateClause: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "bar"), + QueryConstrained: true, + }, + }, + { + name: "insert with returning", + haveStmt: Statement{ + Table: "foo", + QueryValues: x, + SetValues: y, + Mode: Insert, + Returning: &sql.Returning{Field: "id"}, + QueryConstrained: true, + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + _, ok := u.PrepareUpsert(tt.haveStmt).(MultiExecer) + + assert.False(t, ok) + }) + } +}