mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-09-07 09:05:55 +00:00
Merge pull request #724 from niaow/pg-primitive
Add primitive types to pg encoder
This commit is contained in:
Nia
GitHub
commit
96a8fd5cb7
4 changed files with 436 additions and 7 deletions
|
|
@ -2899,6 +2899,8 @@ func (t ExtractedTable) ToRows(callback func(*pb.RowResponse) error) error {
|
|||
for i, r := range c.Rows {
|
||||
var col *pb.ColumnResponse
|
||||
switch r := r.(type) {
|
||||
case nil:
|
||||
col = &pb.ColumnResponse{}
|
||||
case bool:
|
||||
col = &pb.ColumnResponse{
|
||||
ColumnVal: &pb.ColumnResponse_BoolVal{
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ package pgtest
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/pilosa/pilosa/v2/pg"
|
||||
)
|
||||
|
|
@ -29,3 +30,44 @@ func (h HandlerFunc) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q
|
|||
}
|
||||
|
||||
var _ pg.QueryHandler = HandlerFunc(nil)
|
||||
|
||||
// ResultSet is a QueryResultWriter that accumulates results in a slice.
|
||||
type ResultSet struct {
|
||||
Columns []pg.ColumnInfo
|
||||
Data [][]string
|
||||
ResultTag string
|
||||
}
|
||||
|
||||
// WriteHeader writes headers to the result set.
|
||||
func (rs *ResultSet) WriteHeader(cols ...pg.ColumnInfo) error {
|
||||
if rs.Columns != nil {
|
||||
return errors.New("double-write of headers")
|
||||
}
|
||||
|
||||
colsCopy := make([]pg.ColumnInfo, len(cols))
|
||||
copy(colsCopy, cols)
|
||||
rs.Columns = colsCopy
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// WriteRowText writes a row to the result set.
|
||||
func (rs *ResultSet) WriteRowText(vals ...string) error {
|
||||
if rs.Columns == nil {
|
||||
return errors.New("wrote a row without headers")
|
||||
}
|
||||
|
||||
row := make([]string, len(vals))
|
||||
copy(row, vals)
|
||||
|
||||
rs.Data = append(rs.Data, row)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Tag applies a tag to the result set.
|
||||
func (rs *ResultSet) Tag(tag string) {
|
||||
rs.ResultTag = tag
|
||||
}
|
||||
|
||||
var _ pg.QueryResultWriter = (*ResultSet)(nil)
|
||||
|
|
|
|||
72
server/pg.go
72
server/pg.go
|
|
@ -28,6 +28,7 @@ import (
|
|||
"github.com/pilosa/pilosa/v2"
|
||||
"github.com/pilosa/pilosa/v2/logger"
|
||||
"github.com/pilosa/pilosa/v2/pg"
|
||||
"github.com/pilosa/pilosa/v2/pql"
|
||||
pb "github.com/pilosa/pilosa/v2/proto"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
|
@ -49,12 +50,7 @@ func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) *
|
|||
api: api,
|
||||
logger: logger,
|
||||
s: pg.Server{
|
||||
QueryHandler: &queryDecodeHandler{
|
||||
child: &pilosaQueryHandler{
|
||||
api: api,
|
||||
logger: logger,
|
||||
},
|
||||
},
|
||||
QueryHandler: NewPostgresHandler(api, logger),
|
||||
TypeEngine: pg.PrimitiveTypeEngine{},
|
||||
StartupTimeout: 5 * time.Second,
|
||||
ReadTimeout: 10 * time.Second,
|
||||
|
|
@ -66,6 +62,16 @@ func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) *
|
|||
}
|
||||
}
|
||||
|
||||
// NewPostgresHandler creates a postgres query handler wrapping the pilosa API.
|
||||
func NewPostgresHandler(api *pilosa.API, logger logger.Logger) pg.QueryHandler {
|
||||
return &queryDecodeHandler{
|
||||
child: &pilosaQueryHandler{
|
||||
api: api,
|
||||
logger: logger,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Start a postgres endpoint at the specified address.
|
||||
func (s *PostgresServer) Start(addr string) error {
|
||||
l, err := net.Listen("tcp", addr)
|
||||
|
|
@ -197,6 +203,8 @@ func pgFormatVal(val interface{}) string {
|
|||
return strconv.FormatUint(val, 10)
|
||||
case string:
|
||||
return val
|
||||
case pql.Decimal:
|
||||
return val.String()
|
||||
default:
|
||||
data, _ := json.Marshal(val)
|
||||
return string(data)
|
||||
|
|
@ -316,10 +324,15 @@ func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error {
|
|||
for i, col := range row.Columns {
|
||||
var v string
|
||||
switch col := col.ColumnVal.(type) {
|
||||
case nil:
|
||||
v = "null"
|
||||
case *pb.ColumnResponse_BoolVal:
|
||||
v = strconv.FormatBool(col.BoolVal)
|
||||
case *pb.ColumnResponse_DecimalVal:
|
||||
v = col.DecimalVal.String()
|
||||
v = pql.Decimal{
|
||||
Value: col.DecimalVal.Value,
|
||||
Scale: col.DecimalVal.Scale,
|
||||
}.String()
|
||||
case *pb.ColumnResponse_Float64Val:
|
||||
v = strconv.FormatFloat(col.Float64Val, 'g', -1, 64)
|
||||
case *pb.ColumnResponse_Int64Val:
|
||||
|
|
@ -381,6 +394,51 @@ func pgWriteResult(w pg.QueryResultWriter, result interface{}) error {
|
|||
return pgWriteRowser(w, result)
|
||||
case pb.StreamClient:
|
||||
return pgWriteRowser(w, &clientRowser{result})
|
||||
case uint64:
|
||||
err := w.WriteHeader(pg.ColumnInfo{
|
||||
Name: "count",
|
||||
Type: pg.TypeCharoid,
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing headers")
|
||||
}
|
||||
|
||||
err = w.WriteRowText(strconv.FormatUint(result, 10))
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing count")
|
||||
}
|
||||
|
||||
return nil
|
||||
case int64:
|
||||
err := w.WriteHeader(pg.ColumnInfo{
|
||||
Name: "value",
|
||||
Type: pg.TypeCharoid,
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing headers")
|
||||
}
|
||||
|
||||
err = w.WriteRowText(strconv.FormatInt(result, 10))
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing count")
|
||||
}
|
||||
|
||||
return nil
|
||||
case bool:
|
||||
err := w.WriteHeader(pg.ColumnInfo{
|
||||
Name: "result",
|
||||
Type: pg.TypeCharoid,
|
||||
})
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing headers")
|
||||
}
|
||||
|
||||
err = w.WriteRowText(strconv.FormatBool(result))
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "writing count")
|
||||
}
|
||||
|
||||
return nil
|
||||
default:
|
||||
return errors.Errorf("result type %T not yet supported", result)
|
||||
}
|
||||
|
|
|
|||
327
server/pg_test.go
Normal file
327
server/pg_test.go
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
// Copyright 2020 Pilosa Corp.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package server_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/pilosa/pilosa/v2"
|
||||
"github.com/pilosa/pilosa/v2/logger"
|
||||
"github.com/pilosa/pilosa/v2/pg"
|
||||
"github.com/pilosa/pilosa/v2/pg/pgtest"
|
||||
"github.com/pilosa/pilosa/v2/server"
|
||||
"github.com/pilosa/pilosa/v2/test"
|
||||
)
|
||||
|
||||
func TestPostgresHandler(t *testing.T) {
|
||||
m := test.RunCommand(t)
|
||||
defer m.Close()
|
||||
|
||||
pgh := server.NewPostgresHandler(m.API, logger.NewLogfLogger(t))
|
||||
|
||||
m.MustCreateIndex(t, "i", pilosa.IndexOptions{TrackExistence: true})
|
||||
m.MustCreateField(t, "i", "set")
|
||||
m.MustCreateField(t, "i", "keyset", pilosa.OptFieldKeys())
|
||||
m.MustCreateField(t, "i", "mutex", pilosa.OptFieldTypeMutex(pilosa.CacheTypeNone, 0))
|
||||
m.MustCreateField(t, "i", "keymutex", pilosa.OptFieldKeys(), pilosa.OptFieldTypeMutex(pilosa.CacheTypeNone, 0))
|
||||
m.MustCreateField(t, "i", "int", pilosa.OptFieldTypeInt(math.MinInt64, math.MaxInt64))
|
||||
m.MustCreateField(t, "i", "decimal", pilosa.OptFieldTypeDecimal(2))
|
||||
m.MustCreateField(t, "i", "time", pilosa.OptFieldTypeTime("YMDH"))
|
||||
m.MustCreateField(t, "i", "bool", pilosa.OptFieldTypeBool())
|
||||
|
||||
m.MustCreateIndex(t, "j", pilosa.IndexOptions{TrackExistence: true, Keys: true})
|
||||
m.MustCreateField(t, "j", "set")
|
||||
|
||||
storeOK := pgtest.ResultSet{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{
|
||||
Name: "result",
|
||||
Type: pg.TypeCharoid,
|
||||
},
|
||||
},
|
||||
Data: [][]string{
|
||||
{
|
||||
"true",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
Name string
|
||||
Queries []string
|
||||
Results []pgtest.ResultSet
|
||||
}{
|
||||
{
|
||||
Name: "Extract-Nothing",
|
||||
Queries: []string{
|
||||
`[i]Extract(All(), Rows(set))`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{
|
||||
Name: "_id",
|
||||
Type: pg.TypeCharoid,
|
||||
},
|
||||
{
|
||||
Name: "set",
|
||||
Type: pg.TypeCharoid,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Store",
|
||||
Queries: []string{
|
||||
`[i]Store(ConstRow(columns=[1, 2, 3]), set=4)`,
|
||||
`[i]Store(ConstRow(columns=[0, 2, 4]), set=5)`,
|
||||
`[j]Store(ConstRow(columns=[1, 2, 3]), set=4)`,
|
||||
`[j]Store(ConstRow(columns=[0, 2, 4]), set=5)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Set",
|
||||
Queries: []string{
|
||||
`[i]Set(1, keyset="a")`,
|
||||
`[i]Set(2, keyset="b")`,
|
||||
`[i]Set(3, mutex=3)`,
|
||||
`[i]Set(4, keymutex="d")`,
|
||||
`[i]Set(1, int=5)`,
|
||||
`[i]Set(2, decimal=6.01)`,
|
||||
`[i]Set(3, time=7, 2016-01-01T00:00)`,
|
||||
`[i]Set(4, bool=false)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
storeOK,
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Extract",
|
||||
Queries: []string{
|
||||
`[i]Extract(
|
||||
All(),
|
||||
Rows(set), Rows(keyset),
|
||||
Rows(mutex), Rows(keymutex),
|
||||
Rows(int), Rows(decimal),
|
||||
Rows(time),
|
||||
Rows(bool)
|
||||
)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "_id", Type: pg.TypeCharoid},
|
||||
{Name: "set", Type: pg.TypeCharoid},
|
||||
{Name: "keyset", Type: pg.TypeCharoid},
|
||||
{Name: "mutex", Type: pg.TypeCharoid},
|
||||
{Name: "keymutex", Type: pg.TypeCharoid},
|
||||
{Name: "int", Type: pg.TypeCharoid},
|
||||
{Name: "decimal", Type: pg.TypeCharoid},
|
||||
{Name: "time", Type: pg.TypeCharoid},
|
||||
{Name: "bool", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{`1`, `[4]`, `["a"]`, `null`, `null`, `5`, `null`, `[]`, `null`},
|
||||
{`2`, `[4,5]`, `["b"]`, `null`, `null`, `null`, `6.01`, `[]`, `null`},
|
||||
{`3`, `[4]`, `[]`, `3`, `null`, `null`, `null`, `[7]`, `null`},
|
||||
{`4`, `[5]`, `[]`, `null`, `d`, `null`, `null`, `[]`, `false`},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "GroupBy",
|
||||
Queries: []string{
|
||||
`[i]GroupBy(Rows(set))`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "set", Type: pg.TypeCharoid},
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
{Name: "sum", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"4", "3", "0"},
|
||||
{"5", "3", "0"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Count",
|
||||
Queries: []string{
|
||||
`[i]Count(Row(set=4))`,
|
||||
`[i]Count(Row(int > 0))`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{{"3"}},
|
||||
},
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{{"1"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "FieldValue",
|
||||
Queries: []string{
|
||||
`[i]FieldValue(field=int, column=1)`,
|
||||
`[i]FieldValue(field=decimal, column=2)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "value", Type: pg.TypeCharoid},
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"5", "1"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "value", Type: pg.TypeCharoid},
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"6.01", "1"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Rows",
|
||||
Queries: []string{
|
||||
`[i]Rows(set)`,
|
||||
`[i]Rows(keyset)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "set", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"4"},
|
||||
{"5"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "keyset", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"a"},
|
||||
{"b"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "TopN",
|
||||
Queries: []string{
|
||||
`[i]TopN(set)`,
|
||||
`[i]TopN(keyset)`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "set", Type: pg.TypeCharoid},
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"5", "3"},
|
||||
{"4", "3"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "keyset", Type: pg.TypeCharoid},
|
||||
{Name: "count", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{"b", "1"},
|
||||
{"a", "1"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "SQL",
|
||||
Queries: []string{
|
||||
`select _id from i;`,
|
||||
},
|
||||
Results: []pgtest.ResultSet{
|
||||
{
|
||||
Columns: []pg.ColumnInfo{
|
||||
{Name: "_id", Type: pg.TypeCharoid},
|
||||
},
|
||||
Data: [][]string{
|
||||
{`1`},
|
||||
{`2`},
|
||||
{`3`},
|
||||
{`4`},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
c := c
|
||||
t.Run(c.Name, func(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
for i, q := range c.Queries {
|
||||
var res pgtest.ResultSet
|
||||
err := pgh.HandleQuery(ctx, &res, pg.SimpleQuery(q))
|
||||
if err != nil {
|
||||
t.Errorf("query %q failed: %v", q, err)
|
||||
continue
|
||||
}
|
||||
|
||||
expected := c.Results[i]
|
||||
if !reflect.DeepEqual(res, expected) {
|
||||
t.Errorf("query %q returned incorrect results: expected %v but got %v", q, expected, res)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue