mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-09-06 08:35:55 +00:00
3bc0dc28f0
In order to standardize results as streams of RowResponse, this PR introduces two interfaces `StreamClient` and `StreamServer`) which mirror the grpc stream interfaces. Upstream users (sqlmapper, vdsm, etc) can implement instances of these interfaces to ensure that results can stream through the entire sytem in an expected way. This PR also fixes a couple of missing data types.
114 lines
3.3 KiB
Go
114 lines
3.3 KiB
Go
// Copyright 2017 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 client
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
|
|
pb "github.com/pilosa/pilosa/v2/proto"
|
|
"github.com/pkg/errors"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
)
|
|
|
|
// GRPCClient is a client for working with the gRPC server.
|
|
type GRPCClient struct {
|
|
conn *grpc.ClientConn
|
|
}
|
|
|
|
// NewGRPCClient returns a new instance of GRPCClient.
|
|
func NewGRPCClient(dialTarget string, tlsConfig *tls.Config) (*GRPCClient, error) {
|
|
var opts []grpc.DialOption
|
|
if tlsConfig != nil {
|
|
creds := credentials.NewTLS(tlsConfig)
|
|
opts = append(opts, grpc.WithTransportCredentials(creds))
|
|
} else {
|
|
opts = append(opts, grpc.WithInsecure())
|
|
}
|
|
|
|
gconn, err := grpc.Dial(dialTarget, opts...)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "creating new grpc client")
|
|
}
|
|
|
|
return &GRPCClient{
|
|
conn: gconn,
|
|
}, nil
|
|
}
|
|
|
|
// Close closes any connections the client has opened.
|
|
func (c *GRPCClient) Close() error {
|
|
if c.conn != nil {
|
|
return c.conn.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Query returns a stream of RowResponse for the given index and PQL string.
|
|
func (c *GRPCClient) Query(ctx context.Context, index string, pql string) (pb.StreamClient, error) {
|
|
if c.conn == nil {
|
|
return nil, errors.New("client has not established a grpc connection")
|
|
}
|
|
|
|
grpcClient := pb.NewPilosaClient(c.conn)
|
|
|
|
stream, err := grpcClient.QueryPQL(ctx, &pb.QueryPQLRequest{
|
|
Index: index,
|
|
Pql: pql,
|
|
})
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "getting stream")
|
|
} else if stream == nil {
|
|
return nil, errors.New("could not create stream")
|
|
}
|
|
|
|
return stream, err
|
|
}
|
|
|
|
// Inspect returns a stream of RowResponse for the given index, columns, and filters.
|
|
// It is intended to mimic something like "select [fields] from table where recordID IN (...)".
|
|
func (c *GRPCClient) Inspect(ctx context.Context, index string, columnIDs []uint64, columnKeys []string, fieldFilters []string) (pb.StreamClient, error) {
|
|
if c.conn == nil {
|
|
return nil, errors.New("client has not established a grpc connection")
|
|
}
|
|
|
|
if len(columnIDs) > 0 && len(columnKeys) > 0 {
|
|
return nil, errors.New("only provide column ids or keys, not both")
|
|
}
|
|
|
|
// Convert columns to proto type IdsOrKeys.
|
|
idsOrKeys := &pb.IdsOrKeys{}
|
|
if len(columnKeys) > 0 {
|
|
idsOrKeys.Type = &pb.IdsOrKeys_Keys{Keys: &pb.StringArray{Vals: columnKeys}}
|
|
} else {
|
|
idsOrKeys.Type = &pb.IdsOrKeys_Ids{Ids: &pb.Uint64Array{Vals: columnIDs}}
|
|
}
|
|
|
|
grpcClient := pb.NewPilosaClient(c.conn)
|
|
|
|
stream, err := grpcClient.Inspect(ctx, &pb.InspectRequest{
|
|
Index: index,
|
|
Columns: idsOrKeys,
|
|
FilterFields: fieldFilters,
|
|
})
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "getting stream")
|
|
} else if stream == nil {
|
|
return nil, errors.New("could not create stream")
|
|
}
|
|
|
|
return stream, err
|
|
}
|