support with/from/returning and table-like targets in update

Author: deoxxa

sqlbuilder_test.go

@@ -57,6 +57,33 @@ func TestUpdate(t *testing.T) {
 	a.Equal([]interface{}{"jim", 5}, qv)
 }
 
+func TestUpdateWithFromReturning(t *testing.T) {
+	a := assert.New(t)
+
+	tbl := NewTable("jobs", "id", "data", "reserved_to", "completed_at")
+
+	nextJob := CommonTableExpression("next_job").As(
+		Select().From(tbl).Columns(tbl.C("id")).Where(BooleanOperator("and",
+			BooleanOperator("or", IsNull(tbl.C("reserved_to")), Lte(tbl.C("reserved_to"), Func("now"))),
+			IsNull(tbl.C("completed_at")),
+		)).OrderBy(OrderAsc(tbl.C("id"))).OffsetLimit(OffsetLimit(nil, Literal("1"))),
+	)
+
+	updatedJob := AliasTable(tbl, "updated_job")
+
+	expiry := time.Date(2023, time.July, 25, 18, 34, 15, 0, time.UTC)
+
+	q := Update().With(nextJob).Target(updatedJob).Set(UpdateColumns{
+		updatedJob.C("reserved_to"): Bind(expiry),
+	}).Where(Eq(updatedJob.C("id"), nextJob.C("id"))).Returning(updatedJob.C("id"))
+
+	qs, qv, err := NewSerializer(DialectPostgres{}).F(q.AsStatement).ToSQL()
+
+	a.NoError(err)
+	a.Equal(`WITH "next_job" AS (SELECT "jobs"."id" FROM "jobs" WHERE (("jobs"."reserved_to" IS NULL or ("jobs"."reserved_to" <= now())) and "jobs"."completed_at" IS NULL) ORDER BY "jobs"."id" ASC LIMIT 1) UPDATE "jobs" "updated_job" SET "reserved_to" = $1 WHERE ("updated_job"."id" = "next_job"."id") RETURNING "updated_job"."id"`, qs)
+	a.Equal([]interface{}{expiry}, qv)
+}
+
 func TestSelect(t *testing.T) {
 	a := assert.New(t)
 

update.go

@@ -3,29 +3,55 @@ package sqlbuilder
 type UpdateColumns map[*BasicColumn]AsExpr
 
 type UpdateStatement struct {
-	table *Table
-	where AsExpr
-	set   UpdateColumns
+	with      []AsCommonTableExpression
+	target    AsTableOrSubquery
+	set       UpdateColumns
+	from      AsTableOrSubquery
+	where     AsExpr
+	returning []AsExpr
 }
 
 func (s *UpdateStatement) clone() *UpdateStatement {
 	return &UpdateStatement{
-		table: s.table,
-		where: s.where,
-		set:   s.set,
+		with:      s.with,
+		target:    s.target,
+		set:       s.set,
+		from:      s.from,
+		where:     s.where,
+		returning: s.returning,
 	}
 }
 
 func Update() *UpdateStatement {
 	return &UpdateStatement{set: make(UpdateColumns)}
 }
 
+func (s *UpdateStatement) With(with ...AsCommonTableExpression) *UpdateStatement {
+	c := s.clone()
+	c.with = with
+	return c
+}
+
+func (s *UpdateStatement) AndWith(with ...AsCommonTableExpression) *UpdateStatement {
+	c := s.clone()
+	c.with = append(c.with, with...)
+	return c
+}
+
 func (s *UpdateStatement) Table(table *Table) *UpdateStatement {
+	return s.Target(table)
+}
+
+func (s *UpdateStatement) Target(target AsTableOrSubquery) *UpdateStatement {
 	c := s.clone()
-	c.table = table
+	c.target = target
 	return c
 }
 
+func (s *UpdateStatement) GetTarget() AsTableOrSubquery {
+	return s.target
+}
+
 func (s *UpdateStatement) Set(set UpdateColumns) *UpdateStatement {
 	c := s.clone()
 	c.set = set
@@ -47,22 +73,67 @@ func (s *UpdateStatement) AndSet(column *BasicColumn, expr AsExpr) *UpdateStatem
 	return c
 }
 
+func (s *UpdateStatement) From(from AsTableOrSubquery) *UpdateStatement {
+	c := s.clone()
+	c.from = from
+	return c
+}
+
 func (s *UpdateStatement) Where(where AsExpr) *UpdateStatement {
 	c := s.clone()
 	c.where = where
 	return c
 }
 
+func (s *UpdateStatement) Returning(returning ...AsExpr) *UpdateStatement {
+	c := s.clone()
+	c.returning = returning
+	return c
+}
+
 func (q *UpdateStatement) AsStatement(s *Serializer) {
-	s.D("UPDATE ").F(q.table.AsNamed).D(" SET ")
+	if len(q.with) > 0 {
+		s.D("WITH ")
+
+		for _, w := range q.with {
+			if w.IsRecursive() {
+				s.D("RECURSIVE ")
+				break
+			}
+		}
+
+		for i, w := range q.with {
+			s.F(w.AsCommonTableExpression).DC(",", i != len(q.with)-1).D(" ")
+		}
+	}
+
+	s.D("UPDATE ").F(q.target.AsTableOrSubquery).D(" SET ")
 
 	i := 0
 	for k, e := range q.set {
 		s.F(k.AsNamedShort).D(" = ").F(e.AsExpr).DC(", ", i < len(q.set)-1)
 		i++
 	}
 
+	if q.from != nil {
+		s.D(" FROM ").F(q.from.AsTableOrSubquery)
+	}
+
 	if q.where != nil {
 		s.D(" WHERE ").F(q.where.AsExpr)
 	}
+
+	if len(q.returning) > 0 {
+		s.D(" RETURNING ")
+
+		for i, c := range q.returning {
+			if a, ok := c.(AsResultColumn); ok {
+				s.F(a.AsResultColumn)
+			} else {
+				s.F(c.AsExpr)
+			}
+
+			s.DC(", ", i < len(q.returning)-1)
+		}
+	}
 }