473 lines
10 KiB
Go
473 lines
10 KiB
Go
|
package mssql
|
||
|
|
|
||
|
|
import (
|
||
|
|
"database/sql"
|
||
|
|
"database/sql/driver"
|
||
|
|
"encoding/binary"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"log"
|
||
|
|
"math"
|
||
|
|
"net"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
var partnersCache partners = partners{mu: sync.RWMutex{}, v: make(map[string]string)}
|
||
|
|
|
||
|
|
func init() {
|
||
|
|
sql.Register("mssql", &MssqlDriver{})
|
||
|
|
}
|
||
|
|
|
||
|
|
type MssqlDriver struct {
|
||
|
|
log *log.Logger
|
||
|
|
}
|
||
|
|
|
||
|
|
func (d *MssqlDriver) SetLogger(logger *log.Logger) {
|
||
|
|
d.log = logger
|
||
|
|
}
|
||
|
|
|
||
|
|
func CheckBadConn(err error) error {
|
||
|
|
if err == io.EOF {
|
||
|
|
return driver.ErrBadConn
|
||
|
|
}
|
||
|
|
neterr, ok := err.(net.Error)
|
||
|
|
if !ok || (!neterr.Timeout() && neterr.Temporary()) {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return driver.ErrBadConn
|
||
|
|
}
|
||
|
|
|
||
|
|
type MssqlConn struct {
|
||
|
|
sess *tdsSession
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *MssqlConn) Commit() error {
|
||
|
|
headers := []headerStruct{
|
||
|
|
{hdrtype: dataStmHdrTransDescr,
|
||
|
|
data: transDescrHdr{c.sess.tranid, 1}.pack()},
|
||
|
|
}
|
||
|
|
if err := sendCommitXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
tokchan := make(chan tokenStruct, 5)
|
||
|
|
go processResponse(c.sess, tokchan)
|
||
|
|
for tok := range tokchan {
|
||
|
|
switch token := tok.(type) {
|
||
|
|
case error:
|
||
|
|
return token
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *MssqlConn) Rollback() error {
|
||
|
|
headers := []headerStruct{
|
||
|
|
{hdrtype: dataStmHdrTransDescr,
|
||
|
|
data: transDescrHdr{c.sess.tranid, 1}.pack()},
|
||
|
|
}
|
||
|
|
if err := sendRollbackXact(c.sess.buf, headers, "", 0, 0, ""); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
tokchan := make(chan tokenStruct, 5)
|
||
|
|
go processResponse(c.sess, tokchan)
|
||
|
|
for tok := range tokchan {
|
||
|
|
switch token := tok.(type) {
|
||
|
|
case error:
|
||
|
|
return token
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *MssqlConn) Begin() (driver.Tx, error) {
|
||
|
|
headers := []headerStruct{
|
||
|
|
{hdrtype: dataStmHdrTransDescr,
|
||
|
|
data: transDescrHdr{0, 1}.pack()},
|
||
|
|
}
|
||
|
|
if err := sendBeginXact(c.sess.buf, headers, 0, ""); err != nil {
|
||
|
|
return nil, CheckBadConn(err)
|
||
|
|
}
|
||
|
|
tokchan := make(chan tokenStruct, 5)
|
||
|
|
go processResponse(c.sess, tokchan)
|
||
|
|
for tok := range tokchan {
|
||
|
|
switch token := tok.(type) {
|
||
|
|
case error:
|
||
|
|
if c.sess.tranid != 0 {
|
||
|
|
return nil, token
|
||
|
|
}
|
||
|
|
return nil, CheckBadConn(token)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// successful BEGINXACT request will return sess.tranid
|
||
|
|
// for started transaction
|
||
|
|
return c, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func parseConnectionString(dsn string) (res map[string]string) {
|
||
|
|
res = map[string]string{}
|
||
|
|
parts := strings.Split(dsn, ";")
|
||
|
|
for _, part := range parts {
|
||
|
|
if len(part) == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
lst := strings.SplitN(part, "=", 2)
|
||
|
|
name := strings.TrimSpace(strings.ToLower(lst[0]))
|
||
|
|
if len(name) == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var value string = ""
|
||
|
|
if len(lst) > 1 {
|
||
|
|
value = strings.TrimSpace(lst[1])
|
||
|
|
}
|
||
|
|
res[name] = value
|
||
|
|
}
|
||
|
|
return res
|
||
|
|
}
|
||
|
|
|
||
|
|
func (d *MssqlDriver) Open(dsn string) (driver.Conn, error) {
|
||
|
|
params := parseConnectionString(dsn)
|
||
|
|
|
||
|
|
conn, err := openConnection(dsn, params)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
conn.sess.log = (*Logger)(d.log)
|
||
|
|
return conn, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func openConnection(dsn string, params map[string]string) (*MssqlConn, error) {
|
||
|
|
sess, err := connect(params)
|
||
|
|
if err != nil {
|
||
|
|
partner := partnersCache.Get(dsn)
|
||
|
|
if partner == "" {
|
||
|
|
partner = params["failoverpartner"]
|
||
|
|
// remove the failoverpartner entry to prevent infinite recursion
|
||
|
|
delete(params, "failoverpartner")
|
||
|
|
if port, ok := params["failoverport"]; ok {
|
||
|
|
params["port"] = port
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if partner != "" {
|
||
|
|
params["server"] = partner
|
||
|
|
return openConnection(dsn, params)
|
||
|
|
}
|
||
|
|
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
if partner := sess.partner; partner != "" {
|
||
|
|
// append an instance so the port will be ignored when this value is used;
|
||
|
|
// tds does not provide the port number.
|
||
|
|
if !strings.Contains(partner, `\`) {
|
||
|
|
partner += `\.`
|
||
|
|
}
|
||
|
|
partnersCache.Set(dsn, partner)
|
||
|
|
}
|
||
|
|
|
||
|
|
return &MssqlConn{sess}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *MssqlConn) Close() error {
|
||
|
|
return c.sess.buf.transport.Close()
|
||
|
|
}
|
||
|
|
|
||
|
|
type MssqlStmt struct {
|
||
|
|
c *MssqlConn
|
||
|
|
query string
|
||
|
|
paramCount int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *MssqlConn) Prepare(query string) (driver.Stmt, error) {
|
||
|
|
q, paramCount := parseParams(query)
|
||
|
|
return &MssqlStmt{c, q, paramCount}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) Close() error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) NumInput() int {
|
||
|
|
return s.paramCount
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) sendQuery(args []driver.Value) (err error) {
|
||
|
|
headers := []headerStruct{
|
||
|
|
{hdrtype: dataStmHdrTransDescr,
|
||
|
|
data: transDescrHdr{s.c.sess.tranid, 1}.pack()},
|
||
|
|
}
|
||
|
|
if len(args) != s.paramCount {
|
||
|
|
return errors.New(fmt.Sprintf("sql: expected %d parameters, got %d", s.paramCount, len(args)))
|
||
|
|
}
|
||
|
|
if s.c.sess.logFlags&logSQL != 0 {
|
||
|
|
s.c.sess.log.Println(s.query)
|
||
|
|
}
|
||
|
|
if s.c.sess.logFlags&logParams != 0 && len(args) > 0 {
|
||
|
|
for i := 0; i < len(args); i++ {
|
||
|
|
s.c.sess.log.Printf("\t@p%d\t%v\n", i+1, args[i])
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
if len(args) == 0 {
|
||
|
|
if err = sendSqlBatch72(s.c.sess.buf, s.query, headers); err != nil {
|
||
|
|
if s.c.sess.tranid != 0 {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return CheckBadConn(err)
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
params := make([]Param, len(args)+2)
|
||
|
|
decls := make([]string, len(args))
|
||
|
|
params[0], err = s.makeParam(s.query)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
for i, val := range args {
|
||
|
|
params[i+2], err = s.makeParam(val)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
name := fmt.Sprintf("@p%d", i+1)
|
||
|
|
params[i+2].Name = name
|
||
|
|
decls[i] = fmt.Sprintf("%s %s", name, makeDecl(params[i+2].ti))
|
||
|
|
}
|
||
|
|
params[1], err = s.makeParam(strings.Join(decls, ","))
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if err = sendRpc(s.c.sess.buf, headers, Sp_ExecuteSql, 0, params); err != nil {
|
||
|
|
if s.c.sess.tranid != 0 {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return CheckBadConn(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) Query(args []driver.Value) (res driver.Rows, err error) {
|
||
|
|
if err = s.sendQuery(args); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
tokchan := make(chan tokenStruct, 5)
|
||
|
|
go processResponse(s.c.sess, tokchan)
|
||
|
|
// process metadata
|
||
|
|
var cols []string
|
||
|
|
loop:
|
||
|
|
for tok := range tokchan {
|
||
|
|
switch token := tok.(type) {
|
||
|
|
case doneStruct:
|
||
|
|
break loop
|
||
|
|
case []columnStruct:
|
||
|
|
cols = make([]string, len(token))
|
||
|
|
for i, col := range token {
|
||
|
|
cols[i] = col.ColName
|
||
|
|
}
|
||
|
|
break loop
|
||
|
|
case error:
|
||
|
|
if s.c.sess.tranid != 0 {
|
||
|
|
return nil, token
|
||
|
|
}
|
||
|
|
return nil, CheckBadConn(token)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return &MssqlRows{sess: s.c.sess, tokchan: tokchan, cols: cols}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) Exec(args []driver.Value) (res driver.Result, err error) {
|
||
|
|
if err = s.sendQuery(args); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
tokchan := make(chan tokenStruct, 5)
|
||
|
|
go processResponse(s.c.sess, tokchan)
|
||
|
|
var rowCount int64
|
||
|
|
for token := range tokchan {
|
||
|
|
switch token := token.(type) {
|
||
|
|
case doneInProcStruct:
|
||
|
|
if token.Status&doneCount != 0 {
|
||
|
|
rowCount = int64(token.RowCount)
|
||
|
|
}
|
||
|
|
case doneStruct:
|
||
|
|
if token.Status&doneCount != 0 {
|
||
|
|
rowCount = int64(token.RowCount)
|
||
|
|
}
|
||
|
|
case error:
|
||
|
|
if s.c.sess.logFlags&logErrors != 0 {
|
||
|
|
s.c.sess.log.Println("got error:", token)
|
||
|
|
}
|
||
|
|
if s.c.sess.tranid != 0 {
|
||
|
|
return nil, token
|
||
|
|
}
|
||
|
|
return nil, CheckBadConn(token)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return &MssqlResult{s.c, rowCount}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type MssqlRows struct {
|
||
|
|
sess *tdsSession
|
||
|
|
cols []string
|
||
|
|
tokchan chan tokenStruct
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rc *MssqlRows) Close() error {
|
||
|
|
for _ = range rc.tokchan {
|
||
|
|
}
|
||
|
|
rc.tokchan = nil
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rc *MssqlRows) Columns() (res []string) {
|
||
|
|
return rc.cols
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rc *MssqlRows) Next(dest []driver.Value) (err error) {
|
||
|
|
for tok := range rc.tokchan {
|
||
|
|
switch tokdata := tok.(type) {
|
||
|
|
case []columnStruct:
|
||
|
|
return streamErrorf("Unexpected token COLMETADATA")
|
||
|
|
case []interface{}:
|
||
|
|
for i := range dest {
|
||
|
|
dest[i] = tokdata[i]
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
case error:
|
||
|
|
return tokdata
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return io.EOF
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *MssqlStmt) makeParam(val driver.Value) (res Param, err error) {
|
||
|
|
if val == nil {
|
||
|
|
res.ti.TypeId = typeNVarChar
|
||
|
|
res.buffer = nil
|
||
|
|
res.ti.Size = 2
|
||
|
|
return
|
||
|
|
}
|
||
|
|
switch val := val.(type) {
|
||
|
|
case int64:
|
||
|
|
res.ti.TypeId = typeIntN
|
||
|
|
res.buffer = make([]byte, 8)
|
||
|
|
res.ti.Size = 8
|
||
|
|
binary.LittleEndian.PutUint64(res.buffer, uint64(val))
|
||
|
|
case float64:
|
||
|
|
res.ti.TypeId = typeFltN
|
||
|
|
res.ti.Size = 8
|
||
|
|
res.buffer = make([]byte, 8)
|
||
|
|
binary.LittleEndian.PutUint64(res.buffer, math.Float64bits(val))
|
||
|
|
case []byte:
|
||
|
|
res.ti.TypeId = typeBigVarBin
|
||
|
|
res.ti.Size = len(val)
|
||
|
|
res.buffer = val
|
||
|
|
case string:
|
||
|
|
res.ti.TypeId = typeNVarChar
|
||
|
|
res.buffer = str2ucs2(val)
|
||
|
|
res.ti.Size = len(res.buffer)
|
||
|
|
case bool:
|
||
|
|
res.ti.TypeId = typeBitN
|
||
|
|
res.ti.Size = 1
|
||
|
|
res.buffer = make([]byte, 1)
|
||
|
|
if val {
|
||
|
|
res.buffer[0] = 1
|
||
|
|
}
|
||
|
|
case time.Time:
|
||
|
|
if s.c.sess.loginAck.TDSVersion >= verTDS73 {
|
||
|
|
res.ti.TypeId = typeDateTimeOffsetN
|
||
|
|
res.ti.Scale = 7
|
||
|
|
res.ti.Size = 10
|
||
|
|
buf := make([]byte, 10)
|
||
|
|
res.buffer = buf
|
||
|
|
days, ns := dateTime2(val)
|
||
|
|
ns /= 100
|
||
|
|
buf[0] = byte(ns)
|
||
|
|
buf[1] = byte(ns >> 8)
|
||
|
|
buf[2] = byte(ns >> 16)
|
||
|
|
buf[3] = byte(ns >> 24)
|
||
|
|
buf[4] = byte(ns >> 32)
|
||
|
|
buf[5] = byte(days)
|
||
|
|
buf[6] = byte(days >> 8)
|
||
|
|
buf[7] = byte(days >> 16)
|
||
|
|
_, offset := val.Zone()
|
||
|
|
offset /= 60
|
||
|
|
buf[8] = byte(offset)
|
||
|
|
buf[9] = byte(offset >> 8)
|
||
|
|
} else {
|
||
|
|
res.ti.TypeId = typeDateTimeN
|
||
|
|
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 := dur / (24 * time.Hour)
|
||
|
|
tm := (300 * (dur % (24 * time.Hour))) / time.Second
|
||
|
|
binary.LittleEndian.PutUint32(res.buffer[0:4], uint32(days))
|
||
|
|
binary.LittleEndian.PutUint32(res.buffer[4:8], uint32(tm))
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
err = fmt.Errorf("mssql: unknown type for %T", val)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
type MssqlResult struct {
|
||
|
|
c *MssqlConn
|
||
|
|
rowsAffected int64
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *MssqlResult) RowsAffected() (int64, error) {
|
||
|
|
return r.rowsAffected, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *MssqlResult) LastInsertId() (int64, error) {
|
||
|
|
s, err := r.c.Prepare("select cast(@@identity as bigint)")
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
defer s.Close()
|
||
|
|
rows, err := s.Query(nil)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
defer rows.Close()
|
||
|
|
dest := make([]driver.Value, 1)
|
||
|
|
err = rows.Next(dest)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
if dest[0] == nil {
|
||
|
|
return -1, errors.New("There is no generated identity value")
|
||
|
|
}
|
||
|
|
lastInsertId := dest[0].(int64)
|
||
|
|
return lastInsertId, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type partners struct {
|
||
|
|
mu sync.RWMutex
|
||
|
|
v map[string]string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *partners) Set(key, value string) error {
|
||
|
|
p.mu.Lock()
|
||
|
|
defer p.mu.Unlock()
|
||
|
|
if _, ok := p.v[key]; ok {
|
||
|
|
return errors.New("key already exists")
|
||
|
|
}
|
||
|
|
|
||
|
|
p.v[key] = value
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *partners) Get(key string) (value string) {
|
||
|
|
p.mu.RLock()
|
||
|
|
value = p.v[key]
|
||
|
|
p.mu.RUnlock()
|
||
|
|
return
|
||
|
|
}
|