<commit_before>package dbr

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"time"
)

// InsertBuilder contains the clauses for an INSERT statement
type InsertBuilder struct {
	*Session
	runner

	Into string
	Cols []string
	Vals [][]interface{}
	Recs []interface{}
	Maps map[string]interface{}
}

// InsertInto instantiates a InsertBuilder for the given table
func (sess *Session) InsertInto(into string) *InsertBuilder {
	return &InsertBuilder{
		Session: sess,
		runner:  sess.cxn.Db,
		Into:    into,
	}
}

// InsertInto instantiates a InsertBuilder for the given table bound to a transaction
func (tx *Tx) InsertInto(into string) *InsertBuilder {
	return &InsertBuilder{
		Session: tx.Session,
		runner:  tx.Tx,
		Into:    into,
	}
}

// Columns appends columns to insert in the statement
func (b *InsertBuilder) Columns(columns ...string) *InsertBuilder {
	b.Cols = columns
	return b
}

// Values appends a set of values to the statement
func (b *InsertBuilder) Values(vals ...interface{}) *InsertBuilder {
	b.Vals = append(b.Vals, vals)
	return b
}

// Record pulls in values to match Columns from the record
func (b *InsertBuilder) Record(record interface{}) *InsertBuilder {
	b.Recs = append(b.Recs, record)
	return b
}

// Record pulls in values to match Columns from the record
func (b *InsertBuilder) Map(m map[string]interface{}) *InsertBuilder {
	b.Maps = m
	return b
}

// Pair adds a key/value pair to the statement
func (b *InsertBuilder) Pair(column string, value interface{}) *InsertBuilder {
	b.Cols = append(b.Cols, column)
	lenVals := len(b.Vals)
	if lenVals == 0 {
		args := []interface{}{value}
		b.Vals = [][]interface{}{args}
	} else if lenVals == 1 {
		b.Vals[0] = append(b.Vals[0], value)
	} else {
		panic("pair only allows you to specify 1 record to insret")
	}
	return b
}

// ToSql serialized the InsertBuilder to a SQL string
// It returns the string with placeholders and a slice of query arguments
func (b *InsertBuilder) ToSql() (string, []interface{}) {
	if len(b.Into) == 0 {
		panic("no table specified")
	}
	if len(b.Cols) == 0 && len(b.Maps) == 0 {
		panic("no columns or map specified")
	} else if len(b.Maps) == 0 {
		if len(b.Vals) == 0 && len(b.Recs) == 0 {
			panic("no values or records specified")
		}
		if len(b.Cols) == 0 && (len(b.Vals) > 0 || len(b.Recs) > 0) {
			panic("no columns specified")
		}
	}

	var sql bytes.Buffer

	sql.WriteString("INSERT INTO ")
	sql.WriteString(b.Into)
	sql.WriteString(" (")

	if len(b.Maps) != 0 {
		return b.MapToSql(sql)
	}

	var args []interface{}
	var placeholder bytes.Buffer // Build the placeholder like "(?,?,?)"
	// Simulataneously write the cols to the sql buffer, and build a placeholder
	placeholder.WriteRune('(')
	for i, c := range b.Cols {
		if i > 0 {
			sql.WriteRune(',')
			placeholder.WriteRune(',')
		}
		Quoter.writeQuotedColumn(c, &sql)
		placeholder.WriteRune('?')
	}
	sql.WriteString(") VALUES ")
	placeholder.WriteRune(')')
	placeholderStr := placeholder.String()

	// Go thru each value we want to insert. Write the placeholders, and collect args
	for i, row := range b.Vals {
		if i > 0 {
			sql.WriteRune(',')
		}
		sql.WriteString(placeholderStr)

		for _, v := range row {
			args = append(args, v)
		}
	}
	anyVals := len(b.Vals) > 0

	// Go thru the records. Write the placeholders, and do reflection on the records to extract args
	for i, rec := range b.Recs {
		if i > 0 || anyVals {
			sql.WriteRune(',')
		}
		sql.WriteString(placeholderStr)

		ind := reflect.Indirect(reflect.ValueOf(rec))
		vals, err := b.valuesFor(ind.Type(), ind, b.Cols)
		if err != nil {
			panic(err.Error())
		}
		for _, v := range vals {
			args = append(args, v)
		}
	}

	return sql.String(), args
}

