implement CTEBuilder

Author: huandu

README.md

@@ -81,6 +81,7 @@ This package includes following pre-defined builders so far. API document and ex
 - [UpdateBuilder](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#UpdateBuilder): Builder for UPDATE.
 - [DeleteBuilder](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#DeleteBuilder): Builder for DELETE.
 - [UnionBuilder](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#UnionBuilder): Builder for UNION and UNION ALL.
+- [CTEBuilder](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#CTEBuilder): Builder for Common Table Expression (CTE), e.g. `WITH name (col1, col2) AS (SELECT ...)`.
 - [Buildf](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#Buildf): Freestyle builder using `fmt.Sprintf`-like syntax.
 - [Build](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#Build): Advanced freestyle builder using special syntax defined in [Args#Compile](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#Args.Compile).
 - [BuildNamed](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#BuildNamed): Advanced freestyle builder using `${key}` to refer the value of a map by key.

cte.go

@@ -0,0 +1,120 @@
+// Copyright 2024 Huan Du. All rights reserved.
+// Licensed under the MIT license that can be found in the LICENSE file.
+
+package sqlbuilder
+
+const (
+	cteMarkerInit injectionMarker = iota
+	cteMarkerAfterWith
+	cteMarkerAfterAs
+)
+
+// With creates a new CTE builder with default flavor.
+func With(name string, cols ...string) *CTEBuilder {
+	return DefaultFlavor.NewCTEBuilder().With(name, cols...)
+}
+
+func newCTEBuilder() *CTEBuilder {
+	return &CTEBuilder{
+		args:      &Args{},
+		injection: newInjection(),
+	}
+}
+
+// CTEBuilder is a CTE (Common Table Expression) builder.
+type CTEBuilder struct {
+	name       string
+	cols       []string
+	builderVar string
+
+	args *Args
+
+	injection *injection
+	marker    injectionMarker
+}
+
+var _ Builder = new(CTEBuilder)
+
+// With sets the CTE name and columns.
+func (cteb *CTEBuilder) With(name string, cols ...string) *CTEBuilder {
+	cteb.name = name
+	cteb.cols = cols
+	cteb.marker = cteMarkerAfterWith
+	return cteb
+}
+
+// As sets the builder to select data.
+func (cteb *CTEBuilder) As(builder Builder) *CTEBuilder {
+	cteb.builderVar = cteb.args.Add(builder)
+	cteb.marker = cteMarkerAfterAs
+	return cteb
+}
+
+// Select creates a new SelectBuilder to build a SELECT statement using this CTE.
+func (cteb *CTEBuilder) Select(col ...string) *SelectBuilder {
+	sb := cteb.args.Flavor.NewSelectBuilder()
+	return sb.With(cteb).Select(col...)
+}
+
+// String returns the compiled CTE string.
+func (cteb *CTEBuilder) String() string {
+	sql, _ := cteb.Build()
+	return sql
+}
+
+// Build returns compiled CTE string and args.
+func (cteb *CTEBuilder) Build() (sql string, args []interface{}) {
+	return cteb.BuildWithFlavor(cteb.args.Flavor)
+}
+
+// BuildWithFlavor builds a CTE with the specified flavor and initial arguments.
+func (cteb *CTEBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{}) (sql string, args []interface{}) {
+	buf := newStringBuilder()
+	cteb.injection.WriteTo(buf, cteMarkerInit)
+
+	if cteb.name != "" {
+		buf.WriteLeadingString("WITH ")
+		buf.WriteString(cteb.name)
+
+		if len(cteb.cols) > 0 {
+			buf.WriteLeadingString("(")
+			buf.WriteStrings(cteb.cols, ", ")
+			buf.WriteString(")")
+		}
+
+		cteb.injection.WriteTo(buf, cteMarkerAfterWith)
+	}
+
+	if cteb.builderVar != "" {
+		buf.WriteLeadingString("AS (")
+		buf.WriteString(cteb.builderVar)
+		buf.WriteRune(')')
+
+		cteb.injection.WriteTo(buf, cteMarkerAfterAs)
+	}
+
+	return cteb.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
+}
+
+// SetFlavor sets the flavor of compiled sql.
+func (cteb *CTEBuilder) SetFlavor(flavor Flavor) (old Flavor) {
+	old = cteb.args.Flavor
+	cteb.args.Flavor = flavor
+	return
+}
+
+// Var returns a placeholder for value.
+func (cteb *CTEBuilder) Var(arg interface{}) string {
+	return cteb.args.Add(arg)
+}
+
+// SQL adds an arbitrary sql to current position.
+func (cteb *CTEBuilder) SQL(sql string) *CTEBuilder {
+	cteb.injection.SQL(cteb.marker, sql)
+	return cteb
+}
+
+// TableName returns the CTE table name.
+func (cteb *CTEBuilder) TableName() string {
+	return cteb.name
+}

cte_test.go

@@ -0,0 +1,65 @@
+// Copyright 2024 Huan Du. All rights reserved.
+// Licensed under the MIT license that can be found in the LICENSE file.
+
+package sqlbuilder
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/huandu/go-assert"
+)
+
+func ExampleWith() {
+	sb := With("users", "id", "name").As(
+		Select("id", "name").From("users").Where("name IS NOT NULL"),
+	).Select("users.id", "orders.id").Join("orders", "users.id = orders.user_id")
+
+	fmt.Println(sb)
+
+	// Output:
+	// WITH users (id, name) AS (SELECT id, name FROM users WHERE name IS NOT NULL) SELECT users.id, orders.id FROM users JOIN orders ON users.id = orders.user_id
+}
+
+func ExampleCTEBuilder() {
+	usersBuilder := Select("id", "name", "level").From("users")
+	usersBuilder.Where(
+		usersBuilder.GreaterEqualThan("level", 10),
+	)
+	cteb := With("valid_users").As(usersBuilder)
+	fmt.Println(cteb)
+
+	sb := Select("valid_users.id", "valid_users.name", "orders.id").With(cteb)
+	sb.Join("orders", "valid_users.id = orders.user_id")
+	sb.Where(
+		sb.LessEqualThan("orders.price", 200),
+		"valid_users.level < orders.min_level",
+	).OrderBy("orders.price").Desc()
+
+	sql, args := sb.Build()
+	fmt.Println(sql)
+	fmt.Println(args)
+
+	// Output:
+	// WITH valid_users AS (SELECT id, name, level FROM users WHERE level >= ?)
+	// WITH valid_users AS (SELECT id, name, level FROM users WHERE level >= ?) SELECT valid_users.id, valid_users.name, orders.id FROM valid_users JOIN orders ON valid_users.id = orders.user_id WHERE orders.price <= ? AND valid_users.level < orders.min_level ORDER BY orders.price DESC
+	// [10 200]
+}
+
+func TestCTEBuilder(t *testing.T) {
+	a := assert.New(t)
+	cteb := newCTEBuilder()
+	cteb.SQL("/* init */")
+	cteb.With("t", "a", "b")
+	cteb.SQL("/* after with */")
+
+	// Make sure that calling Var() will not affect the As().
+	cteb.Var(123)
+
+	cteb.As(Select("a", "b").From("t"))
+	cteb.SQL("/* after as */")
+
+	sql, args := cteb.Build()
+	a.Equal(sql, "/* init */ WITH t (a, b) /* after with */ AS (SELECT a, b FROM t) /* after as */")
+	a.Assert(args == nil)
+}

flavor.go

@@ -141,6 +141,13 @@ func (f Flavor) NewUnionBuilder() *UnionBuilder {
 	return b
 }
 
+// NewCTEBuilder creates a new CTE builder with flavor.
+func (f Flavor) NewCTEBuilder() *CTEBuilder {
+	b := newCTEBuilder()
+	b.SetFlavor(f)
+	return b
+}
+
 // Quote adds quote for name to make sure the name can be used safely
 // as table name or field name.
 //

select.go

@@ -11,6 +11,7 @@ import (
 
 const (
 	selectMarkerInit injectionMarker = iota
+	selectMarkerAfterWith
 	selectMarkerAfterSelect
 	selectMarkerAfterFrom
 	selectMarkerAfterJoin
@@ -65,6 +66,7 @@ type SelectBuilder struct {
 	whereClauseProxy *whereClauseProxy
 	whereClauseExpr  string
 
+	cteBuilder  string
 	distinct    bool
 	tables      []string
 	selectCols  []string
@@ -92,6 +94,14 @@ func Select(col ...string) *SelectBuilder {
 	return DefaultFlavor.NewSelectBuilder().Select(col...)
 }
 
+// With sets WITH clause (the Common Table Expression) before SELECT.
+func (sb *SelectBuilder) With(builder *CTEBuilder) *SelectBuilder {
+	sb.marker = selectMarkerAfterWith
+	sb.cteBuilder = sb.Var(builder)
+	sb.tables = []string{builder.TableName()}
+	return sb
+}
+
 // Select sets columns in SELECT.
 func (sb *SelectBuilder) Select(col ...string) *SelectBuilder {
 	sb.selectCols = col
@@ -269,6 +279,11 @@ func (sb *SelectBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{
 
 	oraclePage := flavor == Oracle && (sb.limit >= 0 || sb.offset >= 0)
 
+	if sb.cteBuilder != "" {
+		buf.WriteLeadingString(sb.cteBuilder)
+		sb.injection.WriteTo(buf, selectMarkerAfterWith)
+	}
+
 	if len(sb.selectCols) > 0 {
 		buf.WriteLeadingString("SELECT ")
 

struct_test.go

@@ -638,9 +638,7 @@ func ExampleStruct_buildDELETE() {
 
 	// Prepare DELETE query.
 	user := &User{
-		ID:     1234,
-		Name:   "Huan Du",
-		Status: 1,
+		ID: 1234,
 	}
 	b := userStruct.DeleteFrom("user")
 	b.Where(b.Equal("id", user.ID))

union.go

@@ -37,7 +37,7 @@ func newUnionBuilder() *UnionBuilder {
 // UnionBuilder is a builder to build UNION.
 type UnionBuilder struct {
 	opt         string
-	builders    []Builder
+	builderVars []string
 	orderByCols []string
 	order       string
 	limit       int
@@ -72,8 +72,14 @@ func (ub *UnionBuilder) UnionAll(builders ...Builder) *UnionBuilder {
 }
 
 func (ub *UnionBuilder) union(opt string, builders ...Builder) *UnionBuilder {
+	builderVars := make([]string, 0, len(builders))
+
+	for _, b := range builders {
+		builderVars = append(builderVars, ub.Var(b))
+	}
+
 	ub.opt = opt
-	ub.builders = builders
+	ub.builderVars = builderVars
 	ub.marker = unionMarkerAfterUnion
 	return ub
 }
@@ -131,25 +137,25 @@ func (ub *UnionBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{}
 	buf := newStringBuilder()
 	ub.injection.WriteTo(buf, unionMarkerInit)
 
-	if len(ub.builders) > 0 {
+	if len(ub.builderVars) > 0 {
 		needParen := flavor != SQLite
 
 		if needParen {
 			buf.WriteLeadingString("(")
-			buf.WriteString(ub.Var(ub.builders[0]))
+			buf.WriteString(ub.builderVars[0])
 			buf.WriteRune(')')
 		} else {
-			buf.WriteLeadingString(ub.Var(ub.builders[0]))
+			buf.WriteLeadingString(ub.builderVars[0])
 		}
 
-		for _, b := range ub.builders[1:] {
+		for _, b := range ub.builderVars[1:] {
 			buf.WriteString(ub.opt)
 
 			if needParen {
 				buf.WriteRune('(')
 			}
 
-			buf.WriteString(ub.Var(b))
+			buf.WriteString(b)
 
 			if needParen {
 				buf.WriteRune(')')