qb: update builder advanced assignments

qb/cmp.go

@@ -53,15 +53,15 @@ func (c Cmp) writeCql(cql *bytes.Buffer) (names []string) {
 		cql.WriteString(" CONTAINS ")
 	}
 
-	if c.fn == nil {
+	if c.fn != nil {
+		names = append(names, c.fn.writeCql(cql)...)
+	} else {
 		cql.WriteByte('?')
 		if c.name == "" {
 			names = append(names, c.column)
 		} else {
 			names = append(names, c.name)
 		}
-	} else {
-		names = append(names, c.fn.writeCql(cql)...)
 	}
 
 	return
@@ -188,7 +188,7 @@ func GtOrEqNamed(column, name string) Cmp {
 	}
 }
 
-// GtFunc produces column>=someFunc(?...).
+// GtOrEqFunc produces column>=someFunc(?...).
 func GtOrEqFunc(column string, fn *Func) Cmp {
 	return Cmp{
 		op:     geq,

qb/update.go

@@ -9,16 +9,41 @@ package qb
 
 import (
 	"bytes"
+	"fmt"
 )
 
+// assignment specifies an assignment in a set operation.
+type assignment struct {
+	column string
+	name   string
+	expr   bool
+	fn     *Func
+}
+
+func (a assignment) writeCql(cql *bytes.Buffer) (names []string) {
+	cql.WriteString(a.column)
+	switch {
+	case a.expr:
+		names = append(names, a.name)
+	case a.fn != nil:
+		cql.WriteByte('=')
+		names = append(names, a.fn.writeCql(cql)...)
+	default:
+		cql.WriteByte('=')
+		cql.WriteByte('?')
+		names = append(names, a.name)
+	}
+	return
+}
+
 // UpdateBuilder builds CQL UPDATE statements.
 type UpdateBuilder struct {
-	table   string
-	using   using
-	columns columns
-	where   where
-	_if     _if
-	exists  bool
+	table       string
+	using       using
+	assignments []assignment
+	where       where
+	_if         _if
+	exists      bool
 }
 
 // Update returns a new UpdateBuilder with the given table name.
@@ -39,14 +64,12 @@ func (b *UpdateBuilder) ToCql() (stmt string, names []string) {
 	names = append(names, b.using.writeCql(&cql)...)
 
 	cql.WriteString("SET ")
-	for i, c := range b.columns {
-		cql.WriteString(c)
-		cql.WriteString("=?")
-		if i < len(b.columns)-1 {
+	for i, a := range b.assignments {
+		names = append(names, a.writeCql(&cql)...)
+		if i < len(b.assignments)-1 {
 			cql.WriteByte(',')
 		}
 	}
-	names = append(names, b.columns...)
 	cql.WriteByte(' ')
 
 	names = append(names, b.where.writeCql(&cql)...)
@@ -80,7 +103,51 @@ func (b *UpdateBuilder) TTL() *UpdateBuilder {
 
 // Set adds SET clauses to the query.
 func (b *UpdateBuilder) Set(columns ...string) *UpdateBuilder {
-	b.columns = append(b.columns, columns...)
+	for _, c := range columns {
+		b.assignments = append(b.assignments, assignment{
+			column: c,
+			name:   c,
+		})
+	}
+
+	return b
+}
+
+// SetFunc adds SET column=someFunc(?...) clause to the query.
+func (b *UpdateBuilder) SetFunc(column string, fn *Func) *UpdateBuilder {
+	b.assignments = append(b.assignments, assignment{column: column, fn: fn})
+	return b
+}
+
+// Add adds SET column=column+? clauses to the query.
+func (b *UpdateBuilder) Add(column string) *UpdateBuilder {
+	return b.AddNamed(column, column)
+}
+
+// AddNamed adds SET column=column+? clauses to the query with a custom
+// parameter name.
+func (b *UpdateBuilder) AddNamed(column, name string) *UpdateBuilder {
+	b.assignments = append(b.assignments, assignment{
+		column: fmt.Sprint(column, "=", column, "+?"),
+		name:   name,
+		expr:   true,
+	})
+	return b
+}
+
+// Remove adds SET column=column-? clauses to the query.
+func (b *UpdateBuilder) Remove(column string) *UpdateBuilder {
+	return b.RemoveNamed(column, column)
+}
+
+// RemoveNamed adds SET column=column-? clauses to the query with a custom
+// parameter name.
+func (b *UpdateBuilder) RemoveNamed(column, name string) *UpdateBuilder {
+	b.assignments = append(b.assignments, assignment{
+		column: fmt.Sprint(column, "=", column, "-?"),
+		name:   name,
+		expr:   true,
+	})
 	return b
 }
 

qb/update_test.go

@@ -36,6 +36,36 @@ func TestUpdateBuilder(t *testing.T) {
 			S: "UPDATE cycling.cyclist_name SET id=?,user_uuid=?,firstname=?,stars=? WHERE id=? ",
 			N: []string{"id", "user_uuid", "firstname", "stars", "expr"},
 		},
+		// Add SET SetFunc
+		{
+			B: Update("cycling.cyclist_name").SetFunc("user_uuid", Fn("someFunc", "param_0", "param_1")).Where(w).Set("stars"),
+			S: "UPDATE cycling.cyclist_name SET user_uuid=someFunc(?,?),stars=? WHERE id=? ",
+			N: []string{"param_0", "param_1", "stars", "expr"},
+		},
+		// Add SET Add
+		{
+			B: Update("cycling.cyclist_name").Add("total").Where(w),
+			S: "UPDATE cycling.cyclist_name SET total=total+? WHERE id=? ",
+			N: []string{"total", "expr"},
+		},
+		// Add SET AddNamed
+		{
+			B: Update("cycling.cyclist_name").AddNamed("total", "inc").Where(w),
+			S: "UPDATE cycling.cyclist_name SET total=total+? WHERE id=? ",
+			N: []string{"inc", "expr"},
+		},
+		// Add SET Remove
+		{
+			B: Update("cycling.cyclist_name").Remove("total").Where(w),
+			S: "UPDATE cycling.cyclist_name SET total=total-? WHERE id=? ",
+			N: []string{"total", "expr"},
+		},
+		// Add SET RemoveNamed
+		{
+			B: Update("cycling.cyclist_name").RemoveNamed("total", "dec").Where(w),
+			S: "UPDATE cycling.cyclist_name SET total=total-? WHERE id=? ",
+			N: []string{"dec", "expr"},
+		},
 		// Add WHERE
 		{
 			B: Update("cycling.cyclist_name").Set("id", "user_uuid", "firstname").Where(w, Gt("firstname")),