diff --git a/x/sqlbuilder/insert_statement.go b/x/sqlbuilder/insert_statement.go index 95d58c2..d3b217d 100644 --- a/x/sqlbuilder/insert_statement.go +++ b/x/sqlbuilder/insert_statement.go @@ -6,6 +6,8 @@ import ( "slices" "strings" + "github.com/upfluence/errors" + "github.com/upfluence/sql" ) @@ -88,6 +90,12 @@ var Nothing nothing type OnConflictClause struct { Target *OnConflictTarget Action OnConflictAction + + // Where restricts the conflicting rows the Update action is applied to. + // It is ignored for the Nothing action. Columns have to be qualified with + // the table, using ColumnWithTable, as an unqualified one is ambiguous + // with the EXCLUDED row. + Where PredicateClause } func (occ *OnConflictClause) Clone() *OnConflictClause { @@ -98,6 +106,7 @@ func (occ *OnConflictClause) Clone() *OnConflictClause { return &OnConflictClause{ Target: occ.Target.Clone(), Action: occ.Action.Clone(), + Where: clonePredicateClause(occ.Where), } } @@ -210,6 +219,14 @@ func (is InsertStatement) buildQueries(vvs []map[string]any, qvs map[string]any) if err := oc.Action.WriteTo(&qw, qvs); err != nil { return "", nil, err } + + if _, ok := oc.Action.(Update); ok && oc.Where != nil { + qw.WriteString(" WHERE ") + + if err := oc.Where.WriteTo(&qw, qvs); err != nil { + return "", nil, errors.WithStack(err) + } + } } switch rs := is.returnings(); len(rs) { diff --git a/x/sqlbuilder/insert_statement_test.go b/x/sqlbuilder/insert_statement_test.go index b5e3702..55f39f2 100644 --- a/x/sqlbuilder/insert_statement_test.go +++ b/x/sqlbuilder/insert_statement_test.go @@ -104,6 +104,39 @@ func TestInsertQuery(t *testing.T) { stmt: "INSERT INTO foo(buz, bar, fiz) VALUES ($1, $2, $3) ON CONFLICT (buz) DO UPDATE SET bar = EXCLUDED.bar, fiz = EXCLUDED.fiz", args: []any{1, 2, 3}, }, + { + name: "with on conflict update and where", + is: InsertStatement{ + Table: "foo", + Fields: []Marker{Column("buz")}, + OnConfict: &OnConflictClause{ + Target: &OnConflictTarget{ + Fields: []Marker{Column("buz")}, + }, + Action: Update{ + Column("bar"), + }, + Where: StaticEq(Column("biz"), "fiz"), + }, + }, + vs: map[string]any{"buz": 1, "bar": 2}, + stmt: "INSERT INTO foo(buz) VALUES ($1) ON CONFLICT (buz) DO UPDATE SET bar = $2 WHERE biz = $3", + args: []any{1, 2, "fiz"}, + }, + { + name: "with on conflict nothing ignores where", + is: InsertStatement{ + Table: "foo", + Fields: []Marker{Column("buz")}, + OnConfict: &OnConflictClause{ + Action: Nothing, + Where: StaticEq(Column("biz"), "fiz"), + }, + }, + vs: map[string]any{"buz": 1}, + stmt: "INSERT INTO foo(buz) VALUES ($1) ON CONFLICT DO NOTHING", + args: []any{1}, + }, { name: "with returning + isQuery", is: InsertStatement{ diff --git a/x/sqlbuilder/upserter/execer.go b/x/sqlbuilder/upserter/execer.go index 25cd37d..bf085f3 100644 --- a/x/sqlbuilder/upserter/execer.go +++ b/x/sqlbuilder/upserter/execer.go @@ -67,6 +67,7 @@ func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { Fields: stmt.QueryValues, }, Action: action, + Where: stmt.PredicateClause, }, Returning: stmt.Returning, }, @@ -112,6 +113,10 @@ func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { clauses[i] = sqlbuilder.Eq(v) } + if stmt.PredicateClause != nil { + clauses = append(clauses, stmt.PredicateClause) + } + return qb.PrepareUpdate( sqlbuilder.UpdateStatement{ Table: stmt.Table, @@ -194,10 +199,13 @@ func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer { e.is.Fields[len(stmt.QueryValues)+len(stmt.SetValues)+i] = iv } - clause := sqlbuilder.And(clauses...) + e.ss.WhereClause = sqlbuilder.And(clauses...) + + if stmt.PredicateClause != nil { + clauses = append(clauses, stmt.PredicateClause) + } - e.ss.WhereClause = clause - e.us.WhereClause = clause + e.us.WhereClause = sqlbuilder.And(clauses...) return &e } diff --git a/x/sqlbuilder/upserter/upserter.go b/x/sqlbuilder/upserter/upserter.go index caff55e..b5941f7 100644 --- a/x/sqlbuilder/upserter/upserter.go +++ b/x/sqlbuilder/upserter/upserter.go @@ -33,6 +33,13 @@ type Statement struct { Returning *sql.Returning + // PredicateClause restricts the update of an already existing row: the + // row is only updated if the predicate is satisfied. Rows that do not + // exist yet are still inserted. Columns have to be qualified with the + // table, using sqlbuilder.ColumnWithTable, as the unqualified ones are + // ambiguous in an INSERT ... ON CONFLICT. + PredicateClause sqlbuilder.PredicateClause + Mode Mode QueryConstrained bool diff --git a/x/sqlbuilder/upserter/upserter_test.go b/x/sqlbuilder/upserter/upserter_test.go index d184045..a12103d 100644 --- a/x/sqlbuilder/upserter/upserter_test.go +++ b/x/sqlbuilder/upserter/upserter_test.go @@ -7,6 +7,7 @@ import ( _ "github.com/mattn/go-sqlite3" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/upfluence/log" "github.com/upfluence/sql" @@ -292,6 +293,120 @@ func TestUpserterInsertValue(t *testing.T) { }) } +func TestUpserterPredicateClause(t *testing.T) { + 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 + + haveQueryConstrained bool + haveExisting bool + havePredicate sqlbuilder.PredicateClause + + wantAffected int64 + wantZ string + }{ + { + name: "no existing row, failing predicate still inserts", + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "nope"), + wantAffected: 1, + wantZ: "new", + }, + { + name: "no existing row, failing predicate still inserts, query constrained", + haveQueryConstrained: true, + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "nope"), + wantAffected: 1, + wantZ: "new", + }, + { + name: "existing row, valid predicate updates", + haveExisting: true, + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "bar"), + wantAffected: 1, + wantZ: "new", + }, + { + name: "existing row, valid predicate updates, query constrained", + haveQueryConstrained: true, + haveExisting: true, + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "bar"), + wantAffected: 1, + wantZ: "new", + }, + { + name: "existing row, failing predicate does not update", + haveExisting: true, + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "nope"), + wantAffected: 0, + wantZ: "old", + }, + { + name: "existing row, failing predicate does not update, query constrained", + haveQueryConstrained: true, + haveExisting: true, + havePredicate: sqlbuilder.StaticEq(sqlbuilder.ColumnWithTable("y", "foo", "y"), "nope"), + wantAffected: 0, + wantZ: "old", + }, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := db.Exec(ctx, "DELETE FROM foo") + require.NoError(t, err) + + if tt.haveExisting { + _, err := db.Exec( + ctx, + "INSERT INTO foo(x, y, z) VALUES ($1, $2, $3)", + "foo", + "bar", + "old", + ) + require.NoError(t, err) + } + + res, err := u.PrepareUpsert( + UpsertStatement{ + Table: "foo", + QueryValues: []sqlbuilder.Marker{sqlbuilder.Column("x")}, + SetValues: []sqlbuilder.Marker{sqlbuilder.Column("z")}, + InsertValues: []sqlbuilder.Marker{sqlbuilder.Column("y")}, + PredicateClause: tt.havePredicate, + QueryConstrained: tt.haveQueryConstrained, + }, + ).Exec( + ctx, + map[string]any{"x": "foo", "y": "bar", "z": "new"}, + ) + + require.NoError(t, err) + assertResultAffected(t, res, tt.wantAffected) + + var z string + + require.NoError( + t, + db.QueryRow(ctx, "SELECT z FROM foo WHERE x = $1", "foo").Scan(&z), + ) + assert.Equal(t, tt.wantZ, z) + }) + } + }) +} + func TestUpserterUpdateOnly(t *testing.T) { sqltest.NewTestCase( sqltest.WithMigratorFunc(