add SQL to postgres endpoint

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.

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)
 	}

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
+}