1
1
mirror of https://github.com/go-gitea/gitea synced 2024-12-24 01:24:26 +00:00
gitea/vendor/github.com/denisenkom/go-mssqldb/bulkcopy.go

584 lines
14 KiB
Go
Vendored

package mssql
import (
"bytes"
"context"
"encoding/binary"
"fmt"
"math"
"reflect"
"strings"
"time"
"github.com/denisenkom/go-mssqldb/internal/decimal"
)
type Bulk struct {
// ctx is used only for AddRow and Done methods.
// This could be removed if AddRow and Done accepted
// a ctx field as well, which is available with the
// database/sql call.
ctx context.Context
cn *Conn
metadata []columnStruct
bulkColumns []columnStruct
columnsName []string
tablename string
numRows int
headerSent bool
Options BulkOptions
Debug bool
}
type BulkOptions struct {
CheckConstraints bool
FireTriggers bool
KeepNulls bool
KilobytesPerBatch int
RowsPerBatch int
Order []string
Tablock bool
}
type DataValue interface{}
const (
sqlDateFormat = "2006-01-02"
sqlTimeFormat = "2006-01-02 15:04:05.999999999Z07:00"
)
func (cn *Conn) CreateBulk(table string, columns []string) (_ *Bulk) {
b := Bulk{ctx: context.Background(), cn: cn, tablename: table, headerSent: false, columnsName: columns}
b.Debug = false
return &b
}
func (cn *Conn) CreateBulkContext(ctx context.Context, table string, columns []string) (_ *Bulk) {
b := Bulk{ctx: ctx, cn: cn, tablename: table, headerSent: false, columnsName: columns}
b.Debug = false
return &b
}
func (b *Bulk) sendBulkCommand(ctx context.Context) (err error) {
//get table columns info
err = b.getMetadata(ctx)
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.PrepareContext(ctx, query)
if err != nil {
return fmt.Errorf("Prepare failed: %s", err.Error())
}
b.dlogf(query)
_, err = stmt.(*Stmt).ExecContext(ctx, nil)
if err != nil {
return err
}
b.headerSent = true
var buf = b.cn.sess.buf
buf.BeginPacket(packBulkLoadBCP, false)
// 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 *Bulk) AddRow(row []interface{}) (err error) {
if !b.headerSent {
err = b.sendBulkCommand(b.ctx)
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 *Bulk) 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 *Bulk) 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(b.ctx, b.cn.sess, tokchan, nil)
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, b.cn.checkBadConn(token)
}
}
return rowCount, nil
}
func (b *Bulk) 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 *Bulk) getMetadata(ctx context.Context) (err error) {
stmt, err := b.cn.prepareContext(ctx, "SET FMTONLY ON")
if err != nil {
return
}
_, err = stmt.ExecContext(ctx, nil)
if err != nil {
return
}
// Get columns info.
stmt, err = b.cn.prepareContext(ctx, fmt.Sprintf("select * from %s SET FMTONLY OFF", b.tablename))
if err != nil {
return
}
rows, err := stmt.QueryContext(ctx, nil)
if err != nil {
return fmt.Errorf("get columns info failed: %v", err)
}
b.metadata = rows.(*Rows).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 rows.Close()
}
func (b *Bulk) 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: %T", val)
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: %T %s", val, 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: %T %s", val, 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: %T %s", val, 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: %T %s", val, val)
return
}
res.ti.TypeId = typeBitN
res.ti.Size = 1
res.buffer = make([]byte, 1)
if val.(bool) {
res.buffer[0] = 1
}
case typeDateTime2N:
switch val := val.(type) {
case time.Time:
res.buffer = encodeDateTime2(val, int(col.ti.Scale))
res.ti.Size = len(res.buffer)
case string:
var t time.Time
if t, err = time.Parse(sqlTimeFormat, val); err != nil {
return res, fmt.Errorf("bulk: unable to convert string to date: %v", err)
}
res.buffer = encodeDateTime2(t, int(col.ti.Scale))
res.ti.Size = len(res.buffer)
default:
err = fmt.Errorf("mssql: invalid type for datetime2 column: %T %s", val, val)
return
}
case typeDateTimeOffsetN:
switch val := val.(type) {
case time.Time:
res.buffer = encodeDateTimeOffset(val, int(col.ti.Scale))
res.ti.Size = len(res.buffer)
case string:
var t time.Time
if t, err = time.Parse(sqlTimeFormat, val); err != nil {
return res, fmt.Errorf("bulk: unable to convert string to date: %v", err)
}
res.buffer = encodeDateTimeOffset(t, int(col.ti.Scale))
res.ti.Size = len(res.buffer)
default:
err = fmt.Errorf("mssql: invalid type for datetimeoffset column: %T %s", val, val)
return
}
case typeDateN:
switch val := val.(type) {
case time.Time:
res.buffer = encodeDate(val)
res.ti.Size = len(res.buffer)
case string:
var t time.Time
if t, err = time.ParseInLocation(sqlDateFormat, val, time.UTC); err != nil {
return res, fmt.Errorf("bulk: unable to convert string to date: %v", err)
}
res.buffer = encodeDate(t)
res.ti.Size = len(res.buffer)
default:
err = fmt.Errorf("mssql: invalid type for date column: %T %s", val, val)
return
}
case typeDateTime, typeDateTimeN, typeDateTim4:
var t time.Time
switch val := val.(type) {
case time.Time:
t = val
case string:
if t, err = time.Parse(sqlTimeFormat, val); err != nil {
return res, fmt.Errorf("bulk: unable to convert string to date: %v", err)
}
default:
err = fmt.Errorf("mssql: invalid type for datetime column: %T %s", val, val)
return
}
if col.ti.Size == 4 {
res.buffer = encodeDateTim4(t)
res.ti.Size = len(res.buffer)
} else if col.ti.Size == 8 {
res.buffer = encodeDateTime(t)
res.ti.Size = len(res.buffer)
} else {
err = fmt.Errorf("mssql: invalid size of column %d", col.ti.Size)
}
// case typeMoney, typeMoney4, typeMoneyN:
case typeDecimal, typeDecimalN, typeNumeric, typeNumericN:
prec := col.ti.Prec
scale := col.ti.Scale
var dec decimal.Decimal
switch v := val.(type) {
case int:
dec = decimal.Int64ToDecimalScale(int64(v), 0)
case int8:
dec = decimal.Int64ToDecimalScale(int64(v), 0)
case int16:
dec = decimal.Int64ToDecimalScale(int64(v), 0)
case int32:
dec = decimal.Int64ToDecimalScale(int64(v), 0)
case int64:
dec = decimal.Int64ToDecimalScale(int64(v), 0)
case float32:
dec, err = decimal.Float64ToDecimalScale(float64(v), scale)
case float64:
dec, err = decimal.Float64ToDecimalScale(float64(v), scale)
case string:
dec, err = decimal.StringToDecimalScale(v, scale)
default:
return res, fmt.Errorf("unknown value for decimal: %T %#v", v, v)
}
if err != nil {
return res, err
}
dec.SetPrec(prec)
var length byte
switch {
case prec <= 9:
length = 4
case prec <= 19:
length = 8
case prec <= 28:
length = 12
default:
length = 16
}
buf := make([]byte, length+1)
// first byte length written by typeInfo.writer
res.ti.Size = int(length) + 1
// second byte sign
if !dec.IsPositive() {
buf[0] = 0
} else {
buf[0] = 1
}
ub := dec.UnscaledBytes()
l := len(ub)
if l > int(length) {
err = fmt.Errorf("decimal out of range: %s", dec)
return res, err
}
// reverse the bytes
for i, j := 1, l-1; j >= 0; i, j = i+1, j-1 {
buf[i] = ub[j]
}
res.buffer = buf
case typeBigVarBin, typeBigBinary:
switch val := val.(type) {
case []byte:
res.ti.Size = len(val)
res.buffer = val
default:
err = fmt.Errorf("mssql: invalid type for Binary column: %T %s", val, val)
return
}
case typeGuid:
switch val := val.(type) {
case []byte:
res.ti.Size = len(val)
res.buffer = val
default:
err = fmt.Errorf("mssql: invalid type for Guid column: %T %s", val, val)
return
}
default:
err = fmt.Errorf("mssql: type %x not implemented", col.ti.TypeId)
}
return
}
func (b *Bulk) dlogf(format string, v ...interface{}) {
if b.Debug {
b.cn.sess.log.Printf(format, v...)
}
}