From e021f5e6704208fb2e51bcbcf86fe0d4812c922a Mon Sep 17 00:00:00 2001 From: Xiangyu Wang Date: Wed, 26 May 2021 16:57:32 +0800 Subject: [PATCH] Support to use json to construct dummy request (#5430) Signed-off-by: Xiangyu Wang --- internal/proxynode/dummyreq.go | 46 ++++++++++++++++++++++ internal/proxynode/dummyreq_test.go | 59 +++++++++++++++++++++++++++++ internal/proxynode/impl.go | 48 +++++++++++++++++------ 3 files changed, 142 insertions(+), 11 deletions(-) create mode 100644 internal/proxynode/dummyreq.go create mode 100644 internal/proxynode/dummyreq_test.go diff --git a/internal/proxynode/dummyreq.go b/internal/proxynode/dummyreq.go new file mode 100644 index 0000000000..ae5ff2285c --- /dev/null +++ b/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 +} diff --git a/internal/proxynode/dummyreq_test.go b/internal/proxynode/dummyreq_test.go new file mode 100644 index 0000000000..80bd0b1ce8 --- /dev/null +++ b/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"}) +} diff --git a/internal/proxynode/impl.go b/internal/proxynode/impl.go index 81158108b3..4d41a99ba6 100644 --- a/internal/proxynode/impl.go +++ b/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" { - request := &milvuspb.RetrieveRequest{ - DbName: "", - CollectionName: "", - PartitionNames: []string{}, - Ids: &schemapb.IDs{}, - OutputFields: []string{}, + 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 } - _, _ = 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{ - 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) {