mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 02:44:59 +00:00
1789 lines
53 KiB
Go
1789 lines
53 KiB
Go
// Copyright 2022 Molecula Corp. (DBA FeatureBase).
|
||
// SPDX-License-Identifier: Apache-2.0
|
||
package server
|
||
|
||
import (
|
||
"context"
|
||
"crypto/tls"
|
||
"fmt"
|
||
"net"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
pilosa "github.com/featurebasedb/featurebase/v3"
|
||
"github.com/featurebasedb/featurebase/v3/authn"
|
||
"github.com/featurebasedb/featurebase/v3/authz"
|
||
"github.com/featurebasedb/featurebase/v3/logger"
|
||
"github.com/featurebasedb/featurebase/v3/monitor"
|
||
"github.com/featurebasedb/featurebase/v3/pql"
|
||
pb "github.com/featurebasedb/featurebase/v3/proto"
|
||
vdsm_pb "github.com/featurebasedb/featurebase/v3/proto/vdsm"
|
||
"github.com/featurebasedb/featurebase/v3/sql"
|
||
"github.com/improbable-eng/grpc-web/go/grpcweb"
|
||
"github.com/pkg/errors"
|
||
"google.golang.org/grpc"
|
||
"google.golang.org/grpc/codes"
|
||
"google.golang.org/grpc/credentials"
|
||
"google.golang.org/grpc/metadata"
|
||
"google.golang.org/grpc/peer"
|
||
"google.golang.org/grpc/reflection"
|
||
"google.golang.org/grpc/status"
|
||
"vitess.io/vitess/go/vt/sqlparser"
|
||
)
|
||
|
||
// GRPCHandler contains methods which handle the various gRPC requests.
|
||
type GRPCHandler struct {
|
||
api *pilosa.API
|
||
perms *authz.GroupPermissions
|
||
logger logger.Logger
|
||
queryLogger logger.Logger
|
||
inspectDeprecated sync.Once
|
||
}
|
||
|
||
func NewGRPCHandler(api *pilosa.API) *GRPCHandler {
|
||
return &GRPCHandler{api: api, logger: logger.NopLogger}
|
||
}
|
||
|
||
func (h *GRPCHandler) WithLogger(logger logger.Logger) *GRPCHandler {
|
||
h.logger = logger
|
||
return h
|
||
}
|
||
|
||
func (h *GRPCHandler) WithPerms(perms *authz.GroupPermissions) *GRPCHandler {
|
||
h.perms = perms
|
||
return h
|
||
}
|
||
|
||
func (h *GRPCHandler) WithQueryLogger(logger logger.Logger) *GRPCHandler {
|
||
h.queryLogger = logger
|
||
return h
|
||
}
|
||
|
||
// errorToStatusError appends an appropriate grpc status code
|
||
// to the error (returning it as a status.Error).
|
||
func errToStatusError(err error) error {
|
||
if err == nil {
|
||
return status.New(codes.OK, "").Err()
|
||
}
|
||
|
||
// Check error string.
|
||
switch cause := errors.Cause(err); cause {
|
||
case pilosa.ErrIndexNotFound,
|
||
pilosa.ErrFieldNotFound,
|
||
pilosa.ErrForeignIndexNotFound,
|
||
pilosa.ErrBSIGroupNotFound:
|
||
return status.Error(codes.NotFound, err.Error())
|
||
|
||
case pilosa.ErrIndexExists,
|
||
pilosa.ErrFieldExists,
|
||
pilosa.ErrBSIGroupExists:
|
||
return status.Error(codes.AlreadyExists, err.Error())
|
||
|
||
case pilosa.ErrIndexRequired,
|
||
pilosa.ErrFieldRequired,
|
||
pilosa.ErrColumnRequired,
|
||
pilosa.ErrBSIGroupNameRequired,
|
||
pilosa.ErrName,
|
||
pilosa.ErrQueryRequired,
|
||
pilosa.ErrFieldsArgumentRequired,
|
||
pilosa.ErrIntFieldWithKeys,
|
||
pilosa.ErrDecimalFieldWithKeys:
|
||
return status.Error(codes.FailedPrecondition, err.Error())
|
||
|
||
case pilosa.ErrInvalidView,
|
||
pilosa.ErrInvalidBSIGroupType,
|
||
pilosa.ErrInvalidBSIGroupValueType,
|
||
pilosa.ErrInvalidCacheType:
|
||
return status.Error(codes.InvalidArgument, err.Error())
|
||
|
||
case pilosa.ErrDecimalOutOfRange,
|
||
pilosa.ErrBSIGroupValueTooLow,
|
||
pilosa.ErrBSIGroupValueTooHigh,
|
||
pilosa.ErrInvalidRangeOperation,
|
||
pilosa.ErrInvalidBetweenValue:
|
||
return status.Error(codes.OutOfRange, err.Error())
|
||
|
||
case pilosa.ErrQueryTimeout:
|
||
return status.Error(codes.DeadlineExceeded, err.Error())
|
||
|
||
case pilosa.ErrQueryCancelled:
|
||
return status.Error(codes.Canceled, err.Error())
|
||
|
||
case pilosa.ErrNotImplemented:
|
||
return status.Error(codes.Unimplemented, err.Error())
|
||
|
||
case pilosa.ErrAborted:
|
||
return status.Error(codes.Aborted, err.Error())
|
||
|
||
case pilosa.ErrClusterDoesNotOwnShard,
|
||
pilosa.ErrNodeNotPrimary,
|
||
pilosa.ErrTooManyWrites,
|
||
pilosa.ErrNodeIDNotExists:
|
||
return status.Error(codes.Internal, err.Error())
|
||
default:
|
||
if _, ok := cause.(pilosa.ConflictError); ok {
|
||
return status.Error(codes.AlreadyExists, err.Error())
|
||
}
|
||
}
|
||
|
||
return status.Error(codes.Unknown, err.Error())
|
||
}
|
||
|
||
func (h *GRPCHandler) execSQL(ctx context.Context, queryStr string) (pb.ToRowser, error) {
|
||
pilosa.CounterSQLQueries.Inc()
|
||
return execSQL(ctx, h.api, h.logger, queryStr)
|
||
}
|
||
|
||
func isAllowed(requested []string, allowed []string) bool {
|
||
if len(allowed) == 0 {
|
||
return false
|
||
}
|
||
|
||
for _, r := range requested {
|
||
in := false
|
||
for _, a := range allowed {
|
||
if a == r {
|
||
in = true
|
||
}
|
||
}
|
||
if !in {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// QuerySQL handles the SQL request and sends RowResponses to the stream.
|
||
func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQLServer) error {
|
||
ctx := stream.Context()
|
||
if uinfo, ok := authn.GetUserInfo(ctx); ok && uinfo != nil {
|
||
// authz
|
||
m := sql.NewMapper()
|
||
parsed, err := m.MapSQL(req.Sql)
|
||
if err != nil {
|
||
return errors.Wrap(err, "parsing SQL")
|
||
}
|
||
|
||
perm := authz.Read
|
||
switch parsed.Statement.(type) {
|
||
case *sqlparser.DDL: // currently only used for DropTable
|
||
perm = authz.Admin
|
||
}
|
||
|
||
allowed := h.perms.GetAuthorizedIndexList(uinfo.Groups, perm)
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
if !isAllowed(parsed.Tables, allowed) {
|
||
return status.Error(codes.PermissionDenied, "insufficient permissions to access requested tables")
|
||
}
|
||
ctx = authn.WithIndexes(ctx, allowed)
|
||
}
|
||
LogQuery(ctx, "QuerySQL", req, h.queryLogger)
|
||
}
|
||
|
||
span := monitor.StartSpan(ctx, "GRPC", "/pilosa.Pilosa/QuerySQL")
|
||
|
||
span.SetTag("SQL Query", req.Sql)
|
||
start := time.Now()
|
||
results, err := h.execSQL(ctx, req.Sql)
|
||
duration := time.Since(start)
|
||
monitor.Finish(span)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
err = stream.SendHeader(metadata.New(map[string]string{
|
||
"duration": strconv.Itoa(int(duration)),
|
||
}))
|
||
if err != nil {
|
||
return errors.Wrap(err, "sending header")
|
||
}
|
||
|
||
err = newDurationRowser(results, duration).ToRows(stream.Send)
|
||
if err != nil {
|
||
return errors.Wrap(err, "streaming result")
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// QuerySQLUnary is a unary-response (non-streaming) version of QuerySQL, returning a TableResponse.
|
||
//
|
||
// Note regarding QuerySQLUnary and QueryPQLUnary:
|
||
// These methods are not ideal, as gRPC responses are payload-length limited to
|
||
// 4MB, so in most cases, we would recommend users use the QuerySQL and
|
||
// QueryPQL methods, as they stream the response as several small RowResponses.
|
||
// The response size limit is configurable on the client size, but we really
|
||
// only recommend these methods in the case that the payload is known to be
|
||
// quite small (e.g. single counts). These are provided mostly to support gRPC
|
||
// Futures, which are used by python-molecula to perform multiple queries
|
||
// concurrently. There is additional discussion and historical context here:
|
||
// https://github.com/molecula/pilosa/pull/644
|
||
func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest) (*pb.TableResponse, error) {
|
||
start := time.Now()
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
// authz
|
||
m := sql.NewMapper()
|
||
parsed, err := m.MapSQL(req.Sql)
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "parsing SQL")
|
||
}
|
||
|
||
perm := authz.Read
|
||
switch parsed.Statement.(type) {
|
||
case *sqlparser.DDL: // currently only used for DropTable
|
||
perm = authz.Admin
|
||
}
|
||
|
||
allowed := h.perms.GetAuthorizedIndexList(uinfo.Groups, perm)
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
if !isAllowed(parsed.Tables, allowed) {
|
||
return nil, status.Error(codes.PermissionDenied, "insufficient permissions to access requested tables")
|
||
}
|
||
ctx = authn.WithIndexes(ctx, allowed)
|
||
}
|
||
}
|
||
|
||
results, err := h.execSQL(ctx, req.Sql)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var table *pb.TableResponse
|
||
switch results := results.(type) {
|
||
case pb.ToTabler:
|
||
table, err = results.ToTable()
|
||
default:
|
||
table, err = pb.RowsToTable(results, 0)
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
duration := time.Since(start)
|
||
table.Duration = int64(duration)
|
||
err = grpc.SendHeader(ctx, metadata.New(map[string]string{
|
||
"duration": strconv.Itoa(int(duration)),
|
||
}))
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "sending header")
|
||
}
|
||
|
||
return table, nil
|
||
}
|
||
|
||
// QueryPQL handles the PQL request and sends RowResponses to the stream.
|
||
func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQLServer) error {
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: req.Pql,
|
||
}
|
||
|
||
ctx := stream.Context()
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
lperm := authz.Read
|
||
q, err := pql.ParseString(req.Pql)
|
||
if err != nil {
|
||
return status.Error(codes.InvalidArgument, err.Error())
|
||
}
|
||
if q.WriteCallN() > 0 {
|
||
lperm = authz.Write
|
||
}
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.Groups, lperm)) {
|
||
return status.Error(codes.PermissionDenied, "insufficient permissions to access requested indexes")
|
||
}
|
||
}
|
||
LogQuery(ctx, "QueryPQL", req, h.queryLogger)
|
||
}
|
||
span := monitor.StartSpan(ctx, "GRPC", "/pilosa.Pilosa/QueryPQL")
|
||
span.SetTag("PQL Query", req.Pql)
|
||
span.SetTag("Index", req.Index)
|
||
t := time.Now()
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
durQuery := time.Since(t)
|
||
monitor.Finish(span)
|
||
|
||
if err != nil {
|
||
return errToStatusError(err)
|
||
} else if len(resp.Results) != 1 {
|
||
// TODO: make a test for this
|
||
return status.Error(codes.InvalidArgument, "QueryPQL handles exactly one query")
|
||
}
|
||
longQueryTime := h.api.LongQueryTime()
|
||
if longQueryTime > 0 && durQuery > longQueryTime {
|
||
h.logger.Infof("GRPC QueryPQL %v %s", durQuery, query.Query)
|
||
}
|
||
|
||
rslt := resp.Results[0]
|
||
toRowser, err := ToRowserWrapper(rslt)
|
||
if err != nil {
|
||
return errors.Wrap(err, "wrapping as type ToRowser")
|
||
}
|
||
|
||
err = stream.SendHeader(metadata.New(map[string]string{
|
||
"duration": strconv.Itoa(int(durQuery)),
|
||
}))
|
||
if err != nil {
|
||
return errors.Wrap(err, "sending header")
|
||
}
|
||
|
||
t = time.Now()
|
||
if err := newDurationRowser(toRowser, durQuery).ToRows(stream.Send); err != nil {
|
||
return errToStatusError(err)
|
||
}
|
||
durFormat := time.Since(t)
|
||
pilosa.SummaryGRPCStreamQueryDurationSeconds.Observe(durQuery.Seconds())
|
||
pilosa.SummaryGRPCStreamFormatDurationSeconds.Observe(durFormat.Seconds())
|
||
pilosa.CounterPQLQueries.Inc()
|
||
|
||
return errToStatusError(nil)
|
||
}
|
||
|
||
// QueryPQLUnary is a unary-response (non-streaming) version of QueryPQL, returning a TableResponse.
|
||
//
|
||
// Note comment above QuerySQLUnary describing the need for the *Unary methods.
|
||
func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest) (*pb.TableResponse, error) {
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: req.Pql,
|
||
}
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
lperm := authz.Read
|
||
q, err := pql.ParseString(req.Pql)
|
||
if err != nil {
|
||
return nil, status.Error(codes.InvalidArgument, err.Error())
|
||
}
|
||
if q.WriteCallN() > 0 {
|
||
lperm = authz.Write
|
||
}
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.Groups, lperm)) {
|
||
return nil, status.Error(codes.PermissionDenied, fmt.Sprintf("insufficient permissions for %v", req.Index))
|
||
}
|
||
}
|
||
}
|
||
|
||
t := time.Now()
|
||
resp, err := h.api.Query(ctx, &query)
|
||
durQuery := time.Since(t)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
} else if len(resp.Results) != 1 {
|
||
return nil, status.Error(codes.InvalidArgument, "QueryPQLUnary handles exactly one query")
|
||
}
|
||
longQueryTime := h.api.LongQueryTime()
|
||
if longQueryTime > 0 && durQuery > longQueryTime {
|
||
h.logger.Infof("GRPC QueryPQLUnary %v %s", durQuery, query.Query)
|
||
}
|
||
|
||
rslt := resp.Results[0]
|
||
toTabler, err := ToTablerWrapper(rslt)
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "wrapping as type ToTabler")
|
||
}
|
||
|
||
t = time.Now()
|
||
table, err := toTabler.ToTable()
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
durFormat := time.Since(t)
|
||
|
||
duration := durQuery + durFormat
|
||
table.Duration = int64(duration)
|
||
err = grpc.SendHeader(ctx, metadata.New(map[string]string{
|
||
"duration": strconv.Itoa(int(duration)),
|
||
}))
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "sending header")
|
||
}
|
||
|
||
pilosa.SummaryGRPCUnaryQueryDurationSeconds.Observe(durQuery.Seconds())
|
||
pilosa.SummaryGRPCUnaryFormatDurationSeconds.Observe(durFormat.Seconds())
|
||
pilosa.CounterPQLQueries.Inc()
|
||
|
||
return table, errToStatusError(nil)
|
||
}
|
||
|
||
// CreateIndex creates a new Index
|
||
func (h *GRPCHandler) CreateIndex(ctx context.Context, req *pb.CreateIndexRequest) (*pb.CreateIndexResponse, error) {
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
return nil, status.Error(codes.PermissionDenied, "must be admin to create index")
|
||
}
|
||
}
|
||
// Always enable TrackExistence for gRPC-created indexes
|
||
opts := pilosa.IndexOptions{Keys: req.Keys, TrackExistence: true}
|
||
_, err := h.api.CreateIndex(ctx, req.Name, opts)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
if err := grpc.SendHeader(ctx, metadata.MD{}); err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
return &pb.CreateIndexResponse{}, nil
|
||
}
|
||
|
||
// GetIndex returns a single Index given a name
|
||
func (h *GRPCHandler) GetIndex(ctx context.Context, req *pb.GetIndexRequest) (*pb.GetIndexResponse, error) {
|
||
if uinfo, ok := authn.GetUserInfo(ctx); uinfo != nil {
|
||
if !ok {
|
||
return nil, status.Error(codes.InvalidArgument, "malformed auth header")
|
||
}
|
||
p, err := h.perms.GetPermissions(uinfo, req.Name)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if !p.Satisfies(authz.Read) {
|
||
return nil, status.Error(codes.PermissionDenied, fmt.Sprintf("permission denied for index %v", req.Name))
|
||
}
|
||
}
|
||
schema, err := h.api.Schema(ctx, false)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
|
||
for _, index := range schema {
|
||
if req.Name == index.Name {
|
||
return &pb.GetIndexResponse{Index: &pb.Index{Name: index.Name}}, nil
|
||
}
|
||
}
|
||
if err := grpc.SendHeader(ctx, metadata.MD{}); err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
return nil, status.Error(codes.NotFound, fmt.Sprintf("Index with name %s not found", req.Name))
|
||
}
|
||
|
||
// GetIndexes returns a list of all Indexes
|
||
func (h *GRPCHandler) GetIndexes(ctx context.Context, req *pb.GetIndexesRequest) (*pb.GetIndexesResponse, error) {
|
||
userInfo, _ := authn.GetUserInfo(ctx) // nil is a valid userInfo here
|
||
schema, err := h.api.Schema(ctx, false)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
|
||
indexes := make([]*pb.Index, 0)
|
||
for _, index := range schema {
|
||
if userInfo != nil {
|
||
if p, err := h.perms.GetPermissions(userInfo, index.Name); err == nil && p.Satisfies(authz.Read) {
|
||
indexes = append(indexes, &pb.Index{Name: index.Name})
|
||
}
|
||
} else {
|
||
indexes = append(indexes, &pb.Index{Name: index.Name})
|
||
}
|
||
}
|
||
|
||
if err := grpc.SendHeader(ctx, metadata.MD{}); err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
|
||
return &pb.GetIndexesResponse{Indexes: indexes}, nil
|
||
}
|
||
|
||
// DeleteIndex deletes an Index
|
||
func (h *GRPCHandler) DeleteIndex(ctx context.Context, req *pb.DeleteIndexRequest) (*pb.DeleteIndexResponse, error) {
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
if !h.perms.IsAdmin(uinfo.Groups) {
|
||
return nil, status.Error(codes.PermissionDenied, "must be admin to delete index")
|
||
}
|
||
}
|
||
err := h.api.DeleteIndex(ctx, req.Name)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
if err := grpc.SendHeader(ctx, metadata.MD{}); err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
return &pb.DeleteIndexResponse{}, nil
|
||
}
|
||
|
||
// VDSMGRPCHandler contains methods which handle the various gRPC requests, ported from VDSM.
|
||
type VDSMGRPCHandler struct {
|
||
grpcHandler *GRPCHandler
|
||
api *pilosa.API
|
||
logger logger.Logger
|
||
}
|
||
|
||
func NewVDSMGRPCHandler(grpcHandler *GRPCHandler, api *pilosa.API) *VDSMGRPCHandler {
|
||
return &VDSMGRPCHandler{grpcHandler: grpcHandler, api: api, logger: logger.NopLogger}
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) WithLogger(logger logger.Logger) *VDSMGRPCHandler {
|
||
h.logger = logger
|
||
return h
|
||
}
|
||
|
||
// GetVDSs returns a single VDS given a name
|
||
func (h *VDSMGRPCHandler) GetVDS(ctx context.Context, req *vdsm_pb.GetVDSRequest) (*vdsm_pb.GetVDSResponse, error) {
|
||
typedIdOrName := req.GetIdOrName()
|
||
switch idOrName := typedIdOrName.(type) {
|
||
case *vdsm_pb.GetVDSRequest_Id:
|
||
return nil, status.Error(codes.InvalidArgument, "VDS IDs are no longer supported")
|
||
case *vdsm_pb.GetVDSRequest_Name:
|
||
schema, err := h.api.Schema(ctx, false)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
|
||
for _, index := range schema {
|
||
if idOrName.Name == index.Name {
|
||
return &vdsm_pb.GetVDSResponse{Vds: &vdsm_pb.VDS{Name: index.Name}}, nil
|
||
}
|
||
}
|
||
return nil, status.Error(codes.NotFound, fmt.Sprintf("VDS with name %s not found", idOrName.Name))
|
||
default:
|
||
return nil, status.Error(codes.NotFound, "VDS not found")
|
||
}
|
||
}
|
||
|
||
// GetVDSs returns a list of all VDSs
|
||
func (h *VDSMGRPCHandler) GetVDSs(ctx context.Context, req *vdsm_pb.GetVDSsRequest) (*vdsm_pb.GetVDSsResponse, error) {
|
||
schema, err := h.api.Schema(ctx, false)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
|
||
vdss := make([]*vdsm_pb.VDS, len(schema))
|
||
for i, index := range schema {
|
||
vdss[i] = &vdsm_pb.VDS{Name: index.Name}
|
||
}
|
||
return &vdsm_pb.GetVDSsResponse{Vdss: vdss}, nil
|
||
}
|
||
|
||
// PostVDS creates a new VDS
|
||
func (*VDSMGRPCHandler) PostVDS(ctx context.Context, req *vdsm_pb.PostVDSRequest) (*vdsm_pb.PostVDSResponse, error) {
|
||
// Pilosa doesn't implement VDSD files, so this is unimplemented
|
||
return nil, status.Errorf(codes.Unimplemented, "method PostVDS not implemented")
|
||
}
|
||
|
||
// DeleteVDS deletes a VDS
|
||
func (h *VDSMGRPCHandler) DeleteVDS(ctx context.Context, req *vdsm_pb.DeleteVDSRequest) (*vdsm_pb.DeleteVDSResponse, error) {
|
||
typedIdOrName := req.GetIdOrName()
|
||
switch idOrName := typedIdOrName.(type) {
|
||
case *vdsm_pb.DeleteVDSRequest_Id:
|
||
return nil, status.Error(codes.InvalidArgument, "VDS IDs are no longer supported")
|
||
case *vdsm_pb.DeleteVDSRequest_Name:
|
||
err := h.api.DeleteIndex(ctx, idOrName.Name)
|
||
if err != nil {
|
||
return nil, errToStatusError(err)
|
||
}
|
||
return &vdsm_pb.DeleteVDSResponse{}, nil
|
||
default:
|
||
return nil, status.Error(codes.NotFound, "")
|
||
}
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) QuerySQL(req *pb.QuerySQLRequest, srv vdsm_pb.Molecula_QuerySQLServer) error {
|
||
return h.grpcHandler.QuerySQL(req, srv)
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest) (*pb.TableResponse, error) {
|
||
return h.grpcHandler.QuerySQLUnary(ctx, req)
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) QueryPQL(req *vdsm_pb.QueryPQLRequest, srv vdsm_pb.Molecula_QueryPQLServer) error {
|
||
preq := &pb.QueryPQLRequest{Index: req.Vds, Pql: req.Pql}
|
||
return h.grpcHandler.QueryPQL(preq, srv)
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) QueryPQLUnary(ctx context.Context, req *vdsm_pb.QueryPQLRequest) (*pb.TableResponse, error) {
|
||
preq := &pb.QueryPQLRequest{Index: req.Vds, Pql: req.Pql}
|
||
return h.grpcHandler.QueryPQLUnary(ctx, preq)
|
||
}
|
||
|
||
func (h *VDSMGRPCHandler) Inspect(req *vdsm_pb.InspectRequest, srv vdsm_pb.Molecula_InspectServer) error {
|
||
preq := &pb.InspectRequest{Index: req.Vds, Columns: req.Records, FilterFields: req.FilterFields, Limit: req.Limit, Offset: req.Offset, Query: req.Query}
|
||
return h.grpcHandler.Inspect(preq, srv)
|
||
}
|
||
|
||
// ResultUint64 is a wrapper around a uint64 result type
|
||
// so that we can implement the ToTabler and ToRowser
|
||
// interfaces.
|
||
type ResultUint64 uint64
|
||
|
||
// ToTable implements the ToTabler interface.
|
||
func (r ResultUint64) ToTable() (*pb.TableResponse, error) {
|
||
return pb.RowsToTable(&r, 1)
|
||
}
|
||
|
||
// ToRows implements the ToRowser interface.
|
||
func (r ResultUint64) ToRows(callback func(*pb.RowResponse) error) error {
|
||
return callback(&pb.RowResponse{
|
||
Headers: []*pb.ColumnInfo{{Name: "count", Datatype: "uint64"}},
|
||
Columns: []*pb.ColumnResponse{
|
||
{ColumnVal: &pb.ColumnResponse_Uint64Val{Uint64Val: uint64(r)}},
|
||
},
|
||
})
|
||
}
|
||
|
||
// ResultBool is a wrapper around a bool result type
|
||
// so that we can implement the ToTabler and ToRowser
|
||
// interfaces.
|
||
type ResultBool bool
|
||
|
||
// ToTable implements the ToTabler interface.
|
||
func (r ResultBool) ToTable() (*pb.TableResponse, error) {
|
||
return pb.RowsToTable(&r, 1)
|
||
}
|
||
|
||
// ToRows implements the ToRowser interface.
|
||
func (r ResultBool) ToRows(callback func(*pb.RowResponse) error) error {
|
||
return callback(&pb.RowResponse{
|
||
Headers: []*pb.ColumnInfo{{Name: "result", Datatype: "bool"}},
|
||
Columns: []*pb.ColumnResponse{
|
||
{ColumnVal: &pb.ColumnResponse_BoolVal{BoolVal: bool(r)}},
|
||
},
|
||
})
|
||
}
|
||
|
||
// Normally we wouldn't need this wrapper, but since pilosa returns
|
||
// some concrete types for which we can't implement the ToTabler
|
||
// interface, we have to check for those here and then wrap them
|
||
// with a custom type.
|
||
func ToTablerWrapper(result interface{}) (pb.ToTabler, error) {
|
||
toTabler, ok := result.(pb.ToTabler)
|
||
if !ok {
|
||
switch v := result.(type) {
|
||
case []pilosa.GroupCount:
|
||
gc := pilosa.NewGroupCounts("", v...)
|
||
toTabler = gc
|
||
case uint64:
|
||
toTabler = ResultUint64(v)
|
||
case bool:
|
||
toTabler = ResultBool(v)
|
||
default:
|
||
return nil, errors.Errorf("ToTabler interface not implemented by type: %T", result)
|
||
}
|
||
}
|
||
return toTabler, nil
|
||
}
|
||
|
||
// Normally we wouldn't need this wrapper, but since pilosa returns
|
||
// some concrete types for which we can't implement the ToRowser
|
||
// interface, we have to check for those here and then wrap them
|
||
// with a custom type.
|
||
func ToRowserWrapper(result interface{}) (pb.ToRowser, error) {
|
||
toRowser, ok := result.(pb.ToRowser)
|
||
if !ok {
|
||
switch v := result.(type) {
|
||
case []pilosa.GroupCount:
|
||
gc := pilosa.NewGroupCounts("", v...)
|
||
toRowser = gc
|
||
case uint64:
|
||
toRowser = ResultUint64(v)
|
||
case bool:
|
||
toRowser = ResultBool(v)
|
||
default:
|
||
return nil, errors.Errorf("ToRowser interface not implemented by type: %T", result)
|
||
}
|
||
}
|
||
return toRowser, nil
|
||
}
|
||
|
||
// durationRowser is a wrapper for pb.ToRowser that can be used to inject a
|
||
// duration value into the first record in a stream
|
||
type durationRowser struct {
|
||
pb.ToRowser
|
||
duration time.Duration
|
||
once sync.Once
|
||
}
|
||
|
||
func (r *durationRowser) ToRows(callback func(*pb.RowResponse) error) error {
|
||
cb := func(rr *pb.RowResponse) error {
|
||
r.once.Do(func() {
|
||
rr.Duration = int64(r.duration)
|
||
})
|
||
return callback(rr)
|
||
}
|
||
return r.ToRowser.ToRows(cb)
|
||
}
|
||
|
||
func newDurationRowser(orig pb.ToRowser, duration time.Duration) pb.ToRowser {
|
||
return &durationRowser{
|
||
ToRowser: orig,
|
||
duration: duration,
|
||
}
|
||
}
|
||
|
||
// Inspect handles the inspect request and sends an InspectResponse to the stream.
|
||
func (h *GRPCHandler) Inspect(req *pb.InspectRequest, stream pb.Pilosa_InspectServer) error {
|
||
const defaultLimit = 100000
|
||
|
||
h.inspectDeprecated.Do(func() {
|
||
h.logger.Infof("DEPRECATED: Inspect is deprecated, please use Extract() instead.")
|
||
})
|
||
|
||
ctx := stream.Context()
|
||
if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil {
|
||
LogQuery(stream.Context(), "Inspect", req, h.queryLogger)
|
||
}
|
||
|
||
index, err := h.api.Index(stream.Context(), req.Index)
|
||
if err != nil {
|
||
return errToStatusError(err)
|
||
}
|
||
|
||
qcx := h.api.Holder().Txf().NewQcx()
|
||
defer qcx.Abort()
|
||
|
||
var fields []*pilosa.Field
|
||
for _, field := range index.Fields() {
|
||
// exclude internal fields (starting with "_")
|
||
if strings.HasPrefix(field.Name(), "_") {
|
||
continue
|
||
}
|
||
if len(req.FilterFields) > 0 {
|
||
for _, filter := range req.FilterFields {
|
||
if filter == field.Name() {
|
||
fields = append(fields, field)
|
||
break
|
||
}
|
||
}
|
||
} else {
|
||
fields = append(fields, field)
|
||
}
|
||
}
|
||
|
||
if req.Query != "" {
|
||
// Execute the query and use it to select columns.
|
||
if req.Columns != nil && req.Columns.Type != nil {
|
||
l := 0
|
||
switch v := req.Columns.Type.(type) {
|
||
case *pb.IdsOrKeys_Ids:
|
||
l = len(v.Ids.Vals)
|
||
case *pb.IdsOrKeys_Keys:
|
||
l = len(v.Keys.Vals)
|
||
}
|
||
if l > 0 {
|
||
return errors.New("found a list of columns in a query-based inspect call")
|
||
}
|
||
}
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: req.Query,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "querying for columns with %q", req.Query)
|
||
}
|
||
if len(resp.Results) != 1 {
|
||
return errors.Errorf("expected 1 result for inspect query; got %d from %q", len(resp.Results), req.Query)
|
||
}
|
||
row, ok := resp.Results[0].(*pilosa.Row)
|
||
if !ok {
|
||
return errors.Errorf("incorrect query result type %T for query %q", resp.Results[0], req.Query)
|
||
}
|
||
if len(row.Keys) > 0 {
|
||
req.Columns = &pb.IdsOrKeys{
|
||
Type: &pb.IdsOrKeys_Keys{
|
||
Keys: &pb.StringArray{Vals: row.Keys},
|
||
},
|
||
}
|
||
} else {
|
||
req.Columns = &pb.IdsOrKeys{
|
||
Type: &pb.IdsOrKeys_Ids{
|
||
Ids: &pb.Uint64Array{Vals: row.Columns()},
|
||
},
|
||
}
|
||
}
|
||
if !row.Any() {
|
||
// No columns were matched.
|
||
return nil
|
||
}
|
||
}
|
||
|
||
limit := req.Limit
|
||
if limit == 0 {
|
||
limit = defaultLimit
|
||
}
|
||
offset := req.Offset
|
||
|
||
if !index.Keys() {
|
||
var cols []uint64
|
||
if req.Columns != nil {
|
||
ints, ok := req.Columns.Type.(*pb.IdsOrKeys_Ids)
|
||
if !ok {
|
||
return errors.New("invalid int columns")
|
||
}
|
||
cols = ints.Ids.Vals
|
||
}
|
||
ci := []*pb.ColumnInfo{
|
||
{Name: "_id", Datatype: "uint64"},
|
||
}
|
||
for _, field := range fields {
|
||
fdt := fieldDataType(field)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "field %s", field.Name())
|
||
}
|
||
ci = append(ci, &pb.ColumnInfo{Name: field.Name(), Datatype: fdt})
|
||
}
|
||
|
||
// If Columns is empty, then get the _exists list (via All()),
|
||
// from the index and loop over that instead.
|
||
if len(cols) > 0 {
|
||
// Apply limit/offset to the provided columns.
|
||
if int(offset) >= len(cols) {
|
||
return nil
|
||
}
|
||
end := limit + offset
|
||
if int(end) > len(cols) {
|
||
end = uint64(len(cols))
|
||
}
|
||
cols = cols[offset:end]
|
||
} else {
|
||
// Prevent getting too many records by forcing a limit.
|
||
pql := fmt.Sprintf("All(limit=%d, offset=%d)", limit, offset)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "querying for all: %s", pql)
|
||
}
|
||
|
||
ids, ok := resp.Results[0].(*pilosa.Row)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting results as a row")
|
||
}
|
||
|
||
limitedCols := ids.Columns()
|
||
if len(limitedCols) == 0 {
|
||
// If cols is still empty after the limit/offset, then
|
||
// return with no results.
|
||
return nil
|
||
}
|
||
cols = limitedCols
|
||
}
|
||
|
||
for _, col := range cols {
|
||
rowResp := &pb.RowResponse{
|
||
Headers: ci,
|
||
Columns: []*pb.ColumnResponse{
|
||
{ColumnVal: &pb.ColumnResponse_Uint64Val{Uint64Val: col}},
|
||
},
|
||
}
|
||
ci = nil // only include headers with the first row
|
||
|
||
colAdded := 0
|
||
for _, field := range fields {
|
||
// TODO: handle `time` fields
|
||
switch field.Type() {
|
||
case "set":
|
||
pql := fmt.Sprintf("Rows(%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "querying rows for set: %s", pql)
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Keys) > 0 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringArrayVal{StringArrayVal: &pb.StringArray{Vals: ids.Keys}}})
|
||
colAdded++
|
||
} else if len(ids.Rows) > 0 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Uint64ArrayVal{Uint64ArrayVal: &pb.Uint64Array{Vals: ids.Rows}}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "mutex":
|
||
pql := fmt.Sprintf("Rows(%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "querying rows for mutex")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Keys) == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringVal{StringVal: ids.Keys[0]}})
|
||
colAdded++
|
||
} else if len(ids.Rows) == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Uint64Val{Uint64Val: ids.Rows[0]}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "int":
|
||
if field.Keys() {
|
||
var value string
|
||
var exists bool
|
||
var err error
|
||
if fi := field.ForeignIndex(); fi != "" {
|
||
// Get the value from the int field.
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting int field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
vals, err := h.api.TranslateIndexIDs(stream.Context(), fi, []uint64{uint64(valCount.Val)})
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting keys for ids")
|
||
}
|
||
if len(vals) > 0 && vals[0] != "" {
|
||
value = vals[0]
|
||
exists = true
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
value, exists, err = field.StringValue(qcx, col)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting string field value for column")
|
||
}
|
||
}
|
||
if exists {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringVal{StringVal: value}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting int field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Int64Val{Int64Val: valCount.Val}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
}
|
||
|
||
case "decimal":
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting decimal field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
valVal := valCount.DecimalVal.Value()
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_DecimalVal{DecimalVal: &pb.Decimal{Value: (&valVal).Int64(), Scale: valCount.DecimalVal.Scale}}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "bool":
|
||
pql := fmt.Sprintf("Rows(%s, column=%d)", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "querying rows for bool")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Rows) == 1 {
|
||
var bval bool
|
||
if ids.Rows[0] == 1 {
|
||
bval = true
|
||
}
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_BoolVal{BoolVal: bval}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "time":
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
}
|
||
|
||
// For SQL queries like:
|
||
// SELECT * FROM t WHERE _id=garbageID;
|
||
// we don't want to return any rows.
|
||
// So, check here if we added any columns.
|
||
//
|
||
// Because we don't have keys to translate
|
||
// and _id is an artificial field that's why for query:
|
||
// SELECT _id FROM t WHERE _id=existing-id;
|
||
// we return an empty result.
|
||
//
|
||
// TODO(kuba--): We need to find a way to check here if
|
||
// existing-id is not a garbage.
|
||
//
|
||
// A query which will work here is 'SELECT *' or any query with more columns
|
||
// than just _id.
|
||
if colAdded > 0 {
|
||
if err := stream.Send(rowResp); err != nil {
|
||
return errors.Wrap(err, "sending response to stream")
|
||
}
|
||
}
|
||
}
|
||
|
||
} else {
|
||
var cols []string
|
||
|
||
if req.Columns != nil {
|
||
switch keys := req.Columns.Type.(type) {
|
||
case *pb.IdsOrKeys_Ids:
|
||
// The default behavior (in api/client/grpc.go) is to
|
||
// send an empty set of Ids even if the index supports
|
||
// keys, so in that case we just need to ignore it.
|
||
case *pb.IdsOrKeys_Keys:
|
||
cols = keys.Keys.Vals
|
||
default:
|
||
return errToStatusError(errors.New("invalid key columns"))
|
||
}
|
||
}
|
||
|
||
forceSend := false
|
||
ci := []*pb.ColumnInfo{
|
||
{Name: "_id", Datatype: "string"},
|
||
}
|
||
for _, field := range fields {
|
||
fdt := fieldDataType(field)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "field %s", field.Name())
|
||
}
|
||
ci = append(ci, &pb.ColumnInfo{Name: field.Name(), Datatype: fdt})
|
||
}
|
||
|
||
// If Columns is empty, then get the _exists list (via All()),
|
||
// from the index and loop over that instead.
|
||
if len(cols) > 0 {
|
||
// Apply limit/offset to the provided columns.
|
||
if int(offset) >= len(cols) {
|
||
return nil
|
||
}
|
||
end := limit + offset
|
||
if int(end) > len(cols) {
|
||
end = uint64(len(cols))
|
||
}
|
||
cols = cols[offset:end]
|
||
if len(cols) == 1 {
|
||
if id, err := h.api.TranslateIndexKey(stream.Context(), index.Name(), cols[0], false); id != 0 && err == nil {
|
||
forceSend = true
|
||
}
|
||
}
|
||
} else {
|
||
// Prevent getting too many records by forcing a limit.
|
||
pql := fmt.Sprintf("All(limit=%d, offset=%d)", limit, offset)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrapf(err, "querying for all: %s", pql)
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(*pilosa.Row)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting results as a row")
|
||
}
|
||
|
||
limitedCols := ids.Keys
|
||
if len(limitedCols) == 0 {
|
||
// If cols is still empty after the limit/offset, then
|
||
// return with no results.
|
||
return nil
|
||
}
|
||
cols = limitedCols
|
||
} else {
|
||
return errors.Errorf("expected 1 result for inspect query; got %d from %q", len(resp.Results), req.Query)
|
||
}
|
||
}
|
||
|
||
for _, col := range cols {
|
||
rowResp := &pb.RowResponse{
|
||
Headers: ci,
|
||
Columns: []*pb.ColumnResponse{
|
||
{ColumnVal: &pb.ColumnResponse_StringVal{StringVal: col}},
|
||
},
|
||
}
|
||
ci = nil // only include headers with the first row
|
||
|
||
colAdded := 0
|
||
for _, field := range fields {
|
||
// TODO: handle `time` fields
|
||
switch field.Type() {
|
||
case "set":
|
||
pql := fmt.Sprintf("Rows(%s, column=\"%s\")", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "querying set rows(keys)")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Keys) > 0 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringArrayVal{StringArrayVal: &pb.StringArray{Vals: ids.Keys}}})
|
||
colAdded++
|
||
} else if len(ids.Rows) > 0 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Uint64ArrayVal{Uint64ArrayVal: &pb.Uint64Array{Vals: ids.Rows}}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "mutex":
|
||
pql := fmt.Sprintf("Rows(%s, column=\"%s\")", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "querying mutex rows(keys)")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Keys) == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringVal{StringVal: ids.Keys[0]}})
|
||
colAdded++
|
||
} else if len(ids.Rows) == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Uint64Val{Uint64Val: ids.Rows[0]}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "int":
|
||
// Translate column key.
|
||
id, err := h.api.TranslateIndexKey(stream.Context(), index.Name(), col, false)
|
||
if err != nil {
|
||
return errors.Wrap(err, "translating column key")
|
||
}
|
||
|
||
if field.Keys() {
|
||
var value string
|
||
var exists bool
|
||
var err error
|
||
if fi := field.ForeignIndex(); fi != "" {
|
||
// Get the value from the int field.
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column=%d)", field.Name(), id)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting int field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
vals, err := h.api.TranslateIndexIDs(stream.Context(), fi, []uint64{uint64(valCount.Val)})
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting keys for ids")
|
||
}
|
||
if len(vals) > 0 && vals[0] != "" {
|
||
value = vals[0]
|
||
exists = true
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
value, exists, err = field.StringValue(qcx, id)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting string field value for column")
|
||
}
|
||
}
|
||
if exists {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_StringVal{StringVal: value}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column=%d)", field.Name(), id)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting int field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_Int64Val{Int64Val: valCount.Val}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
}
|
||
|
||
case "decimal":
|
||
pql := fmt.Sprintf("FieldValue(field=%s, column='%s')", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "getting decimal field value for column")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
valCount, ok := resp.Results[0].(pilosa.ValCount)
|
||
if ok && valCount.Count == 1 {
|
||
valVal := valCount.DecimalVal.Value()
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_DecimalVal{DecimalVal: &pb.Decimal{Value: (&valVal).Int64(), Scale: valCount.DecimalVal.Scale}}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
case "bool":
|
||
pql := fmt.Sprintf("Rows(%s, column=\"%s\")", field.Name(), col)
|
||
query := pilosa.QueryRequest{
|
||
Index: req.Index,
|
||
Query: pql,
|
||
}
|
||
resp, err := h.api.Query(stream.Context(), &query)
|
||
if err != nil {
|
||
return errors.Wrap(err, "querying bool rows(keys)")
|
||
}
|
||
|
||
if len(resp.Results) > 0 {
|
||
ids, ok := resp.Results[0].(pilosa.RowIdentifiers)
|
||
if !ok {
|
||
return errors.Wrap(err, "getting row identifiers")
|
||
}
|
||
|
||
if len(ids.Rows) == 1 {
|
||
var bval bool
|
||
if ids.Rows[0] == 1 {
|
||
bval = true
|
||
}
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: &pb.ColumnResponse_BoolVal{BoolVal: bval}})
|
||
colAdded++
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
} else {
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
|
||
default:
|
||
rowResp.Columns = append(rowResp.Columns,
|
||
&pb.ColumnResponse{ColumnVal: nil})
|
||
}
|
||
}
|
||
|
||
// For SQL queries like:
|
||
// SELECT _id FROM parent WHERE _id="garbage";
|
||
// we get here without any real columns and fields, and we did not
|
||
// translate any keys. That's why we don't want to send anything back
|
||
// and return fake response like:
|
||
//
|
||
// _id
|
||
// -------
|
||
// <nil>
|
||
// (1 row)
|
||
if colAdded > 0 || forceSend {
|
||
if err := stream.Send(rowResp); err != nil {
|
||
return errors.Wrap(err, "sending response to stream")
|
||
}
|
||
}
|
||
}
|
||
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// fieldDataType returns a useful data type (string,
|
||
// uint64, bool, etc.) based on the Pilosa field type.
|
||
// DO NOT USE THIS IN FUTURE CODE.
|
||
// This remains only for backwards-compatability within inspect.
|
||
// It does not produce sane results in all scenarios.
|
||
func fieldDataType(f *pilosa.Field) string {
|
||
switch f.Type() {
|
||
case "set":
|
||
if f.Keys() {
|
||
return "[]string"
|
||
}
|
||
return "[]uint64"
|
||
case "mutex":
|
||
if f.Keys() {
|
||
return "string"
|
||
}
|
||
return "uint64"
|
||
case "int":
|
||
if f.Keys() {
|
||
return "string"
|
||
}
|
||
return "int64"
|
||
case "decimal":
|
||
return "decimal"
|
||
case "bool":
|
||
return "bool"
|
||
case "time":
|
||
return "int64" // TODO: this is a placeholder
|
||
default:
|
||
panic(fmt.Sprintf("unimplemented fieldDataType: %s", f.Type()))
|
||
}
|
||
}
|
||
|
||
type grpcServer struct {
|
||
api *pilosa.API
|
||
grpcServer *grpc.Server
|
||
ln net.Listener
|
||
tlsConfig *tls.Config
|
||
auth *authn.Auth
|
||
perms *authz.GroupPermissions
|
||
|
||
logger logger.Logger
|
||
queryLogger logger.Logger
|
||
}
|
||
|
||
type grpcServerOption func(s *grpcServer) error
|
||
|
||
func OptGRPCServerAPI(api *pilosa.API) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.api = api
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerListener(ln net.Listener) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.ln = ln
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerTLSConfig(tlsConfig *tls.Config) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.tlsConfig = tlsConfig
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerLogger(logger logger.Logger) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.logger = logger
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerAuth(authn *authn.Auth) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.auth = authn
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerPerm(gp *authz.GroupPermissions) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.perms = gp
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func OptGRPCServerQueryLogger(logger logger.Logger) grpcServerOption {
|
||
return func(s *grpcServer) error {
|
||
s.queryLogger = logger
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func (s *grpcServer) Serve() error {
|
||
s.logger.Infof("enabled grpc listening on %s", s.ln.Addr())
|
||
|
||
// and start...
|
||
if err := s.grpcServer.Serve(s.ln); err != nil {
|
||
return errors.Wrap(err, "starting grpc server")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *grpcServer) middleware(origins []string) func(http.Handler) http.Handler {
|
||
httpOriginFunc := grpcweb.WithOriginFunc(func(origin string) bool {
|
||
for _, x := range origins {
|
||
if origin == x {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
})
|
||
|
||
wrappedGrpc := grpcweb.WrapServer(s.grpcServer, httpOriginFunc)
|
||
|
||
return func(h http.Handler) http.Handler {
|
||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if wrappedGrpc.IsGrpcWebRequest(r) || wrappedGrpc.IsAcceptableGrpcCorsRequest(r) {
|
||
wrappedGrpc.ServeHTTP(w, r)
|
||
} else {
|
||
h.ServeHTTP(w, r)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// Stop stops the GRPC server. There's no error because the underlying GRPC
|
||
// stuff doesn't report an error.
|
||
func (s *grpcServer) Stop() {
|
||
s.grpcServer.Stop()
|
||
}
|
||
|
||
func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) {
|
||
server := &grpcServer{
|
||
logger: logger.NopLogger,
|
||
}
|
||
for _, opt := range opts {
|
||
err := opt(server)
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "applying option")
|
||
}
|
||
}
|
||
|
||
gopts := make([]grpc.ServerOption, 0)
|
||
if server.tlsConfig != nil {
|
||
creds := credentials.NewTLS(server.tlsConfig)
|
||
gopts = append(gopts, grpc.Creds(creds))
|
||
}
|
||
|
||
// gRPC doesn’t allow multiple interceptors so they have to be manually chained.
|
||
var unaryInterceptors []grpc.UnaryServerInterceptor
|
||
var streamInterceptors []grpc.StreamServerInterceptor
|
||
|
||
if server.auth != nil {
|
||
unaryInterceptors = append(unaryInterceptors, func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
|
||
ctx, err := Valid(ctx, server.auth)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
LogQuery(ctx, info.FullMethod, req, server.logger)
|
||
|
||
// reset the molecula-chip cookie just in case the token was refreshed
|
||
md, ok := metadata.FromIncomingContext(ctx)
|
||
if uinfo, yeah := authn.GetUserInfo(ctx); ok && yeah {
|
||
server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken)
|
||
}
|
||
return handler(ctx, req)
|
||
})
|
||
streamInterceptors = append(streamInterceptors, func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
|
||
ctx, err := Valid(ss.Context(), server.auth)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
// reset the molecula-chip cookie just in case the token was refreshed
|
||
md, ok := metadata.FromIncomingContext(ctx)
|
||
if uinfo, yeah := authn.GetUserInfo(ctx); ok && yeah {
|
||
server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken)
|
||
}
|
||
return handler(srv, &wrappedStream{ss, ctx})
|
||
})
|
||
}
|
||
|
||
if monitor.IsOn() {
|
||
unaryInterceptors = append(unaryInterceptors, monitorUnaryInterceptor)
|
||
}
|
||
|
||
gopts = append(gopts, grpc.UnaryInterceptor(ChainUnaryInterceptor(unaryInterceptors...)))
|
||
gopts = append(gopts, grpc.StreamInterceptor(ChainStreamInterceptors(streamInterceptors...)))
|
||
|
||
// create grpc server
|
||
server.grpcServer = grpc.NewServer(gopts...)
|
||
grpcHandler := NewGRPCHandler(server.api).WithLogger(server.logger).WithQueryLogger(server.queryLogger)
|
||
|
||
// add server permissions if we've got 'em
|
||
if server.perms != nil {
|
||
grpcHandler.perms = server.perms
|
||
}
|
||
|
||
pb.RegisterPilosaServer(server.grpcServer, grpcHandler)
|
||
vdsm_pb.RegisterMoleculaServer(server.grpcServer, NewVDSMGRPCHandler(grpcHandler, server.api).WithLogger(server.logger))
|
||
|
||
// register the server so its services are available to grpc_cli and others
|
||
reflection.Register(server.grpcServer)
|
||
|
||
return server, nil
|
||
}
|
||
|
||
// LogQuery logs requests
|
||
func LogQuery(ctx context.Context, method string, req interface{}, logger logger.Logger) {
|
||
uinfo, ok := authn.GetUserInfo(ctx)
|
||
if !ok {
|
||
uinfo = &authn.UserInfo{}
|
||
}
|
||
md, _ := metadata.FromIncomingContext(ctx)
|
||
p, ok := peer.FromContext(ctx)
|
||
ip := ""
|
||
if ok {
|
||
ip = p.Addr.String()
|
||
}
|
||
ua, ok := md["user-agent"]
|
||
if !ok {
|
||
ua = []string{""}
|
||
}
|
||
switch r := req.(type) {
|
||
case *pb.QueryPQLRequest:
|
||
logger.Infof("GRPC: %v, %v, %v, %v, %v, [%s]%s", ip, ua, method, uinfo.UserID, uinfo.UserName, r.Index, r.Pql)
|
||
case *pb.QuerySQLRequest:
|
||
logger.Infof("GRPC: %v, %v, %v, %v, %v, %s", ip, ua, method, uinfo.UserID, uinfo.UserName, r.Sql)
|
||
default:
|
||
logger.Infof("GRPC: %v, %v, %v, %v, %v", ip, ua, method, uinfo.UserID, uinfo.UserName)
|
||
}
|
||
}
|
||
|
||
// wrappedStream wraps around the embedded grpc.ServerStream, and intercepts the RecvMsg and
|
||
// SendMsg method call.
|
||
type wrappedStream struct {
|
||
grpc.ServerStream
|
||
uiContext context.Context
|
||
}
|
||
|
||
func (w *wrappedStream) Context() context.Context {
|
||
return w.uiContext
|
||
}
|
||
|
||
func (w *wrappedStream) RecvMsg(m interface{}) error {
|
||
return w.ServerStream.RecvMsg(m)
|
||
}
|
||
|
||
func (w *wrappedStream) SendMsg(m interface{}) error {
|
||
return w.ServerStream.SendMsg(m)
|
||
}
|
||
|
||
func Valid(ctx context.Context, auth *authn.Auth) (context.Context, error) {
|
||
md, ok := metadata.FromIncomingContext(ctx)
|
||
if !ok {
|
||
return ctx, status.Errorf(codes.InvalidArgument, "missing metadata")
|
||
}
|
||
|
||
access, refresh := getTokensFromMetadata(md)
|
||
if access == "" {
|
||
return ctx, status.Errorf(codes.InvalidArgument, "missing authorization token")
|
||
}
|
||
uinfo, err := auth.Authenticate(access, refresh)
|
||
if err != nil {
|
||
return ctx, status.Errorf(codes.Unauthenticated, err.Error())
|
||
}
|
||
|
||
return authn.WithUserInfo(ctx, uinfo), nil
|
||
}
|
||
|
||
func getTokensFromMetadata(md metadata.MD) (string, string) {
|
||
// We check lowercase and uppercase because some GRPC clients lowercase metadata
|
||
// names. This is the only place we get tokens from metadata in GRPC calls.
|
||
access, ok := md["authorization"]
|
||
if !ok {
|
||
access, ok = md["Authorization"]
|
||
}
|
||
|
||
refresh, ok2 := md[strings.ToLower(authn.RefreshHeaderName)]
|
||
if !ok2 {
|
||
refresh, ok2 = md[authn.RefreshHeaderName]
|
||
}
|
||
|
||
if !ok || !ok2 {
|
||
if cookies, there := md["cookie"]; there {
|
||
for _, c := range cookies {
|
||
for _, cookie := range strings.Split(c, ";") {
|
||
cookie = strings.TrimSpace(cookie)
|
||
if strings.HasPrefix(cookie, authn.AccessCookieName+"=") && len(access) == 0 {
|
||
access = strings.Split(cookie, authn.AccessCookieName+"=")[1:]
|
||
} else if strings.HasPrefix(cookie, authn.RefreshCookieName+"=") && len(refresh) == 0 {
|
||
refresh = strings.Split(cookie, authn.RefreshCookieName+"=")[1:]
|
||
}
|
||
if len(access) > 0 && len(refresh) > 0 {
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if len(access) == 0 {
|
||
access = []string{""}
|
||
}
|
||
if len(refresh) == 0 {
|
||
refresh = []string{""}
|
||
}
|
||
return strings.TrimPrefix(access[0], "Bearer "), refresh[0]
|
||
}
|
||
|
||
func monitorUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
|
||
if !monitor.IsOn() {
|
||
return handler(ctx, req)
|
||
}
|
||
|
||
span := monitor.StartSpan(ctx, "GRPC", info.FullMethod)
|
||
switch r := req.(type) {
|
||
case *pb.QueryPQLRequest:
|
||
span.SetTag("PQL Query", r.Pql)
|
||
span.SetTag("Index", r.Index)
|
||
case *pb.QuerySQLRequest:
|
||
span.SetTag("SQL Query", r.Sql)
|
||
}
|
||
resp, err := handler(ctx, req)
|
||
monitor.Finish(span)
|
||
return resp, err
|
||
}
|
||
|
||
// Chains together multiple unary interceptors.
|
||
func ChainUnaryInterceptor(interceptors ...grpc.UnaryServerInterceptor) grpc.UnaryServerInterceptor {
|
||
chain := func(interceptor grpc.UnaryServerInterceptor, unaryHandler grpc.UnaryHandler, info *grpc.UnaryServerInfo) grpc.UnaryHandler {
|
||
return func(ctx context.Context, req interface{}) (interface{}, error) {
|
||
return interceptor(ctx, req, info, unaryHandler)
|
||
}
|
||
}
|
||
|
||
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
|
||
chained := handler
|
||
for i := len(interceptors) - 1; i >= 0; i-- {
|
||
chained = chain(interceptors[i], chained, info)
|
||
}
|
||
return chained(ctx, req)
|
||
}
|
||
}
|
||
|
||
// Chains together multiple stream interceptors.
|
||
func ChainStreamInterceptors(interceptors ...grpc.StreamServerInterceptor) grpc.StreamServerInterceptor {
|
||
chain := func(interceptor grpc.StreamServerInterceptor, streamHandler grpc.StreamHandler, info *grpc.StreamServerInfo) grpc.StreamHandler {
|
||
return func(srv interface{}, stream grpc.ServerStream) error {
|
||
return interceptor(srv, stream, info, streamHandler)
|
||
}
|
||
}
|
||
return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
|
||
chained := handler
|
||
for i := len(interceptors) - 1; i >= 0; i-- {
|
||
chained = chain(interceptors[i], chained, info)
|
||
}
|
||
return chained(srv, stream)
|
||
}
|
||
}
|