321 lines
7.3 KiB
Go
321 lines
7.3 KiB
Go
package clickhouse
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"reflect"
|
|
"regexp"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/kshvakov/clickhouse/lib/binary"
|
|
"github.com/kshvakov/clickhouse/lib/column"
|
|
"github.com/kshvakov/clickhouse/lib/data"
|
|
"github.com/kshvakov/clickhouse/lib/protocol"
|
|
"github.com/kshvakov/clickhouse/lib/types"
|
|
)
|
|
|
|
type (
|
|
Date = types.Date
|
|
DateTime = types.DateTime
|
|
UUID = types.UUID
|
|
)
|
|
|
|
var (
|
|
ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
|
|
ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
|
|
)
|
|
|
|
var (
|
|
splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
|
|
)
|
|
|
|
type logger func(format string, v ...interface{})
|
|
|
|
type clickhouse struct {
|
|
sync.Mutex
|
|
data.ServerInfo
|
|
data.ClientInfo
|
|
logf logger
|
|
conn *connect
|
|
block *data.Block
|
|
buffer *bufio.Writer
|
|
decoder *binary.Decoder
|
|
encoder *binary.Encoder
|
|
settings *querySettings
|
|
compress bool
|
|
blockSize int
|
|
inTransaction bool
|
|
}
|
|
|
|
func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
|
|
return ch.prepareContext(context.Background(), query)
|
|
}
|
|
|
|
func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
|
|
return ch.prepareContext(ctx, query)
|
|
}
|
|
|
|
func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
|
|
ch.logf("[prepare] %s", query)
|
|
switch {
|
|
case ch.conn.closed:
|
|
return nil, driver.ErrBadConn
|
|
case ch.block != nil:
|
|
return nil, ErrLimitDataRequestInTx
|
|
case isInsert(query):
|
|
if !ch.inTransaction {
|
|
return nil, ErrInsertInNotBatchMode
|
|
}
|
|
return ch.insert(query)
|
|
}
|
|
return &stmt{
|
|
ch: ch,
|
|
query: query,
|
|
numInput: numInput(query),
|
|
}, nil
|
|
}
|
|
|
|
func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
|
|
if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
|
|
return nil, err
|
|
}
|
|
if ch.block, err = ch.readMeta(); err != nil {
|
|
return nil, err
|
|
}
|
|
return &stmt{
|
|
ch: ch,
|
|
isInsert: true,
|
|
}, nil
|
|
}
|
|
|
|
func (ch *clickhouse) Begin() (driver.Tx, error) {
|
|
return ch.beginTx(context.Background(), txOptions{})
|
|
}
|
|
|
|
func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
|
|
return ch.beginTx(ctx, txOptions{
|
|
Isolation: int(opts.Isolation),
|
|
ReadOnly: opts.ReadOnly,
|
|
})
|
|
}
|
|
|
|
type txOptions struct {
|
|
Isolation int
|
|
ReadOnly bool
|
|
}
|
|
|
|
func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
|
|
ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
|
|
switch {
|
|
case ch.inTransaction:
|
|
return nil, sql.ErrTxDone
|
|
case ch.conn.closed:
|
|
return nil, driver.ErrBadConn
|
|
}
|
|
if finish := ch.watchCancel(ctx); finish != nil {
|
|
defer finish()
|
|
}
|
|
ch.block = nil
|
|
ch.inTransaction = true
|
|
return ch, nil
|
|
}
|
|
|
|
func (ch *clickhouse) Commit() error {
|
|
ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
|
|
defer func() {
|
|
if ch.block != nil {
|
|
ch.block.Reset()
|
|
ch.block = nil
|
|
}
|
|
ch.inTransaction = false
|
|
}()
|
|
switch {
|
|
case !ch.inTransaction:
|
|
return sql.ErrTxDone
|
|
case ch.conn.closed:
|
|
return driver.ErrBadConn
|
|
}
|
|
if ch.block != nil {
|
|
if err := ch.writeBlock(ch.block); err != nil {
|
|
return err
|
|
}
|
|
// Send empty block as marker of end of data.
|
|
if err := ch.writeBlock(&data.Block{}); err != nil {
|
|
return err
|
|
}
|
|
if err := ch.encoder.Flush(); err != nil {
|
|
return err
|
|
}
|
|
return ch.process()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ch *clickhouse) Rollback() error {
|
|
ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
|
|
if !ch.inTransaction {
|
|
return sql.ErrTxDone
|
|
}
|
|
if ch.block != nil {
|
|
ch.block.Reset()
|
|
}
|
|
ch.block = nil
|
|
ch.buffer = nil
|
|
ch.inTransaction = false
|
|
return ch.conn.Close()
|
|
}
|
|
|
|
func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
|
|
switch nv.Value.(type) {
|
|
case column.IP, column.UUID:
|
|
return nil
|
|
case nil, []byte, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, string, time.Time:
|
|
return nil
|
|
}
|
|
switch v := nv.Value.(type) {
|
|
case
|
|
[]int, []int8, []int16, []int32, []int64,
|
|
[]uint, []uint8, []uint16, []uint32, []uint64,
|
|
[]float32, []float64,
|
|
[]string:
|
|
return nil
|
|
case net.IP:
|
|
return nil
|
|
case driver.Valuer:
|
|
value, err := v.Value()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
nv.Value = value
|
|
default:
|
|
switch value := reflect.ValueOf(nv.Value); value.Kind() {
|
|
case reflect.Slice:
|
|
return nil
|
|
case reflect.Bool:
|
|
nv.Value = uint8(0)
|
|
if value.Bool() {
|
|
nv.Value = uint8(1)
|
|
}
|
|
case reflect.Int8:
|
|
nv.Value = int8(value.Int())
|
|
case reflect.Int16:
|
|
nv.Value = int16(value.Int())
|
|
case reflect.Int32:
|
|
nv.Value = int32(value.Int())
|
|
case reflect.Int64:
|
|
nv.Value = value.Int()
|
|
case reflect.Uint8:
|
|
nv.Value = uint8(value.Uint())
|
|
case reflect.Uint16:
|
|
nv.Value = uint16(value.Uint())
|
|
case reflect.Uint32:
|
|
nv.Value = uint32(value.Uint())
|
|
case reflect.Uint64:
|
|
nv.Value = uint64(value.Uint())
|
|
case reflect.Float32:
|
|
nv.Value = float32(value.Float())
|
|
case reflect.Float64:
|
|
nv.Value = float64(value.Float())
|
|
case reflect.String:
|
|
nv.Value = value.String()
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ch *clickhouse) Close() error {
|
|
ch.block = nil
|
|
return ch.conn.Close()
|
|
}
|
|
|
|
func (ch *clickhouse) process() error {
|
|
packet, err := ch.decoder.Uvarint()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for {
|
|
switch packet {
|
|
case protocol.ServerPong:
|
|
ch.logf("[process] <- pong")
|
|
return nil
|
|
case protocol.ServerException:
|
|
ch.logf("[process] <- exception")
|
|
return ch.exception()
|
|
case protocol.ServerProgress:
|
|
progress, err := ch.progress()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
|
|
progress.rows,
|
|
progress.bytes,
|
|
progress.totalRows,
|
|
)
|
|
case protocol.ServerProfileInfo:
|
|
profileInfo, err := ch.profileInfo()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
|
|
case protocol.ServerData:
|
|
block, err := ch.readBlock()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
|
|
case protocol.ServerEndOfStream:
|
|
ch.logf("[process] <- end of stream")
|
|
return nil
|
|
default:
|
|
ch.conn.Close()
|
|
return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
|
|
}
|
|
if packet, err = ch.decoder.Uvarint(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ch *clickhouse) cancel() error {
|
|
ch.logf("[cancel request]")
|
|
// even if we fail to write the cancel, we still need to close
|
|
err := ch.encoder.Uvarint(protocol.ClientCancel)
|
|
if err == nil {
|
|
err = ch.encoder.Flush()
|
|
}
|
|
// return the close error if there was one, otherwise return the write error
|
|
if cerr := ch.conn.Close(); cerr != nil {
|
|
return cerr
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (ch *clickhouse) watchCancel(ctx context.Context) func() {
|
|
if done := ctx.Done(); done != nil {
|
|
finished := make(chan struct{})
|
|
go func() {
|
|
select {
|
|
case <-done:
|
|
ch.cancel()
|
|
finished <- struct{}{}
|
|
ch.logf("[cancel] <- done")
|
|
case <-finished:
|
|
ch.logf("[cancel] <- finished")
|
|
}
|
|
}()
|
|
return func() {
|
|
select {
|
|
case <-finished:
|
|
case finished <- struct{}{}:
|
|
}
|
|
}
|
|
}
|
|
return func() {}
|
|
}
|