go-clickhouse/ch/chproto/reader.go

package chproto

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

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

type reader interface {
	io.Reader
	io.ByteReader
	Buffered() int
}

type Reader struct {
	br *bufio.Reader
	zr *lz4Reader
	rd reader // points to br or zr

	buf []byte
}

func NewReader(r io.Reader) *Reader {
	br := bufio.NewReader(r)
	return &Reader{
		br: br,
		zr: newLZ4Reader(br),
		rd: br,

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

func (r *Reader) WithCompression(enabled bool, fn func() error) error {
	if enabled {
		r.rd = r.zr
	}

	firstErr := fn()

	if enabled {
		r.rd = r.br
		if err := r.zr.Release(); err != nil && firstErr == nil {
			firstErr = err
		}
	}

	return firstErr
}

func (r *Reader) Read(buf []byte) (int, error) {
	return r.rd.Read(buf)
}

func (r *Reader) Buffered() int {
	return r.rd.Buffered()
}

func (r *Reader) Bool() (bool, error) {
	c, err := r.rd.ReadByte()
	if err != nil {
		return false, err
	}
	return c == 1, nil
}

func (r *Reader) Uvarint() (uint64, error) {
	return binary.ReadUvarint(r.rd)
}

func (r *Reader) UInt8() (uint8, error) {
	c, err := r.rd.ReadByte()
	if err != nil {
		return 0, err
	}
	return c, nil
}

func (r *Reader) UInt16() (uint16, error) {
	b, err := r.readNTemp(2)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint16(b), nil
}

func (r *Reader) UInt32() (uint32, error) {
	b, err := r.readNTemp(4)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint32(b), nil
}

func (r *Reader) UInt64() (uint64, error) {
	b, err := r.readNTemp(8)
	if err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint64(b), nil
}

func (r *Reader) Int8() (int8, error) {
	num, err := r.UInt8()
	return int8(num), err
}

func (r *Reader) Int16() (int16, error) {
	num, err := r.UInt16()
	return int16(num), err
}

func (r *Reader) Int32() (int32, error) {
	num, err := r.UInt32()
	return int32(num), err
}

func (r *Reader) Int64() (int64, error) {
	num, err := r.UInt64()
	return int64(num), err
}

func (r *Reader) Float32() (float32, error) {
	num, err := r.UInt32()
	if err != nil {
		return 0, err
	}
	return math.Float32frombits(num), nil
}

func (r *Reader) Float64() (float64, error) {
	num, err := r.UInt64()
	if err != nil {
		return 0, err
	}
	return math.Float64frombits(num), nil
}

func (r *Reader) Bytes() ([]byte, error) {
	num, err := r.Uvarint()
	if err != nil {
		return nil, err
	}

	b := make([]byte, int(num))
	_, err = io.ReadFull(r.rd, b)
	if err != nil {
		return nil, err
	}
	return b, nil
}

func (r *Reader) String() (string, error) {
	b, err := r.Bytes()
	if err != nil {
		return "", err
	}
	return internal.String(b), nil
}

func (r *Reader) UUID(b []byte) error {
	if len(b) != uuidLen {
		return fmt.Errorf("got %d bytes, wanted %d", len(b), uuidLen)
	}

	_, err := io.ReadFull(r.rd, b)
	if err != nil {
		return err
	}
	packUUID(b)
	return nil
}

func (r *Reader) readNTemp(n int) ([]byte, error) {
	buf := r.buf[:n]
	_, err := io.ReadFull(r.rd, buf)
	return buf, err
}

func (r *Reader) DateTime() (time.Time, error) {
	sec, err := r.UInt32()
	if err != nil {
		return time.Time{}, err
	}
	if sec == 0 {
		return time.Time{}, nil
	}
	return time.Unix(int64(sec), 0), nil
}

func (r *Reader) Date() (time.Time, error) {
	days, err := r.UInt16()
	if err != nil {
		return time.Time{}, err
	}
	if days == 0 {
		return time.Time{}, nil
	}
	return time.Unix(int64(days)*secsInDay, 0), nil
}