未验证 提交 e021f5e6 编写于 作者: X Xiangyu Wang 提交者: GitHub

Support to use json to construct dummy request (#5430)

Signed-off-by: NXiangyu Wang <xiangyu.wang@zilliz.com>
上级 5feb9ae9
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// 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 proxynode
import (
"encoding/json"
)
type dummyRequestType struct {
RequestType string `json:"request_type"`
}
func parseDummyRequestType(str string) (*dummyRequestType, error) {
drt := &dummyRequestType{}
if err := json.Unmarshal([]byte(str), &drt); err != nil {
return nil, err
}
return drt, nil
}
type dummyRetrieveRequest struct {
RequestType string `json:"request_type"`
DbName string `json:"dbname"`
CollectionName string `json:"collection_name"`
PartitionNames []string `json:"partition_names"`
Ids []int64 `json:"ids"`
OutputFields []string `json:"output_fields"`
}
func parseDummyRetrieveRequest(str string) (*dummyRetrieveRequest, error) {
dr := &dummyRetrieveRequest{}
if err := json.Unmarshal([]byte(str), &dr); err != nil {
return nil, err
}
return dr, nil
}
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// 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 proxynode
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestParseDummyRequestType(t *testing.T) {
invalidStr := `{"request_type":"retrieve"`
_, err := parseDummyRequestType(invalidStr)
assert.NotNil(t, err)
retrievetypeStr := `{"request_type":"retrieve"}`
drt, err := parseDummyRequestType(retrievetypeStr)
assert.Nil(t, err)
assert.Equal(t, drt.RequestType, "retrieve")
}
func TestParseDummyRetrieveRequest(t *testing.T) {
invalidStr := `{"request_type":"retrieve"`
_, err := parseDummyRetrieveRequest(invalidStr)
assert.NotNil(t, err)
onlytypeStr := `{"request_type":"retrieve"}`
drr, err := parseDummyRetrieveRequest(onlytypeStr)
assert.Nil(t, err)
assert.Equal(t, drr.RequestType, "retrieve")
assert.Equal(t, len(drr.DbName), 0)
fulltypeStr := `{
"request_type":"retrieve",
"dbname":"",
"collection_name":"test",
"partition_names": [],
"ids": [100, 101],
"output_fields": ["_id", "age"]
}`
drr2, err := parseDummyRetrieveRequest(fulltypeStr)
assert.Nil(t, err)
assert.Equal(t, drr2.RequestType, "retrieve")
assert.Equal(t, len(drr2.DbName), 0)
assert.Equal(t, drr2.CollectionName, "test")
assert.Equal(t, len(drr2.PartitionNames), 0)
assert.Equal(t, len(drr2.Ids), 2)
assert.Equal(t, drr2.Ids, []int64{100, 101})
assert.Equal(t, drr2.OutputFields, []string{"_id", "age"})
}
......@@ -1437,25 +1437,51 @@ func (node *ProxyNode) getSegmentsOfCollection(ctx context.Context, dbName strin
}
func (node *ProxyNode) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
if req.RequestType == "retrieve" {
failedResponse := &milvuspb.DummyResponse{
Response: `{"status": "fail"}`,
}
// TODO(wxyu): change name RequestType to Request
drt, err := parseDummyRequestType(req.RequestType)
if err != nil {
log.Debug("Failed to parse dummy request type")
return failedResponse, nil
}
if drt.RequestType == "retrieve" {
drr, err := parseDummyRetrieveRequest(req.RequestType)
if err != nil {
log.Debug("Failed to parse dummy retrieve request")
return failedResponse, nil
}
request := &milvuspb.RetrieveRequest{
DbName: "",
CollectionName: "",
PartitionNames: []string{},
Ids: &schemapb.IDs{},
OutputFields: []string{},
DbName: drr.DbName,
CollectionName: drr.CollectionName,
PartitionNames: drr.PartitionNames,
Ids: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: drr.Ids,
},
},
},
OutputFields: drr.OutputFields,
}
_, _ = node.Retrieve(ctx, request)
_, err = node.Retrieve(ctx, request)
if err != nil {
log.Debug("Failed to execute dummy retrieve")
return failedResponse, err
}
return &milvuspb.DummyResponse{
Response: `{"status": "success"}`,
}, nil
}
return &milvuspb.DummyResponse{
Response: `{"status": "fail"}`,
}, nil
log.Debug("cannot find specify dummy request type")
return failedResponse, nil
}
func (node *ProxyNode) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册