1package squirrel
2
3import (
4	"bytes"
5	"database/sql"
6	"errors"
7	"fmt"
8	"io"
9	"sort"
10	"strings"
11
12	"github.com/lann/builder"
13)
14
15type insertData struct {
16	PlaceholderFormat PlaceholderFormat
17	RunWith           BaseRunner
18	Prefixes          []Sqlizer
19	StatementKeyword  string
20	Options           []string
21	Into              string
22	Columns           []string
23	Values            [][]interface{}
24	Suffixes          []Sqlizer
25	Select            *SelectBuilder
26}
27
28func (d *insertData) Exec() (sql.Result, error) {
29	if d.RunWith == nil {
30		return nil, RunnerNotSet
31	}
32	return ExecWith(d.RunWith, d)
33}
34
35func (d *insertData) Query() (*sql.Rows, error) {
36	if d.RunWith == nil {
37		return nil, RunnerNotSet
38	}
39	return QueryWith(d.RunWith, d)
40}
41
42func (d *insertData) QueryRow() RowScanner {
43	if d.RunWith == nil {
44		return &Row{err: RunnerNotSet}
45	}
46	queryRower, ok := d.RunWith.(QueryRower)
47	if !ok {
48		return &Row{err: RunnerNotQueryRunner}
49	}
50	return QueryRowWith(queryRower, d)
51}
52
53func (d *insertData) ToSql() (sqlStr string, args []interface{}, err error) {
54	if len(d.Into) == 0 {
55		err = errors.New("insert statements must specify a table")
56		return
57	}
58	if len(d.Values) == 0 && d.Select == nil {
59		err = errors.New("insert statements must have at least one set of values or select clause")
60		return
61	}
62
63	sql := &bytes.Buffer{}
64
65	if len(d.Prefixes) > 0 {
66		args, err = appendToSql(d.Prefixes, sql, " ", args)
67		if err != nil {
68			return
69		}
70
71		sql.WriteString(" ")
72	}
73
74	if d.StatementKeyword == "" {
75		sql.WriteString("INSERT ")
76	} else {
77		sql.WriteString(d.StatementKeyword)
78		sql.WriteString(" ")
79	}
80
81	if len(d.Options) > 0 {
82		sql.WriteString(strings.Join(d.Options, " "))
83		sql.WriteString(" ")
84	}
85
86	sql.WriteString("INTO ")
87	sql.WriteString(d.Into)
88	sql.WriteString(" ")
89
90	if len(d.Columns) > 0 {
91		sql.WriteString("(")
92		sql.WriteString(strings.Join(d.Columns, ","))
93		sql.WriteString(") ")
94	}
95
96	if d.Select != nil {
97		args, err = d.appendSelectToSQL(sql, args)
98	} else {
99		args, err = d.appendValuesToSQL(sql, args)
100	}
101	if err != nil {
102		return
103	}
104
105	if len(d.Suffixes) > 0 {
106		sql.WriteString(" ")
107		args, err = appendToSql(d.Suffixes, sql, " ", args)
108		if err != nil {
109			return
110		}
111	}
112
113	sqlStr, err = d.PlaceholderFormat.ReplacePlaceholders(sql.String())
114	return
115}
116
117func (d *insertData) appendValuesToSQL(w io.Writer, args []interface{}) ([]interface{}, error) {
118	if len(d.Values) == 0 {
119		return args, errors.New("values for insert statements are not set")
120	}
121
122	io.WriteString(w, "VALUES ")
123
124	valuesStrings := make([]string, len(d.Values))
125	for r, row := range d.Values {
126		valueStrings := make([]string, len(row))
127		for v, val := range row {
128			if vs, ok := val.(Sqlizer); ok {
129				vsql, vargs, err := vs.ToSql()
130				if err != nil {
131					return nil, err
132				}
133				valueStrings[v] = vsql
134				args = append(args, vargs...)
135			} else {
136				valueStrings[v] = "?"
137				args = append(args, val)
138			}
139		}
140		valuesStrings[r] = fmt.Sprintf("(%s)", strings.Join(valueStrings, ","))
141	}
142
143	io.WriteString(w, strings.Join(valuesStrings, ","))
144
145	return args, nil
146}
147
148func (d *insertData) appendSelectToSQL(w io.Writer, args []interface{}) ([]interface{}, error) {
149	if d.Select == nil {
150		return args, errors.New("select clause for insert statements are not set")
151	}
152
153	selectClause, sArgs, err := d.Select.ToSql()
154	if err != nil {
155		return args, err
156	}
157
158	io.WriteString(w, selectClause)
159	args = append(args, sArgs...)
160
161	return args, nil
162}
163
164// Builder
165
166// InsertBuilder builds SQL INSERT statements.
167type InsertBuilder builder.Builder
168
169func init() {
170	builder.Register(InsertBuilder{}, insertData{})
171}
172
173// Format methods
174
175// PlaceholderFormat sets PlaceholderFormat (e.g. Question or Dollar) for the
176// query.
177func (b InsertBuilder) PlaceholderFormat(f PlaceholderFormat) InsertBuilder {
178	return builder.Set(b, "PlaceholderFormat", f).(InsertBuilder)
179}
180
181// Runner methods
182
183// RunWith sets a Runner (like database/sql.DB) to be used with e.g. Exec.
184func (b InsertBuilder) RunWith(runner BaseRunner) InsertBuilder {
185	return setRunWith(b, runner).(InsertBuilder)
186}
187
188// Exec builds and Execs the query with the Runner set by RunWith.
189func (b InsertBuilder) Exec() (sql.Result, error) {
190	data := builder.GetStruct(b).(insertData)
191	return data.Exec()
192}
193
194// Query builds and Querys the query with the Runner set by RunWith.
195func (b InsertBuilder) Query() (*sql.Rows, error) {
196	data := builder.GetStruct(b).(insertData)
197	return data.Query()
198}
199
200// QueryRow builds and QueryRows the query with the Runner set by RunWith.
201func (b InsertBuilder) QueryRow() RowScanner {
202	data := builder.GetStruct(b).(insertData)
203	return data.QueryRow()
204}
205
206// Scan is a shortcut for QueryRow().Scan.
207func (b InsertBuilder) Scan(dest ...interface{}) error {
208	return b.QueryRow().Scan(dest...)
209}
210
211// SQL methods
212
213// ToSql builds the query into a SQL string and bound args.
214func (b InsertBuilder) ToSql() (string, []interface{}, error) {
215	data := builder.GetStruct(b).(insertData)
216	return data.ToSql()
217}
218
219// Prefix adds an expression to the beginning of the query
220func (b InsertBuilder) Prefix(sql string, args ...interface{}) InsertBuilder {
221	return b.PrefixExpr(Expr(sql, args...))
222}
223
224// PrefixExpr adds an expression to the very beginning of the query
225func (b InsertBuilder) PrefixExpr(expr Sqlizer) InsertBuilder {
226	return builder.Append(b, "Prefixes", expr).(InsertBuilder)
227}
228
229// Options adds keyword options before the INTO clause of the query.
230func (b InsertBuilder) Options(options ...string) InsertBuilder {
231	return builder.Extend(b, "Options", options).(InsertBuilder)
232}
233
234// Into sets the INTO clause of the query.
235func (b InsertBuilder) Into(from string) InsertBuilder {
236	return builder.Set(b, "Into", from).(InsertBuilder)
237}
238
239// Columns adds insert columns to the query.
240func (b InsertBuilder) Columns(columns ...string) InsertBuilder {
241	return builder.Extend(b, "Columns", columns).(InsertBuilder)
242}
243
244// Values adds a single row's values to the query.
245func (b InsertBuilder) Values(values ...interface{}) InsertBuilder {
246	return builder.Append(b, "Values", values).(InsertBuilder)
247}
248
249// Suffix adds an expression to the end of the query
250func (b InsertBuilder) Suffix(sql string, args ...interface{}) InsertBuilder {
251	return b.SuffixExpr(Expr(sql, args...))
252}
253
254// SuffixExpr adds an expression to the end of the query
255func (b InsertBuilder) SuffixExpr(expr Sqlizer) InsertBuilder {
256	return builder.Append(b, "Suffixes", expr).(InsertBuilder)
257}
258
259// SetMap set columns and values for insert builder from a map of column name and value
260// note that it will reset all previous columns and values was set if any
261func (b InsertBuilder) SetMap(clauses map[string]interface{}) InsertBuilder {
262	// Keep the columns in a consistent order by sorting the column key string.
263	cols := make([]string, 0, len(clauses))
264	for col := range clauses {
265		cols = append(cols, col)
266	}
267	sort.Strings(cols)
268
269	vals := make([]interface{}, 0, len(clauses))
270	for _, col := range cols {
271		vals = append(vals, clauses[col])
272	}
273
274	b = builder.Set(b, "Columns", cols).(InsertBuilder)
275	b = builder.Set(b, "Values", [][]interface{}{vals}).(InsertBuilder)
276
277	return b
278}
279
280// Select set Select clause for insert query
281// If Values and Select are used, then Select has higher priority
282func (b InsertBuilder) Select(sb SelectBuilder) InsertBuilder {
283	return builder.Set(b, "Select", &sb).(InsertBuilder)
284}
285
286func (b InsertBuilder) statementKeyword(keyword string) InsertBuilder {
287	return builder.Set(b, "StatementKeyword", keyword).(InsertBuilder)
288}
289