cloudflared-mirror/vendor/github.com/kshvakov/clickhouse/clickhouse.go

321 lines
7.3 KiB
Go
Raw Normal View History

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() {}
}