Merge pull request #157 from huandu/feature/cte

Support CTE (Common Table Expression)

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,99 @@
+// 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
+)
+
+// With creates a new CTE builder with default flavor.
+func With(tables ...*CTETableBuilder) *CTEBuilder {
+	return DefaultFlavor.NewCTEBuilder().With(tables...)
+}
+
+func newCTEBuilder() *CTEBuilder {
+	return &CTEBuilder{
+		args:      &Args{},
+		injection: newInjection(),
+	}
+}
+
+// CTEBuilder is a CTE (Common Table Expression) builder.
+type CTEBuilder struct {
+	tableNames       []string
+	tableBuilderVars []string
+
+	args *Args
+
+	injection *injection
+	marker    injectionMarker
+}
+
+var _ Builder = new(CTEBuilder)
+
+// With sets the CTE name and columns.
+func (cteb *CTEBuilder) With(tables ...*CTETableBuilder) *CTEBuilder {
+	tableNames := make([]string, 0, len(tables))
+	tableBuilderVars := make([]string, 0, len(tables))
+
+	for _, table := range tables {
+		tableNames = append(tableNames, table.TableName())
+		tableBuilderVars = append(tableBuilderVars, cteb.args.Add(table))
+	}
+
+	cteb.tableNames = tableNames
+	cteb.tableBuilderVars = tableBuilderVars
+	cteb.marker = cteMarkerAfterWith
+	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 len(cteb.tableBuilderVars) > 0 {
+		buf.WriteLeadingString("WITH ")
+		buf.WriteStrings(cteb.tableBuilderVars, ", ")
+	}
+
+	cteb.injection.WriteTo(buf, cteMarkerAfterWith)
+	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
+}
+
+// SQL adds an arbitrary sql to current position.
+func (cteb *CTEBuilder) SQL(sql string) *CTEBuilder {
+	cteb.injection.SQL(cteb.marker, sql)
+	return cteb
+}
+
+// TableNames returns all table names in a CTE.
+func (cteb *CTEBuilder) TableNames() []string {
+	return cteb.tableNames
+}

cte_test.go

