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