// MapToSql serialized the InsertBuilder to a SQL string
// It goes through the Maps param and combined its keys/values into the SQL query string
// It returns the string with placeholders and a slice of query arguments
func (b *InsertBuilder) MapToSql(sql bytes.Buffer) (string, []interface{}) {
	keys := make([]string, len(b.Maps))
	vals := make([]interface{}, len(b.Maps))
	i := 0
	for k, v := range b.Maps {
		keys[i] = k
		vals[i] = v
		i++
	}
	var args []interface{}
	var placeholder bytes.Buffer // Build the placeholder like "(?,?,?)"

	placeholder.WriteRune('(')
	for i, c := range keys {
		if i > 0 {
			sql.WriteRune(',')
			placeholder.WriteRune(',')
		}
		Quoter.writeQuotedColumn(c, &sql)
		placeholder.WriteRune('?')
	}
	sql.WriteString(") VALUES ")
	placeholder.WriteRune(')')
	sql.WriteString(placeholder.String())

	for _, row := range vals {
		args = append(args, row)
	}

	return sql.String(), args<commit_msg>Package dbr: Insert statements values can now read data from driver.Valuer<commit_after>package dbr

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"reflect"
	"time"
)

// InsertBuilder contains the clauses for an INSERT statement
type InsertBuilder struct {
	*Session
	runner

	Into string
	Cols []string
	Vals [][]interface{}
	Recs []interface{}
	Maps map[string]interface{}
}

// InsertInto instantiates a InsertBuilder for the given table
func (sess *Session) InsertInto(into string) *InsertBuilder {
	return &InsertBuilder{
		Session: sess,
		runner:  sess.cxn.Db,
		Into:    into,
	}
}

// InsertInto instantiates a InsertBuilder for the given table bound to a transaction
func (tx *Tx) InsertInto(into string) *InsertBuilder {
	return &InsertBuilder{
		Session: tx.Session,
		runner:  tx.Tx,
		Into:    into,
	}
}

// Columns appends columns to insert in the statement
func (b *InsertBuilder) Columns(columns ...string) *InsertBuilder {
	b.Cols = columns
	return b
}

// Values appends a set of values to the statement.
// Pro Tip: Use Values() and not Record() to avoid reflection.
// Only this function will consider the driver.Valuer interface when you pass
// a pointer to the value.
func (b *InsertBuilder) Values(vals ...interface{}) *InsertBuilder {
	for i, v := range vals {
		if dbVal, ok := v.(driver.Valuer); ok {
			if val, err := dbVal.Value(); err == nil {
				vals[i] = val // overrides the current value ...
			} else {
				panic(err)
			}
		}
	}
	b.Vals = append(b.Vals, vals)
	return b
}

// Record pulls in values to match Columns from the record. Uses reflection.
func (b *InsertBuilder) Record(record interface{}) *InsertBuilder {
	b.Recs = append(b.Recs, record)
	return b
}

// Record pulls in values to match Columns from the record
func (b *InsertBuilder) Map(m map[string]interface{}) *InsertBuilder {
	b.Maps = m
	return b
}

// Pair adds a key/value pair to the statement. Uses not reflection.
func (b *InsertBuilder) Pair(column string, value interface{}) *InsertBuilder {
	if dbVal, ok := value.(driver.Valuer); ok {
		if val, err := dbVal.Value(); err == nil {
			value = val // overrides the current value ...
		} else {
			panic(err)
		}
	}

	b.Cols = append(b.Cols, column)
	lenVals := len(b.Vals)
	if lenVals == 0 {
		args := []interface{}{value}
		b.Vals = [][]interface{}{args}
	} else if lenVals == 1 {
		b.Vals[0] = append(b.Vals[0], value)
	} else {
		panic("pair only allows you to specify 1 record to insret")
	}
	return b
}

