mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 10:54:59 +00:00
716 lines
15 KiB
Go
716 lines
15 KiB
Go
package wireprotocol
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/binary"
|
|
"io"
|
|
"math"
|
|
"time"
|
|
|
|
"github.com/featurebasedb/featurebase/v3/errors"
|
|
"github.com/featurebasedb/featurebase/v3/pql"
|
|
"github.com/featurebasedb/featurebase/v3/sql3/parser"
|
|
"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
|
|
)
|
|
|
|
const (
|
|
// server --> client
|
|
TOKEN_SCHEMA_INFO int16 = 0xA1
|
|
TOKEN_ROW int16 = 0xA2
|
|
TOKEN_DONE int16 = 0xFD
|
|
TOKEN_INFO_MESSAGE int16 = 0xFE
|
|
TOKEN_ERROR_MESSAGE int16 = 0xFF
|
|
|
|
// client --> server
|
|
TOKEN_SQL int16 = 0x01
|
|
TOKEN_PLAN_OP int16 = 0x02
|
|
)
|
|
|
|
const (
|
|
TYPE_VOID int8 = 0x00
|
|
TYPE_ID int8 = 0x01
|
|
TYPE_BOOL int8 = 0x02
|
|
TYPE_INT int8 = 0x03
|
|
TYPE_DECIMAL int8 = 0x04
|
|
TYPE_TIMESTAMP int8 = 0x05
|
|
TYPE_IDSET int8 = 0x06
|
|
TYPE_STRING int8 = 0x07
|
|
TYPE_STRINGSET int8 = 0x08
|
|
TYPE_VARCHAR int8 = 0x09
|
|
TYPE_VARBINARY int8 = 0x0A
|
|
TYPE_VECTOR int8 = 0x0B
|
|
)
|
|
|
|
func ExpectToken(reader io.Reader, token int16) (int16, error) {
|
|
var tk int16
|
|
err := binary.Read(reader, binary.BigEndian, &tk)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if tk != token {
|
|
return 0, errors.Errorf("expected token found %d", token)
|
|
}
|
|
return tk, nil
|
|
}
|
|
|
|
func ReadToken(reader io.Reader) (int16, error) {
|
|
var tk int16
|
|
err := binary.Read(reader, binary.BigEndian, &tk)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return tk, nil
|
|
}
|
|
|
|
// TOKEN_COLUMN_INFO message
|
|
// length (bytes)
|
|
// token 2
|
|
// column count 2
|
|
//
|
|
// (n) columns
|
|
//
|
|
// name length 1
|
|
// name (from prev)
|
|
// data type 1
|
|
//
|
|
// (optional)
|
|
// if type decimal
|
|
// scale 1
|
|
|
|
// note RelationName and AliasName members from
|
|
// PlannerColumn are not sent over the wire
|
|
func WriteSchema(schema types.Schema) ([]byte, error) {
|
|
buf := new(bytes.Buffer)
|
|
writer := bufio.NewWriter(buf)
|
|
// write token
|
|
writeToken(writer, TOKEN_SCHEMA_INFO)
|
|
|
|
// column count
|
|
writeInt16(writer, int16(len(schema)))
|
|
|
|
// for each column
|
|
for _, s := range schema {
|
|
// name len byte
|
|
writeInt8(writer, int8(len(s.ColumnName)))
|
|
// name
|
|
writer.WriteString(s.ColumnName)
|
|
// type byte
|
|
switch ty := s.Type.(type) {
|
|
case *parser.DataTypeID:
|
|
writeInt8(writer, TYPE_ID)
|
|
|
|
case *parser.DataTypeBool:
|
|
writeInt8(writer, TYPE_BOOL)
|
|
|
|
case *parser.DataTypeInt:
|
|
writeInt8(writer, TYPE_INT)
|
|
|
|
case *parser.DataTypeDecimal:
|
|
writeInt8(writer, TYPE_DECIMAL)
|
|
writeInt8(writer, int8(ty.Scale))
|
|
|
|
case *parser.DataTypeTimestamp:
|
|
writeInt8(writer, TYPE_TIMESTAMP)
|
|
|
|
case *parser.DataTypeIDSet:
|
|
writeInt8(writer, TYPE_IDSET)
|
|
|
|
case *parser.DataTypeString:
|
|
writeInt8(writer, TYPE_STRING)
|
|
|
|
case *parser.DataTypeStringSet:
|
|
writeInt8(writer, TYPE_STRINGSET)
|
|
|
|
case *parser.DataTypeVarchar:
|
|
writeInt8(writer, TYPE_VARCHAR)
|
|
writeInt32(writer, int32(ty.Length))
|
|
|
|
case *parser.DataTypeVector:
|
|
writeInt8(writer, TYPE_VECTOR)
|
|
writeInt32(writer, int32(ty.Length))
|
|
|
|
case *parser.DataTypeVarbinary:
|
|
writeInt8(writer, TYPE_VARBINARY)
|
|
writeInt32(writer, int32(ty.Length))
|
|
|
|
default:
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", s.Type)
|
|
}
|
|
}
|
|
writer.Flush()
|
|
return buf.Bytes(), nil
|
|
}
|
|
|
|
// ReadSchema consumes a schema object from a reader
|
|
// Note you cannot do WriteSchema to a buffer then Read that schame
|
|
// back without consuming or skipping the token.
|
|
// so you can't just to WriteSchema -> ReadSchema it must be WriteSchema->ConsumeToken->ReadSchema
|
|
func ReadSchema(reader io.Reader) (types.Schema, error) {
|
|
var columnCount int16
|
|
err := binary.Read(reader, binary.BigEndian, &columnCount)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var schema types.Schema
|
|
for i := 0; i < int(columnCount); i++ {
|
|
|
|
var nameLen int8
|
|
err = binary.Read(reader, binary.BigEndian, &nameLen)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bname := make([]byte, nameLen)
|
|
err = binary.Read(reader, binary.BigEndian, &bname)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
colName := string(bname)
|
|
|
|
var typ int8
|
|
err = binary.Read(reader, binary.BigEndian, &typ)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var dataType parser.ExprDataType
|
|
|
|
switch typ {
|
|
case TYPE_ID:
|
|
dataType = parser.NewDataTypeID()
|
|
|
|
case TYPE_BOOL:
|
|
dataType = parser.NewDataTypeBool()
|
|
|
|
case TYPE_INT:
|
|
dataType = parser.NewDataTypeInt()
|
|
|
|
case TYPE_DECIMAL:
|
|
var scale int8
|
|
err = binary.Read(reader, binary.BigEndian, &scale)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dataType = parser.NewDataTypeDecimal(int64(scale))
|
|
|
|
case TYPE_TIMESTAMP:
|
|
dataType = parser.NewDataTypeTimestamp()
|
|
|
|
case TYPE_IDSET:
|
|
dataType = parser.NewDataTypeIDSet()
|
|
|
|
case TYPE_STRING:
|
|
dataType = parser.NewDataTypeString()
|
|
|
|
case TYPE_STRINGSET:
|
|
dataType = parser.NewDataTypeStringSet()
|
|
|
|
case TYPE_VARCHAR:
|
|
var length int32
|
|
err = binary.Read(reader, binary.BigEndian, &length)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dataType = parser.NewDataTypeVarchar(int64(length))
|
|
|
|
case TYPE_VECTOR:
|
|
var length int32
|
|
err = binary.Read(reader, binary.BigEndian, &length)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dataType = parser.NewDataTypeVector(int64(length))
|
|
|
|
case TYPE_VARBINARY:
|
|
var length int32
|
|
err = binary.Read(reader, binary.BigEndian, &length)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dataType = parser.NewDataTypeVarbinary(int64(length))
|
|
}
|
|
|
|
schema = append(schema, &types.PlannerColumn{
|
|
ColumnName: colName,
|
|
Type: dataType,
|
|
})
|
|
}
|
|
return schema, nil
|
|
}
|
|
|
|
// TOKEN_ROW message
|
|
// length (bytes)
|
|
// token 2
|
|
//
|
|
// (n) columns
|
|
// if column length is 0 --> null
|
|
//
|
|
// - for ID, INT
|
|
//
|
|
// column length 1
|
|
// value 8
|
|
//
|
|
// - for DECIMAL
|
|
//
|
|
// column length 1
|
|
// value 8
|
|
//
|
|
// - for BOOL
|
|
//
|
|
// column length 1
|
|
// value 1
|
|
//
|
|
// - for TIMESTAMP
|
|
//
|
|
// column length 1
|
|
// value 8
|
|
//
|
|
// - for IDSET
|
|
//
|
|
// set length 2
|
|
// (n) values
|
|
// value 8
|
|
//
|
|
// - for STRING
|
|
//
|
|
// column length 2
|
|
// value (from prev)
|
|
//
|
|
// - for STRINGSET
|
|
//
|
|
// set length 2
|
|
// (n) values
|
|
// value len 2
|
|
// value (from prev)
|
|
//
|
|
|
|
func WriteRow(row types.Row, schema types.Schema) ([]byte, error) {
|
|
buf := new(bytes.Buffer)
|
|
writer := bufio.NewWriter(buf)
|
|
// write token
|
|
writeToken(writer, TOKEN_ROW)
|
|
|
|
// for each column
|
|
for i, s := range schema {
|
|
val := row[i]
|
|
switch s.Type.(type) {
|
|
case *parser.DataTypeID, *parser.DataTypeInt:
|
|
if val == nil {
|
|
writeInt8(writer, 0)
|
|
} else {
|
|
writeInt8(writer, 8)
|
|
v, ok := row[i].(int64)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt64(writer, v)
|
|
}
|
|
|
|
case *parser.DataTypeDecimal:
|
|
if val == nil {
|
|
writeInt8(writer, 0)
|
|
} else {
|
|
writeInt8(writer, 8)
|
|
v, ok := row[i].(pql.Decimal)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt64(writer, v.ToInt64(v.Scale))
|
|
}
|
|
|
|
case *parser.DataTypeBool:
|
|
if val == nil {
|
|
writeInt8(writer, 0)
|
|
} else {
|
|
writeInt8(writer, 1)
|
|
v, ok := row[i].(bool)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
if v {
|
|
writeInt8(writer, 1)
|
|
} else {
|
|
writeInt8(writer, 0)
|
|
}
|
|
}
|
|
|
|
case *parser.DataTypeTimestamp:
|
|
if val == nil {
|
|
writeInt8(writer, 0)
|
|
} else {
|
|
writeInt8(writer, 8)
|
|
v, ok := row[i].(time.Time)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt64(writer, v.UnixNano())
|
|
}
|
|
|
|
case *parser.DataTypeIDSet:
|
|
if val == nil {
|
|
writeInt16(writer, 0)
|
|
} else {
|
|
v, ok := row[i].([]int64)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt16(writer, int16(len(v)))
|
|
for _, s := range v {
|
|
writeInt64(writer, s)
|
|
}
|
|
}
|
|
|
|
case *parser.DataTypeString:
|
|
if val == nil {
|
|
writeInt16(writer, 0)
|
|
} else {
|
|
v, ok := row[i].(string)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt16(writer, int16(len(v)))
|
|
writer.WriteString(v)
|
|
}
|
|
|
|
case *parser.DataTypeStringSet:
|
|
if val == nil {
|
|
writeInt16(writer, 0)
|
|
} else {
|
|
v, ok := row[i].([]string)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt16(writer, int16(len(v)))
|
|
for _, s := range v {
|
|
writeInt16(writer, int16(len(s)))
|
|
writer.WriteString(s)
|
|
}
|
|
}
|
|
|
|
case *parser.DataTypeVarchar:
|
|
if val == nil {
|
|
writeInt32(writer, 0)
|
|
} else {
|
|
v, ok := row[i].(string)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt32(writer, int32(len(v)))
|
|
writer.WriteString(v)
|
|
}
|
|
|
|
case *parser.DataTypeVector:
|
|
if val == nil {
|
|
writeInt32(writer, 0)
|
|
} else {
|
|
v, ok := row[i].([]float64)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt32(writer, int32(len(v)))
|
|
for _, f := range v {
|
|
var fbuf [8]byte
|
|
binary.BigEndian.PutUint64(fbuf[:], math.Float64bits(f))
|
|
writer.Write(fbuf[:])
|
|
}
|
|
}
|
|
|
|
case *parser.DataTypeVarbinary:
|
|
if val == nil {
|
|
writeInt32(writer, 0)
|
|
} else {
|
|
v, ok := row[i].([]byte)
|
|
if !ok {
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", row[i])
|
|
}
|
|
writeInt32(writer, int32(len(v)))
|
|
writer.Write(v)
|
|
}
|
|
|
|
default:
|
|
return []byte{}, errors.Errorf("unexpected type '%T'", s.Type)
|
|
}
|
|
}
|
|
writer.Flush()
|
|
return buf.Bytes(), nil
|
|
}
|
|
|
|
func ReadRow(reader io.Reader, schema types.Schema) (types.Row, error) {
|
|
row := make(types.Row, len(schema))
|
|
|
|
for idx, s := range schema {
|
|
switch t := s.Type.(type) {
|
|
case *parser.DataTypeID, *parser.DataTypeInt:
|
|
var len int8
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
var value int64
|
|
err := binary.Read(reader, binary.BigEndian, &value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = value
|
|
}
|
|
|
|
case *parser.DataTypeDecimal:
|
|
var len int8
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
var value int64
|
|
err := binary.Read(reader, binary.BigEndian, &value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = pql.NewDecimal(value, t.Scale)
|
|
}
|
|
|
|
case *parser.DataTypeBool:
|
|
var len int8
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
var value int8
|
|
err := binary.Read(reader, binary.BigEndian, &value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = value == 1
|
|
}
|
|
|
|
case *parser.DataTypeTimestamp:
|
|
var len int8
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
var value int64
|
|
err := binary.Read(reader, binary.BigEndian, &value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = time.Unix(0, value)
|
|
}
|
|
|
|
case *parser.DataTypeIDSet:
|
|
var len int16
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
set := make([]int64, len)
|
|
for j := range set {
|
|
var value int64
|
|
err := binary.Read(reader, binary.BigEndian, &value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
set[j] = value
|
|
}
|
|
row[idx] = set
|
|
}
|
|
|
|
case *parser.DataTypeString:
|
|
var len int16
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
bvalue := make([]byte, len)
|
|
err = binary.Read(reader, binary.BigEndian, &bvalue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = string(bvalue)
|
|
}
|
|
|
|
case *parser.DataTypeStringSet:
|
|
var len int16
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
set := make([]string, len)
|
|
for j := range set {
|
|
var vlen int16
|
|
err = binary.Read(reader, binary.BigEndian, &vlen)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bvalue := make([]byte, vlen)
|
|
err = binary.Read(reader, binary.BigEndian, &bvalue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
set[j] = string(bvalue)
|
|
}
|
|
row[idx] = set
|
|
}
|
|
|
|
case *parser.DataTypeVarchar:
|
|
var len int32
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
bvalue := make([]byte, len)
|
|
err = binary.Read(reader, binary.BigEndian, &bvalue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = string(bvalue)
|
|
}
|
|
|
|
case *parser.DataTypeVector:
|
|
var len int32
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
bvalue := make([]float64, len)
|
|
for j := 0; j < int(len); j++ {
|
|
var fvalue float64
|
|
err = binary.Read(reader, binary.BigEndian, &fvalue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bvalue[j] = fvalue
|
|
}
|
|
|
|
row[idx] = bvalue
|
|
}
|
|
|
|
case *parser.DataTypeVarbinary:
|
|
var len int32
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len == 0 {
|
|
row[idx] = nil
|
|
} else {
|
|
bvalue := make([]byte, len)
|
|
err = binary.Read(reader, binary.BigEndian, &bvalue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
row[idx] = bvalue
|
|
}
|
|
|
|
default:
|
|
return nil, errors.Errorf("unexpected type '%T'", s.Type)
|
|
}
|
|
}
|
|
|
|
return row, nil
|
|
}
|
|
|
|
// TOKEN_DONE message
|
|
// length (bytes)
|
|
// token 2
|
|
|
|
func WriteDone() []byte {
|
|
buf := new(bytes.Buffer)
|
|
writer := bufio.NewWriter(buf)
|
|
// write token
|
|
writeToken(writer, TOKEN_DONE)
|
|
writer.Flush()
|
|
return buf.Bytes()
|
|
}
|
|
|
|
// TOKEN_ERROR_MESSAGE message
|
|
// length (bytes)
|
|
// token 2
|
|
//
|
|
// message len 4
|
|
// message (from prev)
|
|
|
|
func WriteError(err error) []byte {
|
|
buf := new(bytes.Buffer)
|
|
writer := bufio.NewWriter(buf)
|
|
// write token
|
|
writeToken(writer, TOKEN_ERROR_MESSAGE)
|
|
// write error len
|
|
t := err.Error()
|
|
b := make([]byte, 4)
|
|
binary.BigEndian.PutUint32(b, uint32(len(t)))
|
|
writer.Write(b)
|
|
// write error
|
|
writer.WriteString(t)
|
|
writer.Flush()
|
|
|
|
return buf.Bytes()
|
|
}
|
|
|
|
func ReadError(reader io.Reader) (string, error) {
|
|
var len int16
|
|
err := binary.Read(reader, binary.BigEndian, &len)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
bvalue := make([]byte, len)
|
|
err = binary.Read(reader, binary.BigEndian, &bvalue)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return string(bvalue), nil
|
|
}
|
|
|
|
func writeToken(w io.Writer, token int16) {
|
|
writeInt16(w, token)
|
|
}
|
|
|
|
func writeInt8(w io.Writer, i int8) {
|
|
b := make([]byte, 1)
|
|
b[0] = byte(i)
|
|
w.Write(b)
|
|
}
|
|
|
|
func writeInt16(w io.Writer, i int16) {
|
|
b := make([]byte, 2)
|
|
binary.BigEndian.PutUint16(b, uint16(i))
|
|
w.Write(b)
|
|
}
|
|
|
|
func writeInt32(w io.Writer, i int32) {
|
|
b := make([]byte, 4)
|
|
binary.BigEndian.PutUint32(b, uint32(i))
|
|
w.Write(b)
|
|
}
|
|
|
|
func writeInt64(w io.Writer, i int64) {
|
|
b := make([]byte, 8)
|
|
binary.BigEndian.PutUint64(b, uint64(i))
|
|
w.Write(b)
|
|
}
|