diff --git a/server/grpc.go b/server/grpc.go index 3dcc420a1..af11b6f3c 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -26,7 +26,6 @@ import ( "github.com/pilosa/pilosa/v2" "github.com/pilosa/pilosa/v2/logger" pb "github.com/pilosa/pilosa/v2/proto" - "github.com/pilosa/pilosa/v2/sql" "github.com/pilosa/pilosa/v2/stats" "github.com/pkg/errors" "google.golang.org/grpc" @@ -120,24 +119,7 @@ func (h *GRPCHandler) DeleteVDS(ctx context.Context, req *pb.DeleteVDSRequest) ( } func (h *GRPCHandler) execSQL(ctx context.Context, queryStr string) (pb.StreamClient, error) { - mapper := sql.NewMapper() - mapper.Logger = h.logger - query, err := mapper.MapSQL(queryStr) - if err != nil { - return nil, errors.Wrap(err, "failed to map SQL") - } - var results pb.StreamClient - switch query.SQLType { - case sql.SQLTypeSelect: - handler := sql.NewSelectHandler(h.api) - results, err = handler.Handle(ctx, query) - if err != nil { - return nil, errors.Wrap(err, "failed to start SQL query") - } - default: - return nil, status.Errorf(codes.Unimplemented, "query type not supported") - } - return results, nil + return execSQL(ctx, h.api, h.logger, queryStr) } // QuerySQL handles the SQL request and sends RowResponses to the stream. diff --git a/server/pg.go b/server/pg.go index 98117c374..28bbef7af 100644 --- a/server/pg.go +++ b/server/pg.go @@ -19,6 +19,7 @@ import ( "crypto/tls" "encoding/json" "fmt" + "io" "net" "strconv" "strings" @@ -50,7 +51,8 @@ func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) * s: pg.Server{ QueryHandler: &queryDecodeHandler{ child: &pilosaQueryHandler{ - api: api, + api: api, + logger: logger, }, }, TypeEngine: pg.PrimitiveTypeEngine{}, @@ -125,7 +127,8 @@ func pgDecodePQL(str string) (q pg.Query, err error) { } type pilosaQueryHandler struct { - api *pilosa.API + api *pilosa.API + logger logger.Logger } func pgWriteRow(w pg.QueryResultWriter, row *pilosa.Row) error { @@ -342,6 +345,28 @@ func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error { }) } +type clientRowser struct { + pb.StreamClient +} + +func (cr *clientRowser) ToRows(f func(*pb.RowResponse) error) error { + for { + resp, err := cr.StreamClient.Recv() + if err != nil { + if err == io.EOF { + return nil + } + + return err + } + + err = f(resp) + if err != nil { + return err + } + } +} + func pgWriteResult(w pg.QueryResultWriter, result interface{}) error { switch result := result.(type) { case *pilosa.Row: @@ -354,6 +379,8 @@ func pgWriteResult(w pg.QueryResultWriter, result interface{}) error { return pgWriteGroupCount(w, result) case pb.ToRowser: // we should avoid protobuf where we can... return pgWriteRowser(w, result) + case pb.StreamClient: + return pgWriteRowser(w, &clientRowser{result}) default: return errors.Errorf("result type %T not yet supported", result) } @@ -374,6 +401,13 @@ func (pqh *pilosaQueryHandler) HandleQuery(ctx context.Context, w pg.QueryResult } return errors.Wrap(pgWriteResult(w, resp.Results[0]), "writing query result") + case pg.SimpleQuery: + resp, err := execSQL(ctx, pqh.api, pqh.logger, string(q)) + if err != nil { + return errors.Wrap(err, "executing query") + } + return errors.Wrap(pgWriteResult(w, resp), "writing query result") + default: return errors.Errorf("query type %T not yet supported (query: %s)", q, q) } diff --git a/server/sql.go b/server/sql.go new file mode 100644 index 000000000..1e648b9e3 --- /dev/null +++ b/server/sql.go @@ -0,0 +1,48 @@ +// 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 + +import ( + "context" + + "github.com/pilosa/pilosa/v2" + "github.com/pilosa/pilosa/v2/logger" + pb "github.com/pilosa/pilosa/v2/proto" + "github.com/pilosa/pilosa/v2/sql" + "github.com/pkg/errors" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func execSQL(ctx context.Context, api *pilosa.API, logger logger.Logger, queryStr string) (pb.StreamClient, error) { + mapper := sql.NewMapper() + mapper.Logger = logger + query, err := mapper.MapSQL(queryStr) + if err != nil { + return nil, errors.Wrap(err, "failed to map SQL") + } + var results pb.StreamClient + switch query.SQLType { + case sql.SQLTypeSelect: + handler := sql.NewSelectHandler(api) + results, err = handler.Handle(ctx, query) + if err != nil { + return nil, errors.Wrap(err, "failed to start SQL query") + } + default: + return nil, status.Errorf(codes.Unimplemented, "query type not supported") + } + return results, nil +}