From 9de519d027a4f0cc263c56181dc6aed3b75ae64d Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Tue, 27 Jul 2021 10:58:13 +0300
Subject: [PATCH] dialect/sql: add expression function support

---
 dialect/sql/builder.go      | 43 +++++++++++++++++++++++++++++++++----
 dialect/sql/builder_test.go | 18 ++++++++++++++++
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/dialect/sql/builder.go b/dialect/sql/builder.go
index c683d8d74..c61839a1c 100644
--- a/dialect/sql/builder.go
+++ b/dialect/sql/builder.go
@@ -2543,6 +2543,29 @@ type expr struct {
 
 func (e *expr) Query() (string, []interface{}) { return e.s, e.args }
 
+// ExprFunc returns an expression function that implements the Querier interface.
+//
+//	Update("users").
+//		Set("x", ExprFunc(func(b *Builder) {
+//			// The sql.Builder config (argc and dialect)
+//			// was set before the function was executed.
+//			b.Ident("x").WriteOp(OpAdd).Arg(1)
+//		}))
+//
+func ExprFunc(fn func(*Builder)) Querier {
+	return &exprFunc{fn: fn}
+}
+
+type exprFunc struct {
+	Builder
+	fn func(*Builder)
+}
+
+func (e *exprFunc) Query() (string, []interface{}) {
+	e.fn(&e.Builder)
+	return e.Builder.Query()
+}
+
 // Queries are list of queries join with space between them.
 type Queries []Querier
 
@@ -2688,12 +2711,12 @@ func (b *Builder) Err() error {
 	return fmt.Errorf(br.String())
 }
 
-// An Op represents a predicate operator.
+// An Op represents an operator.
 type Op int
 
-// Predicate operators
 const (
-	OpEQ      Op = iota // logical and.
+	// Predicate operators.
+	OpEQ      Op = iota // =
 	OpNEQ               // <>
 	OpGT                // >
 	OpGTE               // >=
@@ -2704,6 +2727,13 @@ const (
 	OpLike              // LIKE
 	OpIsNull            // IS NULL
 	OpNotNull           // IS NOT NULL
+
+	// Arithmetic operators.
+	OpAdd // +
+	OpSub // -
+	OpMul // *
+	OpDiv // / (Quotient)
+	OpMod // % (Reminder)
 )
 
 var ops = [...]string{
@@ -2718,12 +2748,17 @@ var ops = [...]string{
 	OpLike:    "LIKE",
 	OpIsNull:  "IS NULL",
 	OpNotNull: "IS NOT NULL",
+	OpAdd:     "+",
+	OpSub:     "-",
+	OpMul:     "*",
+	OpDiv:     "/",
+	OpMod:     "%",
 }
 
 // WriteOp writes an operator to the builder.
 func (b *Builder) WriteOp(op Op) *Builder {
 	switch {
-	case op >= OpEQ && op <= OpLike:
+	case op >= OpEQ && op <= OpLike || op >= OpAdd && op <= OpMod:
 		b.Pad().WriteString(ops[op]).Pad()
 	case op == OpIsNull || op == OpNotNull:
 		b.Pad().WriteString(ops[op])
diff --git a/dialect/sql/builder_test.go b/dialect/sql/builder_test.go
index 8b60abcff..520237730 100644
--- a/dialect/sql/builder_test.go
+++ b/dialect/sql/builder_test.go
@@ -1715,3 +1715,21 @@ func TestSelector_UnionOrderBy(t *testing.T) {
 		Query()
 	require.Equal(t, `SELECT * FROM "users" WHERE "active" = $1 UNION SELECT * FROM "old_users1" ORDER BY "users"."whatever"`, query)
 }
+
+func TestUpdateBuilder_SetExpr(t *testing.T) {
+	d := Dialect(dialect.Postgres)
+	excluded := d.Table("excluded")
+	query, args := d.Update("users").
+		Set("name", "Ariel").
+		Set("active", Expr("NOT(active)")).
+		Set("age", Expr(excluded.C("age"))).
+		Set("x", ExprFunc(func(b *Builder) {
+			b.WriteString(excluded.C("x")).WriteString(" || ' (formerly ' || ").Ident("x").WriteString(" || ')'")
+		})).
+		Set("y", ExprFunc(func(b *Builder) {
+			b.Arg("~").WriteOp(OpAdd).WriteString(excluded.C("y")).WriteOp(OpAdd).Arg("~")
+		})).
+		Query()
+	require.Equal(t, `UPDATE "users" SET "name" = $1, "active" = NOT(active), "age" = "excluded"."age", "x" = "excluded"."x" || ' (formerly ' || "x" || ')', "y" = $2 + "excluded"."y" + $3`, query)
+	require.Equal(t, []interface{}{"Ariel", "~", "~"}, args)
+}