mirror of https://github.com/milvus-io/milvus.git
implements restful entities api (#15916)
Signed-off-by: shaoyue.chen <shaoyue.chen@zilliz.com>pull/16118/head
parent
15a3fe41c5
commit
97e5d77953
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue