x/sqlbuilder/upserter: add PredicateClause to restrict updates

x/sqlbuilder/insert_statement.go

@@ -6,6 +6,8 @@
 	"slices"
 	"strings"
 
+	"github.com/upfluence/errors"
+
 	"github.com/upfluence/sql"
 )
 
@@ -32,10 +34,10 @@
 }
 
 func (oct *OnConflictTarget) WriteTo(qw QueryWriter, vs map[string]any) error {
 	io.WriteString(qw, "(")
 
 	for i, f := range oct.Fields {
 		io.WriteString(qw, ColumnName(f))
 
 		if i < len(oct.Fields)-1 {
 			io.WriteString(qw, ", ")
@@ -88,6 +90,12 @@
 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 @@
 	return &OnConflictClause{
 		Target: occ.Target.Clone(),
 		Action: occ.Action.Clone(),
+		Where:  clonePredicateClause(occ.Where),
 	}
 }
 
@@ -151,7 +160,7 @@
 	return is.buildQueries([]map[string]any{qvs}, qvs)
 }
 
 func (is InsertStatement) buildQueries(vvs []map[string]any, qvs map[string]any) (string, []any, error) {
 	var qw queryWriter
 
 	if len(is.Fields) == 0 {
@@ -210,6 +219,14 @@
 		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) {

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{

x/sqlbuilder/upserter/execer.go

@@ -36,7 +36,7 @@
 	mode Mode
 }
 
 func newExecer(te txExecutor, stmt Statement) sqlbuilder.Execer {
 	if len(stmt.QueryValues) == 0 {
 		return errExecer{errNoQueryValues}
 	}
@@ -67,6 +67,7 @@
 							Fields: stmt.QueryValues,
 						},
 						Action: action,
+						Where:  stmt.PredicateClause,
 					},
 					Returning: stmt.Returning,
 				},
@@ -112,6 +113,10 @@
 				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 @@
 		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
 }
@@ -217,7 +225,7 @@
 	return err
 }
 
 func (e *execer) Exec(ctx context.Context, vs map[string]any) (sql.Result, error) {
 	var (
 		res    sql.Result
 		lastID int64

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

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(