implements restful entities api (#15916)

Signed-off-by: shaoyue.chen <shaoyue.chen@zilliz.com>
pull/16118/head
shaoyue 2022-03-21 11:21:23 +08:00 committed by GitHub
parent 15a3fe41c5
commit 97e5d77953
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 230 additions and 4 deletions

View File

@ -24,6 +24,14 @@ func NewHandlers(proxy types.ProxyComponent) *Handlers {
func (h *Handlers) RegisterRoutesTo(router gin.IRouter) {
router.GET("/health", wrapHandler(h.handleGetHealth))
router.POST("/dummy", wrapHandler(h.handlePostDummy))
router.POST("/entities", wrapHandler(h.handleInsert))
router.DELETE("/entities", wrapHandler(h.handleDelete))
router.POST("/search", wrapHandler(h.handleSearch))
router.POST("/query", wrapHandler(h.handleQuery))
router.POST("/persist", wrapHandler(h.handleFlush))
router.GET("/distance", wrapHandler(h.handleCalcDistance))
}
func (h *Handlers) handleGetHealth(c *gin.Context) (interface{}, error) {
@ -33,9 +41,63 @@ func (h *Handlers) handleGetHealth(c *gin.Context) (interface{}, error) {
func (h *Handlers) handlePostDummy(c *gin.Context) (interface{}, error) {
req := milvuspb.DummyRequest{}
// use ShouldBind to supports binding JSON, XML, YAML, and protobuf.
err := c.ShouldBind(&req)
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse json failed: %v", errBadRequest, err)
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Dummy(c, &req)
}
func (h *Handlers) handleInsert(c *gin.Context) (interface{}, error) {
req := milvuspb.InsertRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Insert(c, &req)
}
func (h *Handlers) handleDelete(c *gin.Context) (interface{}, error) {
req := milvuspb.DeleteRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Delete(c, &req)
}
func (h *Handlers) handleSearch(c *gin.Context) (interface{}, error) {
req := milvuspb.SearchRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Search(c, &req)
}
func (h *Handlers) handleQuery(c *gin.Context) (interface{}, error) {
req := milvuspb.QueryRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Query(c, &req)
}
func (h *Handlers) handleFlush(c *gin.Context) (interface{}, error) {
req := milvuspb.FlushRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.Flush(c, &req)
}
func (h *Handlers) handleCalcDistance(c *gin.Context) (interface{}, error) {
req := milvuspb.CalcDistanceRequest{}
err := shouldBind(c, &req)
if err != nil {
return nil, fmt.Errorf("%w: parse body failed: %v", errBadRequest, err)
}
return h.proxy.CalcDistance(c, &req)
}

View File

@ -3,13 +3,18 @@ package httpserver
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/gin-gonic/gin/binding"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/milvuspb"
"github.com/milvus-io/milvus/internal/proto/schemapb"
"github.com/milvus-io/milvus/internal/types"
"github.com/stretchr/testify/assert"
)
@ -25,8 +30,72 @@ func (mockProxyComponent) Dummy(ctx context.Context, request *milvuspb.DummyRequ
return nil, nil
}
func (mockProxyComponent) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) {
if request.CollectionName == "" {
return nil, errors.New("body parse err")
}
return &milvuspb.MutationResult{Acknowledged: true}, nil
}
func (mockProxyComponent) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) {
if request.Expr == "" {
return nil, errors.New("body parse err")
}
return &milvuspb.MutationResult{Acknowledged: true}, nil
}
var searchResult = milvuspb.SearchResults{
Results: &schemapb.SearchResultData{
TopK: 10,
},
}
func (mockProxyComponent) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
if request.Dsl == "" {
return nil, errors.New("body parse err")
}
return &searchResult, nil
}
var queryResult = milvuspb.QueryResults{
CollectionName: "test",
}
func (mockProxyComponent) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
if request.Expr == "" {
return nil, errors.New("body parse err")
}
return &queryResult, nil
}
var flushResult = milvuspb.FlushResponse{
DbName: "default",
}
func (mockProxyComponent) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
if len(request.CollectionNames) < 1 {
return nil, errors.New("body parse err")
}
return &flushResult, nil
}
var calcDistanceResult = milvuspb.CalcDistanceResults{
Array: &milvuspb.CalcDistanceResults_IntDist{
IntDist: &schemapb.IntArray{
Data: []int32{1, 2, 3},
},
},
}
func (mockProxyComponent) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
if len(request.Params) < 1 {
return nil, errors.New("body parse err")
}
return &calcDistanceResult, nil
}
func TestHandlers(t *testing.T) {
mockProxy := mockProxyComponent{}
mockProxy := &mockProxyComponent{}
h := NewHandlers(mockProxy)
testEngine := gin.New()
h.RegisterRoutesTo(testEngine)
@ -86,4 +155,83 @@ func TestHandlers(t *testing.T) {
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
})
type testCase struct {
httpMethod string
path string
body interface{}
expectedStatus int
expectedBody interface{}
}
testCases := []testCase{
{
http.MethodPost, "/entities", &milvuspb.InsertRequest{CollectionName: "c1"},
http.StatusOK, &milvuspb.MutationResult{Acknowledged: true},
},
{
http.MethodPost, "/entities", []byte("bad request"),
http.StatusBadRequest, nil,
},
{
http.MethodDelete, "/entities", milvuspb.DeleteRequest{Expr: "some expr"},
http.StatusOK, &milvuspb.MutationResult{Acknowledged: true},
},
{
http.MethodDelete, "/entities", []byte("bad request"),
http.StatusBadRequest, nil,
},
{
http.MethodPost, "/search", milvuspb.SearchRequest{Dsl: "some dsl"},
http.StatusOK, &searchResult,
},
{
http.MethodPost, "/search", []byte("bad request"),
http.StatusBadRequest, nil,
},
{
http.MethodPost, "/query", milvuspb.QueryRequest{Expr: "some expr"},
http.StatusOK, &queryResult,
},
{
http.MethodPost, "/query", []byte("bad request"),
http.StatusBadRequest, nil,
},
{
http.MethodPost, "/persist", milvuspb.FlushRequest{CollectionNames: []string{"c1"}},
http.StatusOK, flushResult,
},
{
http.MethodPost, "/persist", []byte("bad request"),
http.StatusBadRequest, nil,
},
{
http.MethodGet, "/distance", milvuspb.CalcDistanceRequest{
Params: []*commonpb.KeyValuePair{
{Key: "key", Value: "val"},
}},
http.StatusOK, calcDistanceResult,
},
{
http.MethodGet, "/distance", []byte("bad request"),
http.StatusBadRequest, nil,
},
}
for _, tt := range testCases {
t.Run(fmt.Sprintf("%s %s %d", tt.httpMethod, tt.path, tt.expectedStatus), func(t *testing.T) {
body := []byte{}
if tt.body != nil {
body, _ = json.Marshal(tt.body)
}
req := httptest.NewRequest(tt.httpMethod, tt.path, bytes.NewReader(body))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, tt.expectedStatus, w.Code)
if tt.expectedBody != nil {
bodyBytes, err := json.Marshal(tt.expectedBody)
assert.NoError(t, err)
assert.Equal(t, bodyBytes, w.Body.Bytes())
}
})
}
}

View File

@ -25,7 +25,7 @@ func wrapHandler(handle handlerFunc) gin.HandlerFunc {
data, err := handle(c)
// format body by accept header, protobuf marshal not supported by gin by default
// TODO: add marshal handler to support protobuf response
formatOffered := []string{binding.MIMEJSON, binding.MIMEYAML, binding.MIMEXML}
formatOffered := []string{binding.MIMEJSON, binding.MIMEYAML}
bodyFormatNegotiate := gin.Negotiate{
Offered: formatOffered,
Data: data,
@ -51,3 +51,19 @@ func wrapHandler(handle handlerFunc) gin.HandlerFunc {
c.Negotiate(http.StatusOK, bodyFormatNegotiate)
}
}
// gin.ShouldBind() default as `form`, but we want JSON
func shouldBind(c *gin.Context, obj interface{}) error {
b := getBinding(c.ContentType())
return c.ShouldBindWith(obj, b)
}
func getBinding(contentType string) binding.Binding {
// ref: binding.Default
switch contentType {
case binding.MIMEYAML:
return binding.YAML
default:
return binding.JSON
}
}