@@ -0,0 +1,81 @@
+// 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(
+		CTETable("users", "id", "name").As(
+			Select("id", "name").From("users").Where("name IS NOT NULL"),
+		),
+		CTETable("devices").As(
+			Select("device_id").From("devices"),
+		),
+	).Select("users.id", "orders.id", "devices.device_id").Join(
+		"orders",
+		"users.id = orders.user_id",
+		"devices.device_id = orders.device_id",
+	)
+
+	fmt.Println(sb)
+
+	// Output:
+	// WITH users (id, name) AS (SELECT id, name FROM users WHERE name IS NOT NULL), devices AS (SELECT device_id FROM devices) SELECT users.id, orders.id, devices.device_id FROM users, devices JOIN orders ON users.id = orders.user_id AND devices.device_id = orders.device_id
+}
+
+func ExampleCTEBuilder() {
+	usersBuilder := Select("id", "name", "level").From("users")
+	usersBuilder.Where(
+		usersBuilder.GreaterEqualThan("level", 10),
+	)
+	cteb := With(
+		CTETable("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()
+	ctetb := newCTETableBuilder()
+	cteb.SQL("/* init */")
+	cteb.With(ctetb)
+	cteb.SQL("/* after with */")
+
+	ctetb.SQL("/* table init */")
+	ctetb.Table("t", "a", "b")
+	ctetb.SQL("/* after table */")
+
+	ctetb.As(Select("a", "b").From("t"))
+	ctetb.SQL("/* after table as */")
+
+	sql, args := cteb.Build()
+	a.Equal(sql, "/* init */ WITH /* table init */ t (a, b) /* after table */ AS (SELECT a, b FROM t) /* after table as */ /* after with */")
+	a.Assert(args == nil)
+
+	sql = ctetb.String()
+	a.Equal(sql, "/* table init */ t (a, b) /* after table */ AS (SELECT a, b FROM t) /* after table as */")
+}

ctetable.go

@@ -0,0 +1,106 @@
+// Copyright 2024 Huan Du. All rights reserved.
+// Licensed under the MIT license that can be found in the LICENSE file.
+
+package sqlbuilder
+
+const (
+	cteTableMarkerInit injectionMarker = iota
+	cteTableMarkerAfterTable
+	cteTableMarkerAfterAs
+)
+
+// CTETable creates a new CTE table builder with default flavor.
+func CTETable(name string, cols ...string) *CTETableBuilder {
+	return DefaultFlavor.NewCTETableBuilder().Table(name, cols...)
+}
+
+func newCTETableBuilder() *CTETableBuilder {
+	return &CTETableBuilder{
+		args:      &Args{},
+		injection: newInjection(),
+	}
+}
+
+// CTETableBuilder is a builder to build one table in CTE (Common Table Expression).
+type CTETableBuilder struct {
+	name       string
+	cols       []string
+	builderVar string
+
+	args *Args
+
+	injection *injection
+	marker    injectionMarker
+}
+
+// Table sets the table name and columns in a CTE table.
+func (ctetb *CTETableBuilder) Table(name string, cols ...string) *CTETableBuilder {
+	ctetb.name = name
+	ctetb.cols = cols
+	ctetb.marker = cteTableMarkerAfterTable
+	return ctetb
+}
+
+// As sets the builder to select data.
+func (ctetb *CTETableBuilder) As(builder Builder) *CTETableBuilder {
+	ctetb.builderVar = ctetb.args.Add(builder)
+	ctetb.marker = cteTableMarkerAfterAs
+	return ctetb
+}
+
+// String returns the compiled CTE string.
+func (ctetb *CTETableBuilder) String() string {
+	sql, _ := ctetb.Build()
+	return sql
+}
+
+// Build returns compiled CTE string and args.
+func (ctetb *CTETableBuilder) Build() (sql string, args []interface{}) {
+	return ctetb.BuildWithFlavor(ctetb.args.Flavor)
+}
+
+// BuildWithFlavor builds a CTE with the specified flavor and initial arguments.
+func (ctetb *CTETableBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{}) (sql string, args []interface{}) {
+	buf := newStringBuilder()
+	ctetb.injection.WriteTo(buf, cteTableMarkerInit)
+
+	if ctetb.name != "" {
+		buf.WriteLeadingString(ctetb.name)
+
+		if len(ctetb.cols) > 0 {
+			buf.WriteLeadingString("(")
+			buf.WriteStrings(ctetb.cols, ", ")
+			buf.WriteString(")")
+		}
+
+		ctetb.injection.WriteTo(buf, cteTableMarkerAfterTable)
+	}
+
+	if ctetb.builderVar != "" {
+		buf.WriteLeadingString("AS (")
+		buf.WriteString(ctetb.builderVar)
+		buf.WriteRune(')')
+
+		ctetb.injection.WriteTo(buf, cteTableMarkerAfterAs)
+	}
+
+	return ctetb.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
+}
+
+// SetFlavor sets the flavor of compiled sql.
+func (ctetb *CTETableBuilder) SetFlavor(flavor Flavor) (old Flavor) {
+	old = ctetb.args.Flavor
+	ctetb.args.Flavor = flavor
+	return
+}
+
+// SQL adds an arbitrary sql to current position.
+func (ctetb *CTETableBuilder) SQL(sql string) *CTETableBuilder {
+	ctetb.injection.SQL(ctetb.marker, sql)
+	return ctetb
+}
+
+// TableName returns the CTE table name.
+func (ctetb *CTETableBuilder) TableName() string {
+	return ctetb.name
+}

flavor.go

@@ -141,6 +141,20 @@ 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
+}
+
+// NewCTETableBuilder creates a new CTE table builder with flavor.
+func (f Flavor) NewCTETableBuilder() *CTETableBuilder {
+	b := newCTETableBuilder()
+	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 = builder.TableNames()
+	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))