fix #211: Support clone (deep copy) in all builders

Author: huandu

args.go

@@ -6,9 +6,12 @@ package sqlbuilder
 import (
 	"database/sql"
 	"fmt"
+	"reflect"
 	"sort"
 	"strconv"
 	"strings"
+
+	"github.com/huandu/go-clone"
 )
 
 // Args stores arguments associated with a SQL.
@@ -17,7 +20,7 @@ type Args struct {
 	Flavor Flavor
 
 	indexBase    int
-	argValues    []interface{}
+	argValues    *valueStore
 	namedArgs    map[string]int
 	sqlNamedArgs map[string]int
 	onlyNamed    bool
@@ -48,7 +51,7 @@ func (args *Args) Add(arg interface{}) string {
 }
 
 func (args *Args) add(arg interface{}) int {
-	idx := len(args.argValues) + args.indexBase
+	idx := args.argValues.Len() + args.indexBase
 
 	switch a := arg.(type) {
 	case sql.NamedArg:
@@ -57,7 +60,7 @@ func (args *Args) add(arg interface{}) int {
 		}
 
 		if p, ok := args.sqlNamedArgs[a.Name]; ok {
-			arg = args.argValues[p]
+			arg = args.argValues.Load(p)
 			break
 		}
 
@@ -68,7 +71,7 @@ func (args *Args) add(arg interface{}) int {
 		}
 
 		if p, ok := args.namedArgs[a.name]; ok {
-			arg = args.argValues[p]
+			arg = args.argValues.Load(p)
 			break
 		}
 
@@ -78,10 +81,31 @@ func (args *Args) add(arg interface{}) int {
 		return idx
 	}
 
-	args.argValues = append(args.argValues, arg)
+	if args.argValues == nil {
+		args.argValues = &valueStore{}
+	}
+
+	args.argValues.Add(arg)
 	return idx
 }
 
+// Replace replaces the placeholder with arg.
+//
+// The placeholder must be the value returned by `Add`, e.g. "$1".
+// If the placeholder is not found, this method does nothing.
+func (args *Args) Replace(placeholder string, arg interface{}) {
+	dollar := strings.IndexRune(placeholder, '$')
+
+	if dollar != 0 {
+		return
+	}
+
+	if i, err := strconv.Atoi(placeholder[1:]); err == nil {
+		i -= args.indexBase
+		args.argValues.Set(i, arg)
+	}
+}
+
 // Compile compiles builder's format to standard sql and returns associated args.
 //
 // The format string uses a special syntax to represent arguments.
@@ -201,14 +225,14 @@ func (args *Args) compileDigits(ctx *argsCompileContext, format string, offset i
 }
 
 func (args *Args) compileSuccessive(ctx *argsCompileContext, format string, offset int) (string, int) {
-	if offset < 0 || offset >= len(args.argValues) {
+	if offset < 0 || offset >= args.argValues.Len() {
 		ctx.WriteString("/* INVALID ARG $")
 		ctx.WriteString(strconv.Itoa(offset))
 		ctx.WriteString(" */")
 		return format, offset
 	}
 
-	arg := args.argValues[offset]
+	arg := args.argValues.Load(offset)
 	ctx.WriteValue(arg)
 
 	return format, offset + 1
@@ -245,7 +269,7 @@ func (args *Args) mergeSQLNamedArgs(ctx *argsCompileContext) []interface{} {
 	sort.Ints(ints)
 
 	for _, i := range ints {
-		values = append(values, args.argValues[i])
+		values = append(values, args.argValues.Load(i))
 	}
 
 	return values
@@ -364,3 +388,50 @@ func (ctx *argsCompileContext) WriteValues(values []interface{}, sep string) {
 		ctx.WriteValue(v)
 	}
 }
+
+type valueStore struct {
+	Values []interface{}
+}
+
+func init() {
+	// The values in valueStore should be shadow-copied to avoid unnecessary cost.
+	t := reflect.TypeOf(valueStore{})
+	clone.SetCustomFunc(t, func(allocator *clone.Allocator, old, new reflect.Value) {
+		values := old.FieldByName("Values")
+		newValues := allocator.Clone(values)
+		new.FieldByName("Values").Set(newValues)
+	})
+}
+
+func (as *valueStore) Len() int {
+	if as == nil {
+		return 0
+	}
+
+	return len(as.Values)
+}
+
+// Add adds an arg to argsValues and returns its index.
+func (as *valueStore) Add(arg interface{}) int {
+	as.Values = append(as.Values, arg)
+	return len(as.Values) - 1
+}
+
+// Set sets the arg value by index.
+func (as *valueStore) Set(index int, arg interface{}) {
+	if as == nil || index < 0 || index >= len(as.Values) {
+		return
+	}
+
+	as.Values[index] = arg
+}
+
+// Load returns the arg value by index.
+// Returns nil if index is out of range or as itself is nil.
+func (as *valueStore) Load(index int) interface{} {
+	if as == nil || index < 0 || index >= len(as.Values) {
+		return nil
+	}
+
+	return as.Values[index]
+}

createtable.go

@@ -5,6 +5,8 @@ package sqlbuilder
 
 import (
 	"strings"
+
+	"github.com/huandu/go-clone"
 )
 
 const (
@@ -29,6 +31,12 @@ func newCreateTableBuilder() *CreateTableBuilder {
 	}
 }
 
+// Clone returns a deep copy of CreateTableBuilder.
+// It's useful when you want to create a base builder and clone it to build similar queries.
+func (ctb *CreateTableBuilder) Clone() *CreateTableBuilder {
+	return clone.Clone(ctb).(*CreateTableBuilder)
+}
+
 // CreateTableBuilder is a builder to build CREATE TABLE.
 type CreateTableBuilder struct {
 	verb        string

createtable_test.go

@@ -102,3 +102,22 @@ func TestCreateTableGetFlavor(t *testing.T) {
 	flavor = ctbClick.Flavor()
 	a.Equal(ClickHouse, flavor)
 }
+
+func TestCreateTableClone(t *testing.T) {
+	a := assert.New(t)
+	ctb := CreateTable("demo.user").IfNotExists().
+		Define("id", "BIGINT(20)", "NOT NULL", "AUTO_INCREMENT", "PRIMARY KEY", `COMMENT "user id"`).
+		Option("DEFAULT CHARACTER SET", "utf8mb4")
+
+	ctb2 := ctb.Clone()
+	ctb2.Define("name", "VARCHAR(255)", "NOT NULL", `COMMENT "user name"`)
+
+	sql1, args1 := ctb.Build()
+	sql2, args2 := ctb2.Build()
+
+	a.Equal("CREATE TABLE IF NOT EXISTS demo.user (id BIGINT(20) NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT \"user id\") DEFAULT CHARACTER SET utf8mb4", sql1)
+	a.Equal(0, len(args1))
+
+	a.Equal("CREATE TABLE IF NOT EXISTS demo.user (id BIGINT(20) NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT \"user id\", name VARCHAR(255) NOT NULL COMMENT \"user name\") DEFAULT CHARACTER SET utf8mb4", sql2)
+	a.Equal(0, len(args2))
+}

cte.go

@@ -3,6 +3,12 @@
 
 package sqlbuilder
 
+import (
+	"reflect"
+
+	"github.com/huandu/go-clone"
+)
+
 const (
 	cteMarkerInit injectionMarker = iota
 	cteMarkerAfterWith
@@ -25,6 +31,25 @@ func newCTEBuilder() *CTEBuilder {
 	}
 }
 
+// Clone returns a deep copy of CTEBuilder.
+// It's useful when you want to create a base builder and clone it to build similar queries.
+func (cteb *CTEBuilder) Clone() *CTEBuilder {
+	return clone.Clone(cteb).(*CTEBuilder)
+}
+
+func init() {
+	t := reflect.TypeOf(CTEBuilder{})
+	clone.SetCustomFunc(t, func(allocator *clone.Allocator, old, new reflect.Value) {
+		cloned := allocator.CloneSlowly(old)
+		new.Set(cloned)
+
+		cteb := cloned.Addr().Interface().(*CTEBuilder)
+		for i, b := range cteb.queries {
+			cteb.args.Replace(cteb.queryBuilderVars[i], b)
+		}
+	})
+}
+
 // CTEBuilder is a CTE (Common Table Expression) builder.
 type CTEBuilder struct {
 	recursive        bool
@@ -47,7 +72,7 @@ func (cteb *CTEBuilder) With(queries ...*CTEQueryBuilder) *CTEBuilder {
 		queryBuilderVars = append(queryBuilderVars, cteb.args.Add(query))
 	}
 
-	cteb.queries = queries
+	cteb.queries = append([]*CTEQueryBuilder(nil), queries...)
 	cteb.queryBuilderVars = queryBuilderVars
 	cteb.marker = cteMarkerAfterWith
 	return cteb

cte_test.go

@@ -193,3 +193,53 @@ func TestCTEQueryBuilderGetFlavor(t *testing.T) {
 	flavor = ctetbClick.Flavor()
 	a.Equal(ClickHouse, flavor)
 }
+
+func TestCTEBuildClone(t *testing.T) {
+	a := assert.New(t)
+	cteb := With(
+		CTETable("users", "id", "name").As(
+			Select("id", "name").From("users").Where("name IS NOT NULL"),
+		),
+	)
+	ctebClone := cteb.Clone()
+	a.Equal(cteb.String(), ctebClone.String())
+
+	sb := Select("users.id", "users.name", "orders.id").From("users").With(cteb)
+	sbClone := sb.Clone()
+	a.Equal(sb.String(), sbClone.String())
+
+	sql, args := sb.Build()
+	sqlClone, argsClone := sbClone.Build()
+	a.Equal(sql, sqlClone)
+	a.Equal(args, argsClone)
+
+	sb.Limit(20)
+	a.NotEqual(sb.String(), sbClone.String())
+}
+
+func TestCTEQueryBuilderClone(t *testing.T) {
+	a := assert.New(t)
+	ctetb := CTETable("users", "id", "name").As(Select("id", "name").From("users").Where("id > 0"))
+	// Ensure AddToTableList flag is respected in clone as default for CTETable
+	clone := ctetb.Clone()
+	a.Equal(ctetb.String(), clone.String())
+	a.Equal(ctetb.ShouldAddToTableList(), clone.ShouldAddToTableList())
+}
+
+func TestCTEBuilderClone(t *testing.T) {
+	a := assert.New(t)
+	q1 := CTETable("u", "id").As(Select("id").From("users"))
+	q2 := CTEQuery("o", "id").As(Select("id").From("orders"))
+	cte := With(q1, q2)
+	clone := cte.Clone()
+
+	s1, args1 := cte.Build()
+	s2, args2 := clone.Build()
+	a.Equal(s1, s2)
+	a.Equal(args1, args2)
+
+	// mutate clone and verify original unchanged
+	q3 := CTETable("p", "id").As(Select("id").From("profiles"))
+	clone.With(q3)
+	a.NotEqual(cte.String(), clone.String())
+}

ctequery.go

@@ -3,6 +3,12 @@
 
 package sqlbuilder
 
+import (
+	"reflect"
+
+	"github.com/huandu/go-clone"
+)
+
 const (
 	cteQueryMarkerInit injectionMarker = iota
 	cteQueryMarkerAfterTable
@@ -29,10 +35,28 @@ func newCTEQueryBuilder() *CTEQueryBuilder {
 	}
 }
 
+// Clone returns a deep copy of CTEQueryBuilder.
+// It's useful when you want to create a base builder and clone it to build similar queries.
+func (ctetb *CTEQueryBuilder) Clone() *CTEQueryBuilder {
+	return clone.Clone(ctetb).(*CTEQueryBuilder)
+}
+
+func init() {
+	t := reflect.TypeOf(CTEQueryBuilder{})
+	clone.SetCustomFunc(t, func(allocator *clone.Allocator, old, new reflect.Value) {
+		cloned := allocator.CloneSlowly(old)
+		new.Set(cloned)
+
+		ctetb := cloned.Addr().Interface().(*CTEQueryBuilder)
+		ctetb.args.Replace(ctetb.builderVar, ctetb.builder)
+	})
+}
+
 // CTEQueryBuilder is a builder to build one table in CTE (Common Table Expression).
 type CTEQueryBuilder struct {
 	name       string
 	cols       []string
+	builder    Builder
 	builderVar string
 
 	// if true, this query's table name will be automatically added to the table list
@@ -62,6 +86,7 @@ func (ctetb *CTEQueryBuilder) Table(name string, cols ...string) *CTEQueryBuilde
 
 // As sets the builder to select data.
 func (ctetb *CTEQueryBuilder) As(builder Builder) *CTEQueryBuilder {
+	ctetb.builder = builder
 	ctetb.builderVar = ctetb.args.Add(builder)
 	ctetb.marker = cteQueryMarkerAfterAs
 	return ctetb

delete.go

@@ -3,6 +3,12 @@
 
 package sqlbuilder
 
+import (
+	"reflect"
+
+	"github.com/huandu/go-clone"
+)
+
 const (
 	deleteMarkerInit injectionMarker = iota
 	deleteMarkerAfterWith
@@ -33,6 +39,24 @@ func newDeleteBuilder() *DeleteBuilder {
 	}
 }
 
+// Clone returns a deep copy of DeleteBuilder.
+// It's useful when you want to create a base builder and clone it to build similar queries.
+func (db *DeleteBuilder) Clone() *DeleteBuilder {
+	return clone.Clone(db).(*DeleteBuilder)
+}
+
+func init() {
+	t := reflect.TypeOf(DeleteBuilder{})
+	clone.SetCustomFunc(t, func(allocator *clone.Allocator, old, new reflect.Value) {
+		cloned := allocator.CloneSlowly(old)
+		new.Set(cloned)
+
+		db := cloned.Addr().Interface().(*DeleteBuilder)
+		db.args.Replace(db.whereClauseExpr, db.whereClauseProxy)
+		db.args.Replace(db.cteBuilderVar, db.cteBuilder)
+	})
+}
+
 // DeleteBuilder is a builder to build DELETE.
 type DeleteBuilder struct {
 	*WhereClause