Merge pull request #58 from huandu/feature-sql-injection

Fix #48 #49 #53: Add new method SQL in all builders

README.md

@@ -19,7 +19,20 @@ go get github.com/huandu/go-sqlbuilder
 
 ### Basic usage
 
-Here is a sample to demonstrate how to build a SELECT query.
+We can build a SQL really quick with this package.
+
+```go
+sql := sqlbuilder.Select("id", "name").From("demo.user").
+    Where("status = 1").Limit(10).
+    String()
+
+fmt.Println(sql)
+
+// Output:
+// SELECT id, name FROM demo.user WHERE status = 1 LIMIT 10
+```
+
+In the most cases, we need to escape all input from user. In this case, create a builder before starting.
 
 ```go
 sb := sqlbuilder.NewSelectBuilder()
@@ -37,6 +50,8 @@ fmt.Println(args)
 // [1 2 5]
 ```
 
+### Pre-defined SQL builders
+
 Following builders are implemented right now. API document and examples are provided in the `godoc` document.
 
 - [Struct](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#Struct): Builder factory for a struct.
@@ -50,6 +65,29 @@ Following builders are implemented right now. API document and examples are prov
 - [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.
 
+There is a method `SQL(sql string)` implemented by all statement builders like `SelectBuilder`. We can use this method to insert any arbitrary SQL fragment when building a SQL. It's quite useful to build SQL containing non-standard syntax supported by a OLTP or OLAP system.
+
+```go
+// Build a SQL to create a HIVE table.
+sql := sqlbuilder.CreateTable("users").
+    SQL("PARTITION BY (year)").
+    SQL("AS").
+    SQL(
+        sqlbuilder.Select("columns[0] id", "columns[1] name", "columns[2] year").
+            From("`all-users.csv`").
+            Limit(100).
+            String(),
+    ).
+    String()
+
+fmt.Println(sql)
+
+// Output:
+// CREATE TABLE users PARTITION BY (year) AS SELECT columns[0] id, columns[1] name, columns[2] year FROM `all-users.csv` LIMIT 100
+```
+
+To learn how to use builders, check out [examples on GoDoc](https://pkg.go.dev/github.com/huandu/go-sqlbuilder#pkg-examples).
+
 ### Build SQL for MySQL, PostgreSQL or SQLite
 
 Parameter markers are different in MySQL, PostgreSQL and SQLite. This package provides some methods to set the type of markers (we call it "flavor") in all builders.

createtable.go

@@ -8,6 +8,13 @@ import (
 	"strings"
 )
 
+const (
+	createTableMarkerInit injectionMarker = iota
+	createTableMarkerAfterCreate
+	createTableMarkerAfterDefine
+	createTableMarkerAfterOption
+)
+
 // NewCreateTableBuilder creates a new CREATE TABLE builder.
 func NewCreateTableBuilder() *CreateTableBuilder {
 	return DefaultFlavor.NewCreateTableBuilder()
@@ -16,8 +23,10 @@ func NewCreateTableBuilder() *CreateTableBuilder {
 func newCreateTableBuilder() *CreateTableBuilder {
 	args := &Args{}
 	return &CreateTableBuilder{
-		verb: "CREATE TABLE",
-		args: args,
+		verb:      "CREATE TABLE",
+		args:      args,
+		injection: newInjection(),
+		marker:    createTableMarkerInit,
 	}
 }
 
@@ -30,18 +39,30 @@ type CreateTableBuilder struct {
 	options     [][]string
 
 	args *Args
+
+	injection *injection
+	marker    injectionMarker
+}
+
+var _ Builder = new(CreateTableBuilder)
+
+// CreateTable sets the table name in CREATE TABLE.
+func CreateTable(table string) *CreateTableBuilder {
+	return DefaultFlavor.NewCreateTableBuilder().CreateTable(table)
 }
 
 // CreateTable sets the table name in CREATE TABLE.
 func (ctb *CreateTableBuilder) CreateTable(table string) *CreateTableBuilder {
 	ctb.table = Escape(table)
+	ctb.marker = createTableMarkerAfterCreate
 	return ctb
 }
 
 // CreateTempTable sets the table name and changes the verb of ctb to CREATE TEMPORARY TABLE.
 func (ctb *CreateTableBuilder) CreateTempTable(table string) *CreateTableBuilder {
 	ctb.verb = "CREATE TEMPORARY TABLE"
 	ctb.table = Escape(table)
+	ctb.marker = createTableMarkerAfterCreate
 	return ctb
 }
 
@@ -54,12 +75,14 @@ func (ctb *CreateTableBuilder) IfNotExists() *CreateTableBuilder {
 // Define adds definition of a column or index in CREATE TABLE.
 func (ctb *CreateTableBuilder) Define(def ...string) *CreateTableBuilder {
 	ctb.defs = append(ctb.defs, def)
+	ctb.marker = createTableMarkerAfterDefine
 	return ctb
 }
 
 // Option adds a table option in CREATE TABLE.
 func (ctb *CreateTableBuilder) Option(opt ...string) *CreateTableBuilder {
 	ctb.options = append(ctb.options, opt)
+	ctb.marker = createTableMarkerAfterOption
 	return ctb
 }
 
@@ -79,6 +102,7 @@ func (ctb *CreateTableBuilder) Build() (sql string, args []interface{}) {
 // They can be used in `DB#Query` of package `database/sql` directly.
 func (ctb *CreateTableBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{}) (sql string, args []interface{}) {
 	buf := &bytes.Buffer{}
+	ctb.injection.WriteTo(buf, createTableMarkerInit)
 	buf.WriteString(ctb.verb)
 
 	if ctb.ifNotExists {
@@ -87,16 +111,22 @@ func (ctb *CreateTableBuilder) BuildWithFlavor(flavor Flavor, initialArg ...inte
 
 	buf.WriteRune(' ')
 	buf.WriteString(ctb.table)
-	buf.WriteString(" (")
+	ctb.injection.WriteTo(buf, createTableMarkerAfterCreate)
 
-	defs := make([]string, 0, len(ctb.defs))
+	if len(ctb.defs) > 0 {
+		buf.WriteString(" (")
 
-	for _, def := range ctb.defs {
-		defs = append(defs, strings.Join(def, " "))
-	}
+		defs := make([]string, 0, len(ctb.defs))
 
-	buf.WriteString(strings.Join(defs, ", "))
-	buf.WriteRune(')')
+		for _, def := range ctb.defs {
+			defs = append(defs, strings.Join(def, " "))
+		}
+
+		buf.WriteString(strings.Join(defs, ", "))
+		buf.WriteRune(')')
+
+		ctb.injection.WriteTo(buf, createTableMarkerAfterDefine)
+	}
 
 	if len(ctb.options) > 0 {
 		buf.WriteRune(' ')
@@ -108,6 +138,7 @@ func (ctb *CreateTableBuilder) BuildWithFlavor(flavor Flavor, initialArg ...inte
 		}
 
 		buf.WriteString(strings.Join(opts, ", "))
+		ctb.injection.WriteTo(buf, createTableMarkerAfterOption)
 	}
 
 	return ctb.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
@@ -119,3 +150,14 @@ func (ctb *CreateTableBuilder) SetFlavor(flavor Flavor) (old Flavor) {
 	ctb.args.Flavor = flavor
 	return
 }
+
+// Var returns a placeholder for value.
+func (ctb *CreateTableBuilder) Var(arg interface{}) string {
+	return ctb.args.Add(arg)
+}
+
+// SQL adds an arbitrary sql to current position.
+func (ctb *CreateTableBuilder) SQL(sql string) *CreateTableBuilder {
+	ctb.injection.SQL(ctb.marker, sql)
+	return ctb
+}

createtable_test.go

@@ -7,6 +7,17 @@ import (
 	"fmt"
 )
 
+func ExampleCreateTable() {
+	sql := CreateTable("demo.user").IfNotExists().
+		Define("id", "BIGINT(20)", "NOT NULL", "AUTO_INCREMENT", "PRIMARY KEY", `COMMENT "user id"`).
+		String()
+
+	fmt.Println(sql)
+
+	// Output:
+	// CREATE TABLE IF NOT EXISTS demo.user (id BIGINT(20) NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT "user id")
+}
+
 func ExampleCreateTableBuilder() {
 	ctb := NewCreateTableBuilder()
 	ctb.CreateTable("demo.user").IfNotExists()
@@ -38,3 +49,23 @@ func ExampleCreateTableBuilder_tempTable() {
 	// Output:
 	// CREATE TEMPORARY 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", created_at DATETIME NOT NULL COMMENT "user create time", modified_at DATETIME NOT NULL COMMENT "user modify time", KEY idx_name_modified_at name, modified_at) DEFAULT CHARACTER SET utf8mb4
 }
+
+func ExampleCreateTableBuilder_SQL() {
+	ctb := NewCreateTableBuilder()
+	ctb.SQL(`/* before */`)
+	ctb.CreateTempTable("demo.user").IfNotExists()
+	ctb.SQL("/* after create */")
+	ctb.Define("id", "BIGINT(20)", "NOT NULL", "AUTO_INCREMENT", "PRIMARY KEY", `COMMENT "user id"`)
+	ctb.Define("name", "VARCHAR(255)", "NOT NULL", `COMMENT "user name"`)
+	ctb.SQL("/* after define */")
+	ctb.Option("DEFAULT CHARACTER SET", "utf8mb4")
+	ctb.SQL(ctb.Var(Build("AS SELECT * FROM old.user WHERE name LIKE $?", "%Huan%")))
+
+	sql, args := ctb.Build()
+	fmt.Println(sql)
+	fmt.Println(args)
+
+	// Output:
+	// /* before */ CREATE TEMPORARY TABLE IF NOT EXISTS demo.user /* after create */ (id BIGINT(20) NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT "user id", name VARCHAR(255) NOT NULL COMMENT "user name") /* after define */ DEFAULT CHARACTER SET utf8mb4 AS SELECT * FROM old.user WHERE name LIKE ?
+	// [%Huan%]
+}

delete.go

@@ -5,9 +5,18 @@ package sqlbuilder
 
 import (
 	"bytes"
+	"strconv"
 	"strings"
 )
 
+const (
+	deleteMarkerInit injectionMarker = iota
+	deleteMarkerAfterDeleteFrom
+	deleteMarkerAfterWhere
+	deleteMarkerAfterOrderBy
+	deleteMarkerAfterLimit
+)
+
 // NewDeleteBuilder creates a new DELETE builder.
 func NewDeleteBuilder() *DeleteBuilder {
 	return DefaultFlavor.NewDeleteBuilder()
@@ -19,31 +28,74 @@ func newDeleteBuilder() *DeleteBuilder {
 		Cond: Cond{
 			Args: args,
 		},
-		args: args,
+		limit:     -1,
+		args:      args,
+		injection: newInjection(),
 	}
 }
 
 // DeleteBuilder is a builder to build DELETE.
 type DeleteBuilder struct {
 	Cond
 
-	table      string
-	whereExprs []string
+	table       string
+	whereExprs  []string
+	orderByCols []string
+	order       string
+	limit       int
 
 	args *Args
+
+	injection *injection
+	marker    injectionMarker
 }
 
 var _ Builder = new(DeleteBuilder)
 
+// DeleteFrom sets table name in DELETE.
+func DeleteFrom(table string) *DeleteBuilder {
+	return DefaultFlavor.NewDeleteBuilder().DeleteFrom(table)
+}
+
 // DeleteFrom sets table name in DELETE.
 func (db *DeleteBuilder) DeleteFrom(table string) *DeleteBuilder {
 	db.table = Escape(table)
+	db.marker = deleteMarkerAfterDeleteFrom
 	return db
 }
 
 // Where sets expressions of WHERE in DELETE.
 func (db *DeleteBuilder) Where(andExpr ...string) *DeleteBuilder {
 	db.whereExprs = append(db.whereExprs, andExpr...)
+	db.marker = deleteMarkerAfterWhere
+	return db
+}
+
+// OrderBy sets columns of ORDER BY in DELETE.
+func (db *DeleteBuilder) OrderBy(col ...string) *DeleteBuilder {
+	db.orderByCols = col
+	db.marker = deleteMarkerAfterOrderBy
+	return db
+}
+
+// Asc sets order of ORDER BY to ASC.
+func (db *DeleteBuilder) Asc() *DeleteBuilder {
+	db.order = "ASC"
+	db.marker = deleteMarkerAfterOrderBy
+	return db
+}
+
+// Desc sets order of ORDER BY to DESC.
+func (db *DeleteBuilder) Desc() *DeleteBuilder {
+	db.order = "DESC"
+	db.marker = deleteMarkerAfterOrderBy
+	return db
+}
+
+// Limit sets the LIMIT in DELETE.
+func (db *DeleteBuilder) Limit(limit int) *DeleteBuilder {
+	db.limit = limit
+	db.marker = deleteMarkerAfterLimit
 	return db
 }
 
@@ -63,12 +115,35 @@ func (db *DeleteBuilder) Build() (sql string, args []interface{}) {
 // They can be used in `DB#Query` of package `database/sql` directly.
 func (db *DeleteBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{}) (sql string, args []interface{}) {
 	buf := &bytes.Buffer{}
+	db.injection.WriteTo(buf, deleteMarkerInit)
 	buf.WriteString("DELETE FROM ")
 	buf.WriteString(db.table)
+	db.injection.WriteTo(buf, deleteMarkerAfterDeleteFrom)
 
 	if len(db.whereExprs) > 0 {
 		buf.WriteString(" WHERE ")
 		buf.WriteString(strings.Join(db.whereExprs, " AND "))
+
+		db.injection.WriteTo(buf, deleteMarkerAfterWhere)
+	}
+
+	if len(db.orderByCols) > 0 {
+		buf.WriteString(" ORDER BY ")
+		buf.WriteString(strings.Join(db.orderByCols, ", "))
+
+		if db.order != "" {
+			buf.WriteRune(' ')
+			buf.WriteString(db.order)
+		}
+
+		db.injection.WriteTo(buf, deleteMarkerAfterOrderBy)
+	}
+
+	if db.limit >= 0 {
+		buf.WriteString(" LIMIT ")
+		buf.WriteString(strconv.Itoa(db.limit))
+
+		db.injection.WriteTo(buf, deleteMarkerAfterLimit)
 	}
 
 	return db.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
@@ -80,3 +155,9 @@ func (db *DeleteBuilder) SetFlavor(flavor Flavor) (old Flavor) {
 	db.args.Flavor = flavor
 	return
 }
+
+// SQL adds an arbitrary sql to current position.
+func (db *DeleteBuilder) SQL(sql string) *DeleteBuilder {
+	db.injection.SQL(db.marker, sql)
+	return db
+}