mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 10:54:59 +00:00
315 lines
7.2 KiB
Go
315 lines
7.2 KiB
Go
// Copyright 2022 Molecula Corp. (DBA FeatureBase).
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
package proto
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
|
|
"github.com/pkg/errors"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
)
|
|
|
|
// StreamClient is an interface for a stream
|
|
// which can return a RowResponse sent to a
|
|
// stream via Send().
|
|
type StreamClient interface {
|
|
Recv() (*RowResponse, error)
|
|
}
|
|
|
|
// EmptyStream implements StreamClient interface.
|
|
// It always returns empty RowResponse
|
|
type EmptyStream struct{}
|
|
|
|
// Recv returns io.EOF
|
|
func (EmptyStream) Recv() (*RowResponse, error) {
|
|
return nil, io.EOF
|
|
}
|
|
|
|
// StreamServer is an interface for a stream
|
|
// which can accept a RowResponse to be later
|
|
// returned by the stream via Recv().
|
|
type StreamServer interface {
|
|
Send(*RowResponse) error
|
|
}
|
|
|
|
// ToTabler is an interface for any type that can
|
|
// represent itself as a TableResponse.
|
|
type ToTabler interface {
|
|
ToTable() (*TableResponse, error)
|
|
}
|
|
|
|
// ToRowser is an interface for any type that can
|
|
// represent itself as one or more RowResponses.
|
|
// ToRows takes a callback function which should be
|
|
// called for each row in the response.
|
|
type ToRowser interface {
|
|
ToRows(func(*RowResponse) error) error
|
|
}
|
|
|
|
// RowsToTable is a helper function which takes a ToRowser,
|
|
// along with the number of rows, and returns a TableResponse.
|
|
// The number of rows is treated as a hint.
|
|
func RowsToTable(tr ToRowser, n int) (*TableResponse, error) {
|
|
var headers []*ColumnInfo
|
|
rows := make([]*Row, 0, n)
|
|
|
|
// This callback gets called for every "row" in r.
|
|
// Each row populates its position in the pre-allocated
|
|
// `rows`. The headers get set based on those received
|
|
// in the first row.
|
|
cb := func(rr *RowResponse) error {
|
|
if len(rows) == 0 {
|
|
headers = rr.GetHeaders()
|
|
}
|
|
rows = append(rows, &Row{Columns: rr.GetColumns()})
|
|
return nil
|
|
}
|
|
|
|
if err := tr.ToRows(cb); err != nil {
|
|
return nil, errors.Wrap(err, "calling callback")
|
|
}
|
|
|
|
return &TableResponse{
|
|
Headers: headers,
|
|
Rows: rows,
|
|
}, nil
|
|
}
|
|
|
|
// Error is a helper function to create a RowResponse
|
|
// based on an error message. If the error is a grpc
|
|
// Status, then the status code is passed through.
|
|
func Error(err error) *RowResponse {
|
|
status, _ := status.FromError(err)
|
|
return &RowResponse{
|
|
StatusError: &StatusError{
|
|
Code: uint32(status.Code()),
|
|
Message: status.Err().Error(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// ErrorWrap prepends a message to the existing status
|
|
// error message.
|
|
func ErrorWrap(err error, message string) *RowResponse {
|
|
status, _ := status.FromError(err)
|
|
return &RowResponse{
|
|
StatusError: &StatusError{
|
|
Code: uint32(status.Code()),
|
|
Message: message + ": " + status.Err().Error(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// ErrorWrapf prepends a message to the existing status
|
|
// error message with the format specifier.
|
|
func ErrorWrapf(err error, format string, args ...interface{}) *RowResponse {
|
|
status, _ := status.FromError(err)
|
|
return &RowResponse{
|
|
StatusError: &StatusError{
|
|
Code: uint32(status.Code()),
|
|
Message: fmt.Sprintf(format, args...) + ": " + status.Err().Error(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// ErrorCode is a helper function to create a RowResponse
|
|
// based on a grpc status code and an error message.
|
|
func ErrorCode(err error, c codes.Code) *RowResponse {
|
|
return &RowResponse{
|
|
StatusError: &StatusError{
|
|
Code: uint32(c),
|
|
Message: err.Error(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// RowResponseSorter implements the sort interface for a
|
|
// provided []RowResponse based on the column index, type,
|
|
// and sort direction.
|
|
type RowResponseSorter struct {
|
|
colIdx []int
|
|
colDescending []bool
|
|
colType []string
|
|
|
|
rrs []*RowResponse
|
|
}
|
|
|
|
// NewRowResponseSorter return a new RowResponseSorter. It
|
|
// does input validation and returns an error if the inputs
|
|
// aren't compatible.
|
|
func NewRowResponseSorter(idxs []int, dirs []bool, typs []string, rrs []*RowResponse) (*RowResponseSorter, error) {
|
|
// Ensure the input slices are non-empty and equal size.
|
|
if len(idxs) == 0 {
|
|
return nil, errors.New("index list cannot be empty")
|
|
}
|
|
if len(dirs) != len(idxs) || len(typs) != len(idxs) {
|
|
return nil, errors.New("index, direction, and type lists must be the same size")
|
|
}
|
|
|
|
// Ensure the provided data types are supported by the sorter.
|
|
for i := range typs {
|
|
switch typs[i] {
|
|
case "[]uint64", "[]string", "bool", "float64", "int64", "string", "uint64":
|
|
// pass
|
|
default:
|
|
return nil, fmt.Errorf("unsupported data type: %s", typs[i])
|
|
}
|
|
}
|
|
|
|
// Ensure max(colIdx) is within size of rr.Columns.
|
|
if len(rrs) > 0 {
|
|
var maxColIdx int
|
|
for i := range idxs {
|
|
if idxs[i] > maxColIdx {
|
|
maxColIdx = idxs[i]
|
|
}
|
|
}
|
|
if maxColIdx >= len(rrs[0].Columns) {
|
|
return nil, fmt.Errorf("column index is out of range: %d", maxColIdx)
|
|
}
|
|
}
|
|
|
|
return &RowResponseSorter{
|
|
colIdx: idxs,
|
|
colDescending: dirs,
|
|
colType: typs,
|
|
rrs: rrs,
|
|
}, nil
|
|
|
|
}
|
|
|
|
func (r RowResponseSorter) Len() int { return len(r.rrs) }
|
|
func (r RowResponseSorter) Swap(i, j int) { r.rrs[i], r.rrs[j] = r.rrs[j], r.rrs[i] }
|
|
func (r RowResponseSorter) Less(i, j int) bool {
|
|
ri := r.rrs[i]
|
|
rj := r.rrs[j]
|
|
|
|
for i, idx := range r.colIdx {
|
|
coli := ri.Columns[idx]
|
|
colj := rj.Columns[idx]
|
|
var comp int
|
|
switch r.colType[i] {
|
|
case "[]uint64":
|
|
ai := coli.GetUint64ArrayVal().Vals
|
|
aj := colj.GetUint64ArrayVal().Vals
|
|
comp = func() int {
|
|
for ii := 0; ii < len(ai); ii++ {
|
|
if len(aj) == ii {
|
|
return 1
|
|
}
|
|
piv := ai[ii]
|
|
pjv := aj[ii]
|
|
if piv == pjv {
|
|
continue
|
|
} else if piv < pjv {
|
|
return -1
|
|
} else {
|
|
return 1
|
|
}
|
|
}
|
|
if len(aj) > len(ai) {
|
|
return -1
|
|
}
|
|
return 0
|
|
}()
|
|
case "[]string":
|
|
ai := coli.GetStringArrayVal().Vals
|
|
aj := colj.GetStringArrayVal().Vals
|
|
comp = func() int {
|
|
for ii := 0; ii < len(ai); ii++ {
|
|
if len(aj) == ii {
|
|
return 1
|
|
}
|
|
sComp := strings.Compare(ai[ii], aj[ii])
|
|
if sComp == 0 {
|
|
continue
|
|
} else {
|
|
return sComp
|
|
}
|
|
}
|
|
if len(aj) > len(ai) {
|
|
return -1
|
|
}
|
|
return 0
|
|
}()
|
|
case "bool":
|
|
bi := coli.GetBoolVal()
|
|
bj := colj.GetBoolVal()
|
|
if bi == bj {
|
|
comp = 0
|
|
} else if !bi && bj {
|
|
comp = -1
|
|
} else {
|
|
comp = 1
|
|
}
|
|
case "float64":
|
|
fi := coli.GetFloat64Val()
|
|
fj := colj.GetFloat64Val()
|
|
if fi == fj {
|
|
comp = 0
|
|
} else if fi < fj {
|
|
comp = -1
|
|
} else {
|
|
comp = 1
|
|
}
|
|
case "int64":
|
|
ni := coli.GetInt64Val()
|
|
nj := colj.GetInt64Val()
|
|
if ni == nj {
|
|
comp = 0
|
|
} else if ni < nj {
|
|
comp = -1
|
|
} else {
|
|
comp = 1
|
|
}
|
|
case "string":
|
|
comp = strings.Compare(coli.GetStringVal(), colj.GetStringVal())
|
|
case "uint64":
|
|
ni := coli.GetUint64Val()
|
|
nj := colj.GetUint64Val()
|
|
if ni == nj {
|
|
comp = 0
|
|
} else if ni < nj {
|
|
comp = -1
|
|
} else {
|
|
comp = 1
|
|
}
|
|
}
|
|
|
|
isDescending := r.colDescending[i]
|
|
|
|
switch comp {
|
|
case 0:
|
|
continue
|
|
case -1:
|
|
if isDescending {
|
|
return false
|
|
}
|
|
return true
|
|
case 1:
|
|
if isDescending {
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// ConstRowser implements ToRowser with a slice of row responses.
|
|
type ConstRowser []RowResponse
|
|
|
|
// ToRows calls a function with a pointer to each element of the slice.
|
|
func (c ConstRowser) ToRows(fn func(*RowResponse) error) error {
|
|
for i := range c {
|
|
err := fn(&c[i])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|