// ToSql serialized the InsertBuilder to a SQL string
// It returns the string with placeholders and a slice of query arguments
func (b *InsertBuilder) ToSql() (string, []interface{}) {
	if len(b.Into) == 0 {
		panic("no table specified")
	}
	if len(b.Cols) == 0 && len(b.Maps) == 0 {
		panic("no columns or map specified")
	} else if len(b.Maps) == 0 {
		if len(b.Vals) == 0 && len(b.Recs) == 0 {
			panic("no values or records specified")
		}
		if len(b.Cols) == 0 && (len(b.Vals) > 0 || len(b.Recs) > 0) {
			panic("no columns specified")
		}
	}

	var sql bytes.Buffer

	sql.WriteString("INSERT INTO ")
	sql.WriteString(b.Into)
	sql.WriteString(" (")

	if len(b.Maps) != 0 {
		return b.MapToSql(sql)
	}

	var args []interface{}
	var placeholder bytes.Buffer // Build the placeholder like "(?,?,?)"
	// Simulataneously write the cols to the sql buffer, and build a placeholder
	placeholder.WriteRune('(')
	for i, c := range b.Cols {
		if i > 0 {
			sql.WriteRune(',')
			placeholder.WriteRune(',')
		}
		Quoter.writeQuotedColumn(c, &sql)
		placeholder.WriteRune('?')
	}
	sql.WriteString(") VALUES ")
	placeholder.WriteRune(')')
	placeholderStr := placeholder.String()

	// Go thru each value we want to insert. Write the placeholders, and collect args
	for i, row := range b.Vals {
		if i > 0 {
			sql.WriteRune(',')
		}
		sql.WriteString(placeholderStr)

		for _, v := range row {
			args = append(args, v)
		}
	}
	anyVals := len(b.Vals) > 0

	// Go thru the records. Write the placeholders, and do reflection on the records to extract args
	for i, rec := range b.Recs {
		if i > 0 || anyVals {
			sql.WriteRune(',')
		}
		sql.WriteString(placeholderStr)

		ind := reflect.Indirect(reflect.ValueOf(rec))
		vals, err := b.valuesFor(ind.Type(), ind, b.Cols)
		if err != nil {
			panic(err.Error())
		}
		for _, v := range vals {
			args = append(args, v)
		}
	}

	return sql.String(), args
}

// MapToSql serialized the InsertBuilder to a SQL string
// It goes through the Maps param and combined its keys/values into the SQL query string
// It returns the string with placeholders and a slice of query arguments
func (b *InsertBuilder) MapToSql(sql bytes.Buffer) (string, []interface{}) {
	keys := make([]string, len(b.Maps))
	vals := make([]interface{}, len(b.Maps))
	i := 0
	for k, v := range b.Maps {
		keys[i] = k
		if dbVal, ok := v.(driver.Valuer); ok {
			if val, err := dbVal.Value(); err == nil {
				vals[i] = val
			} else {
				panic(err)
			}
		} else {
			vals[i] = v
		}
		i++
	}
	var args []interface{}
	var placeholder bytes.Buffer // Build the placeholder like "(?,?,?)"

	placeholder.WriteRune('(')
	for i, c := range keys {
		if i > 0 {
			sql.WriteRune(',')
			placeholder.WriteRune(',')
		}
		Quoter.writeQuotedColumn(c, &sql)
		placeholder.WriteRune('?')
	}
	sql.WriteString(") VALUES ")
	placeholder.WriteRune(')')
	sql.WriteString(placeholder.String())

	for _, row := range vals {
		args = append(args, row)
	}

	return sql.String(), args
{'fraction_non_alphanumeric': 0.14288545382481174, 'fraction_numerical': 0.004161712247324614, 'mean_word_length': 3.523980277902286, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3475928', 'n_tokens_mistral': 3738, 'n_tokens_neox': 3470, 'n_words': 1110}