package mssql
import (
"bytes"
_ "database/sql/driver"
"encoding/binary"
"fmt"
"golang.org/x/net/context" // use the "x/net/context" for backwards compatibility.
"math"
"reflect"
"strings"
"time"
)
type MssqlBulk struct {
cn *MssqlConn
metadata []columnStruct
bulkColumns []columnStruct
columnsName []string
tablename string
numRows int
headerSent bool
Options MssqlBulkOptions
Debug bool
}
type MssqlBulkOptions struct {
CheckConstraints bool
FireTriggers bool
KeepNulls bool
KilobytesPerBatch int
RowsPerBatch int
Order []string
Tablock bool
}
type DataValue interface{}
func (cn *MssqlConn) CreateBulk(table string, columns []string) (_ *MssqlBulk) {
b := MssqlBulk{cn: cn, tablename: table, headerSent: false, columnsName: columns}
b.Debug = false
return &b
}
func (b *MssqlBulk) sendBulkCommand() (err error) {
//get table columns info
err = b.getMetadata()
if err != nil {
return err
}
//match the columns
for _, colname := range b.columnsName {
var bulkCol *columnStruct
for _, m := range b.metadata {
if m.ColName == colname {
bulkCol = &m
break
}
}
if bulkCol != nil {
if bulkCol.ti.TypeId == typeUdt {
//send udt as binary
bulkCol.ti.TypeId = typeBigVarBin
}
b.bulkColumns = append(b.bulkColumns, *bulkCol)
b.dlogf("Adding column %s %s %#x", colname, bulkCol.ColName, bulkCol.ti.TypeId)
} else {
return fmt.Errorf("Column %s does not exist in destination table %s", colname, b.tablename)
}
}
//create the bulk command
//columns definitions
var col_defs bytes.Buffer
for i, col := range b.bulkColumns {
if i != 0 {
col_defs.WriteString(", ")
}
col_defs.WriteString("[" + col.ColName + "] " + makeDecl(col.ti))
}
//options
var with_opts []string
if b.Options.CheckConstraints {
with_opts = append(with_opts, "CHECK_CONSTRAINTS")
}
if b.Options.FireTriggers {
with_opts = append(with_opts, "FIRE_TRIGGERS")
}
if b.Options.KeepNulls {
with_opts = append(with_opts, "KEEP_NULLS")
}
if b.Options.KilobytesPerBatch > 0 {
with_opts = append(with_opts, fmt.Sprintf("KILOBYTES_PER_BATCH = %d", b.Options.KilobytesPerBatch))
}
if b.Options.RowsPerBatch > 0 {
with_opts = append(with_opts, fmt.Sprintf("ROWS_PER_BATCH = %d", b.Options.RowsPerBatch))
}
if len(b.Options.Order) > 0 {
with_opts = append(with_opts, fmt.Sprintf("ORDER(%s)", strings.Join(b.Options.Order, ",")))
}
if b.Options.Tablock {
with_opts = append(with_opts, "TABLOCK")
}
var with_part string
if len(with_opts) > 0 {
with_part = fmt.Sprintf("WITH (%s)", strings.Join(with_opts, ","))
}
query := fmt.Sprintf("INSERT BULK %s (%s) %s", b.tablename, col_defs.String(), with_part)
stmt, err := b.cn.Prepare(query)
if err != nil {
return fmt.Errorf("Prepare failed: %s", err.Error())
}
b.dlogf(query)
_, err = stmt.Exec(nil)
if err != nil {
return err
}
b.headerSent = true
var buf = b.cn.sess.buf
buf.BeginPacket(packBulkLoadBCP)
// send the columns metadata
columnMetadata := b.createColMetadata()
_, err = buf.Write(columnMetadata)
return
}
// AddRow immediately writes the row to the destination table.
// The arguments are the row values in the order they were specified.
func (b *MssqlBulk) AddRow(row []interface{}) (err error) {
if b.headerSent == false {
err = b.sendBulkCommand()
if err != nil {
return
}
}
if len(row) != len(b.bulkColumns) {
return fmt.Errorf("Row does not have the same number of columns than the destination table %d %d",
len(row), len(b.bulkColumns))
}
bytes, err := b.makeRowData(row)
if err != nil {
return
}
_, err = b.cn.sess.buf.Write(bytes)
if err != nil {
return
}
b.numRows = b.numRows + 1
return
}
func (b *MssqlBulk) makeRowData(row []interface{}) ([]byte, error) {
buf := new(bytes.Buffer)
buf.WriteByte(byte(tokenRow))
var logcol bytes.Buffer
for i, col := range b.bulkColumns {
if b.Debug {
logcol.WriteString(fmt.Sprintf(" col[%d]='%v' ", i, row[i]))
}
param, err := b.makeParam(row[i], col)
if err != nil {
return nil, fmt.Errorf("bulkcopy: %s", err.Error())
}
if col.ti.Writer == nil {
return nil, fmt.Errorf("no writer for column: %s, TypeId: %#x",
col.ColName, col.ti.TypeId)
}
err = col.ti.Writer(buf, param.ti, param.buffer)
if err != nil {
return nil, fmt.Errorf("bulkcopy: %s", err.Error())
}
}
b.dlogf("row[%d] %s\n", b.numRows, logcol.String())
return buf.Bytes(), nil
}
func (b *MssqlBulk) Done() (rowcount int64, err error) {
if b.headerSent == false {
//no rows had been sent
return 0, nil
}
var buf = b.cn.sess.buf
buf.WriteByte(byte(tokenDone))
binary.Write(buf, binary.LittleEndian, uint16(doneFinal))
binary.Write(buf, binary.LittleEndian, uint16(0)) // curcmd
if b.cn.sess.loginAck.TDSVersion >= verTDS72 {
binary.Write(buf, binary.LittleEndian, uint64(0)) //rowcount 0
} else {
binary.Write(buf, binary.LittleEndian, uint32(0)) //rowcount 0
}
buf.FinishPacket()
tokchan := make(chan tokenStruct, 5)
go processResponse(context.Background(), b.cn.sess, tokchan)
var rowCount int64
for token := range tokchan {
switch token := token.(type) {
case doneStruct:
if token.Status&doneCount != 0 {
rowCount = int64(token.RowCount)
}
if token.isError() {
return 0, token.getError()
}
case error:
return 0, token
}
}
return rowCount, nil
}
func (b *MssqlBulk) createColMetadata() []byte {
buf := new(bytes.Buffer)
buf.WriteByte(byte(tokenColMetadata)) // token
binary.Write(buf, binary.LittleEndian, uint16(len(b.bulkColumns))) // column count
for i, col := range b.bulkColumns {
if b.cn.sess.loginAck.TDSVersion >= verTDS72 {
binary.Write(buf, binary.LittleEndian, uint32(col.UserType)) // usertype, always 0?
} else {
binary.Write(buf, binary.LittleEndian, uint16(col.UserType))
}
binary.Write(buf, binary.LittleEndian, uint16(col.Flags))
writeTypeInfo(buf, &b.bulkColumns[i].ti)
if col.ti.TypeId == typeNText ||
col.ti.TypeId == typeText ||
col.ti.TypeId == typeImage {
tablename_ucs2 := str2ucs2(b.tablename)
binary.Write(buf, binary.LittleEndian, uint16(len(tablename_ucs2)/2))
buf.Write(tablename_ucs2)
}
colname_ucs2 := str2ucs2(col.ColName)
buf.WriteByte(uint8(len(colname_ucs2) / 2))
buf.Write(colname_ucs2)
}
return buf.Bytes()
}
func (b *MssqlBulk) getMetadata() (err error) {
stmt, err := b.cn.Prepare("SET FMTONLY ON")
if err != nil {
return
}
_, err = stmt.Exec(nil)
if err != nil {
return
}
//get columns info
stmt, err = b.cn.Prepare(fmt.Sprintf("select * from %s SET FMTONLY OFF", b.tablename))
if err != nil {
return
}
stmt2 := stmt.(*MssqlStmt)
cols, err := stmt2.QueryMeta()
if err != nil {
return fmt.Errorf("get columns info failed: %v", err.Error())
}
b.metadata = cols
if b.Debug {
for _, col := range b.metadata {
b.dlogf("col: %s typeId: %#x size: %d scale: %d prec: %d flags: %d lcid: %#x\n",
col.ColName, col.ti.TypeId, col.ti.Size, col.ti.Scale, col.ti.Prec,
col.Flags, col.ti.Collation.lcidAndFlags)
}
}
return nil
}
// QueryMeta is almost the same as MssqlStmt.Query, but returns all the columns info.
func (s *MssqlStmt) QueryMeta() (cols []columnStruct, err error) {
if err = s.sendQuery(nil); err != nil {
return
}
tokchan := make(chan tokenStruct, 5)
go processResponse(context.Background(), s.c.sess, tokchan)
loop:
for tok := range tokchan {
switch token := tok.(type) {
case doneStruct:
break loop
case []columnStruct:
cols = token
break loop
case error:
return nil, token
}
}
return cols, nil
}
func (b *MssqlBulk) makeParam(val DataValue, col columnStruct) (res Param, err error) {
res.ti.Size = col.ti.Size
res.ti.TypeId = col.ti.TypeId
if val == nil {
res.ti.Size = 0
return
}
switch col.ti.TypeId {
case typeInt1, typeInt2, typeInt4, typeInt8, typeIntN:
var intvalue int64
switch val := val.(type) {
case int:
intvalue = int64(val)
case int32:
intvalue = int64(val)
case int64:
intvalue = val
default:
err = fmt.Errorf("mssql: invalid type for int column")
return
}
res.buffer = make([]byte, res.ti.Size)
if col.ti.Size == 1 {
res.buffer[0] = byte(intvalue)
} else if col.ti.Size == 2 {
binary.LittleEndian.PutUint16(res.buffer, uint16(intvalue))
} else if col.ti.Size == 4 {
binary.LittleEndian.PutUint32(res.buffer, uint32(intvalue))
} else if col.ti.Size == 8 {
binary.LittleEndian.PutUint64(res.buffer, uint64(intvalue))
}
case typeFlt4, typeFlt8, typeFltN:
var floatvalue float64
switch val := val.(type) {
case float32:
floatvalue = float64(val)
case float64:
floatvalue = val
case int:
floatvalue = float64(val)
case int64:
floatvalue = float64(val)
default:
err = fmt.Errorf("mssql: invalid type for float column", val)
return
}
if col.ti.Size == 4 {
res.buffer = make([]byte, 4)
binary.LittleEndian.PutUint32(res.buffer, math.Float32bits(float32(floatvalue)))
} else if col.ti.Size == 8 {
res.buffer = make([]byte, 8)
binary.LittleEndian.PutUint64(res.buffer, math.Float64bits(floatvalue))
}
case typeNVarChar, typeNText, typeNChar:
switch val := val.(type) {
case string:
res.buffer = str2ucs2(val)
case []byte:
res.buffer = val
default:
err = fmt.Errorf("mssql: invalid type for nvarchar column", val)
return
}
res.ti.Size = len(res.buffer)
case typeVarChar, typeBigVarChar, typeText, typeChar, typeBigChar:
switch val := val.(type) {
case string:
res.buffer = []byte(val)
case []byte:
res.buffer = val
default:
err = fmt.Errorf("mssql: invalid type for varchar column", val)
return
}
res.ti.Size = len(res.buffer)
case typeBit, typeBitN:
if reflect.TypeOf(val).Kind() != reflect.Bool {
err = fmt.Errorf("mssql: invalid type for bit column", val)
return
}
res.ti.TypeId = typeBitN
res.ti.Size = 1
res.buffer = make([]byte, 1)
if val.(bool) {
res.buffer[0] = 1
}
case typeDateTime2N, typeDateTimeOffsetN:
switch val := val.(type) {
case time.Time:
days, ns := dateTime2(val)
ns /= int64(math.Pow10(int(col.ti.Scale)*-1) * 1000000000)
var data = make([]byte, 5)
data[0] = byte(ns)
data[1] = byte(ns >> 8)
data[2] = byte(ns >> 16)
data[3] = byte(ns >> 24)
data[4] = byte(ns >> 32)
if col.ti.Scale <= 2 {
res.ti.Size = 6
} else if col.ti.Scale <= 4 {
res.ti.Size = 7
} else {
res.ti.Size = 8
}
var buf []byte
buf = make([]byte, res.ti.Size)
copy(buf, data[0:res.ti.Size-3])
buf[res.ti.Size-3] = byte(days)
buf[res.ti.Size-2] = byte(days >> 8)
buf[res.ti.Size-1] = byte(days >> 16)
if col.ti.TypeId == typeDateTimeOffsetN {
_, offset := val.Zone()
var offsetMinute = uint16(offset / 60)
buf = append(buf, byte(offsetMinute))
buf = append(buf, byte(offsetMinute>>8))
res.ti.Size = res.ti.Size + 2
}
res.buffer = buf
default:
err = fmt.Errorf("mssql: invalid type for datetime2 column", val)
return
}
case typeDateN:
switch val := val.(type) {
case time.Time:
days, _ := dateTime2(val)
res.ti.Size = 3
res.buffer = make([]byte, 3)
res.buffer[0] = byte(days)
res.buffer[1] = byte(days >> 8)
res.buffer[2] = byte(days >> 16)
default:
err = fmt.Errorf("mssql: invalid type for date column", val)
return
}
case typeDateTime, typeDateTimeN, typeDateTim4:
switch val := val.(type) {
case time.Time:
if col.ti.Size == 4 {
res.ti.Size = 4
res.buffer = make([]byte, 4)
ref := time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)
dur := val.Sub(ref)
days := dur / (24 * time.Hour)
if days < 0 {
err = fmt.Errorf("mssql: Date %s is out of range", val)
return
}
mins := val.Hour()*60 + val.Minute()
binary.LittleEndian.PutUint16(res.buffer[0:2], uint16(days))
binary.LittleEndian.PutUint16(res.buffer[2:4], uint16(mins))
} else if col.ti.Size == 8 {
res.ti.Size = 8
res.buffer = make([]byte, 8)
ref := time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)
dur := val.Sub(ref)
days := math.Floor(float64(dur) / float64(24*time.Hour))
tm := (val.Hour()*60*60+val.Minute()*60+val.Second())*300 + int(val.Nanosecond()/10000000*3)
binary.LittleEndian.PutUint32(res.buffer[0:4], uint32(days))
binary.LittleEndian.PutUint32(res.buffer[4:8], uint32(tm))
} else {
err = fmt.Errorf("mssql: invalid size of column")
}
default:
err = fmt.Errorf("mssql: invalid type for datetime column", val)
}
/*
case typeDecimal, typeDecimalN:
switch val := val.(type) {
case float64:
dec, err := Float64ToDecimal(val)
if err != nil {
return res, err
}
dec.scale = col.ti.Scale
dec.prec = col.ti.Prec
//res.buffer = make([]byte, 3)
res.buffer = dec.Bytes()
res.ti.Size = len(res.buffer)
default:
err = fmt.Errorf("mssql: invalid type for decimal column", val)
return
}
case typeMoney, typeMoney4, typeMoneyN:
if col.ti.Size == 4 {
res.ti.Size = 4
res.buffer = make([]byte, 4)
} else if col.ti.Size == 8 {
res.ti.Size = 8
res.buffer = make([]byte, 8)
} else {
err = fmt.Errorf("mssql: invalid size of money column")
}
*/
case typeBigVarBin:
switch val := val.(type) {
case []byte:
res.ti.Size = len(val)
res.buffer = val
default:
err = fmt.Errorf("mssql: invalid type for Binary column", val)
return
}
default:
err = fmt.Errorf("mssql: type %x not implemented!", col.ti.TypeId)
}
return
}
func (b *MssqlBulk) dlogf(format string, v ...interface{}) {
if b.Debug {
b.cn.sess.log.Printf(format, v...)
}
}