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

Signed-off-by: Xiangyu Wang <xiangyu.wang@zilliz.com>

internal/proxynode/dummyreq.go

@@ -0,0 +1,46 @@
+// 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
+}

internal/proxynode/dummyreq_test.go

@@ -0,0 +1,59 @@
+// 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"})
+}

internal/proxynode/impl.go

@@ -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{
-		request := &milvuspb.RetrieveRequest{
+		Response: `{"status": "fail"}`,
-			DbName:         "",
+	}
-			CollectionName: "",
+
-			PartitionNames: []string{},
+	// TODO(wxyu): change name RequestType to Request
-			Ids:            &schemapb.IDs{},
+	drt, err := parseDummyRequestType(req.RequestType)
-			OutputFields:   []string{},
+	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
 		}
 
-		_, _ = node.Retrieve(ctx, request)
+		request := &milvuspb.RetrieveRequest{
+			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,
+		}
+
+		_, 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{
+	log.Debug("cannot find specify dummy request type")
-		Response: `{"status": "fail"}`,
+	return failedResponse, nil
-	}, nil
 }
 
 func (node *ProxyNode) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {