go-clickhouse/ch/chproto/writer.go

package chproto

import (
	"bufio"
	"encoding/binary"
	"io"
	"math"
	"time"

	"github.com/uptrace/go-clickhouse/ch/internal"
)

const (
	uuidLen   = 16
	secsInDay = 24 * 3600
)

type writer interface {
	io.Writer
	io.ByteWriter
	Flush() error
}

type Writer struct {
	bw *bufio.Writer
	zw *lz4Writer
	wr writer // points to bw or zw

	err error

	buf []byte
}

func NewWriter(w io.Writer) *Writer {
	bw := bufio.NewWriter(w)
	return &Writer{
		bw: bw,
		zw: newLZ4Writer(bw),
		wr: bw,

		buf: make([]byte, uuidLen),
	}
}

func (w *Writer) WithCompression(enabled bool, fn func() error) {
	if w.err != nil {
		return
	}

	if enabled {
		w.wr = w.zw
	}

	w.err = fn()

	if enabled {
		if err := w.zw.Close(); err != nil && w.err == nil {
			w.err = err
		}
		w.wr = w.bw
	}
}

func (w *Writer) Flush() (err error) {
	if w.err != nil {
		err, w.err = w.err, nil
		return err
	}

	return w.wr.Flush()
}

func (w *Writer) Write(b []byte) {
	if w.err != nil {
		return
	}

	_, err := w.wr.Write(b)
	w.err = err
}

func (w *Writer) WriteByte(c byte) {
	if w.err != nil {
		return
	}
	w.err = w.wr.WriteByte(c)
}

func (w *Writer) Bool(flag bool) {
	var num uint8
	if flag {
		num = 1
	}
	w.UInt8(num)
}

func (w *Writer) Uvarint(num uint64) {
	n := binary.PutUvarint(w.buf, num)
	w.Write(w.buf[:n])
}

func (w *Writer) UInt8(num uint8) {
	w.WriteByte(num)
}

func (w *Writer) UInt16(num uint16) {
	binary.LittleEndian.PutUint16(w.buf, num)
	w.Write(w.buf[:2])
}

func (w *Writer) UInt32(num uint32) {
	binary.LittleEndian.PutUint32(w.buf, num)
	w.Write(w.buf[:4])
}

func (w *Writer) UInt64(num uint64) {
	binary.LittleEndian.PutUint64(w.buf, num)
	w.Write(w.buf[:8])
}

func (w *Writer) Int8(num int8) {
	w.UInt8(uint8(num))
}

func (w *Writer) Int16(num int16) {
	w.UInt16(uint16(num))
}

func (w *Writer) Int32(num int32) {
	w.UInt32(uint32(num))
}

func (w *Writer) Int64(num int64) {
	w.UInt64(uint64(num))
}

func (w *Writer) Float32(num float32) {
	w.UInt32(math.Float32bits(num))
}

func (w *Writer) Float64(num float64) {
	w.UInt64(math.Float64bits(num))
}

func (w *Writer) String(s string) {
	w.Uvarint(uint64(len(s)))
	w.Write(internal.Bytes(s))
}

func (w *Writer) Bytes(b []byte) {
	w.Uvarint(uint64(len(b)))
	w.Write(b)
}

func (w *Writer) UUID(b []byte) {
	if len(b) != uuidLen {
		panic("not reached")
	}

	buf := w.buf[:uuidLen]
	copy(buf, b)
	packUUID(buf)
	w.Write(buf)
}

// 2 int64 in little endian order?
func packUUID(b []byte) []byte {
	_ = b[15]
	b[0], b[7] = b[7], b[0]
	b[1], b[6] = b[6], b[1]
	b[2], b[5] = b[5], b[2]
	b[3], b[4] = b[4], b[3]
	b[8], b[15] = b[15], b[8]
	b[9], b[14] = b[14], b[9]
	b[10], b[13] = b[13], b[10]
	b[11], b[12] = b[12], b[11]
	return b
}

func (w *Writer) DateTime(tm time.Time) {
	w.UInt32(uint32(unixTime(tm)))
}

func (w *Writer) Date(tm time.Time) {
	w.UInt16(uint16(unixTime(tm) / secsInDay))
}

func unixTime(tm time.Time) int64 {
	if tm.IsZero() {
		return 0
	}
	return tm.Unix()
}