Repository: milvus-io/milvus-sdk-go
Pull request #135 (2021-07-13 10:56:52)

Add collection check
Add collection check

Related #133 

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/135#" title="Author self-approved">congqixia</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sunby: /lgtm
/approve
sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/135#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/135#issuecomment-878993451" title="Approved">sunby</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia,sunby]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client_grpc_dd.go b/client/client_grpc_dd.go
index 1122f40..b788c68 100644
--- a/client/client_grpc_dd.go
+++ b/client/client_grpc_dd.go
@@ -134,6 +134,14 @@ func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Sc
 	if err := validateSchema(collSchema); err != nil {
 		return err
 	}
+
+	has, err := c.HasCollection(ctx, collSchema.CollectionName)
+	if err != nil {
+		return err
+	}
+	if has {
+		return fmt.Errorf("collection %s already exist", collSchema.CollectionName)
+	}
 	sch := collSchema.ProtoMessage()
 	bs, err := proto.Marshal(sch)
 	if err != nil {
@@ -202,6 +210,14 @@ func (c *grpcClient) DescribeCollection(ctx context.Context, collName string) (*
 		return nil, ErrClientNotReady
 	}
 
+	has, err := c.HasCollection(ctx, collName)
+	if err != nil {
+		return nil, err
+	}
+	if !has {
+		return nil, collNotExistsErr(collName)
+	}
+
 	req := &server.DescribeCollectionRequest{
 		CollectionName: collName,
 	}
@@ -229,6 +245,14 @@ func (c *grpcClient) DropCollection(ctx context.Context, collName string) error
 		return ErrClientNotReady
 	}
 
+	has, err := c.HasCollection(ctx, collName)
+	if err != nil {
+		return err
+	}
+	if !has {
+		return collNotExistsErr(collName)
+	}
+
 	req := &server.DropCollectionRequest{
 		CollectionName: collName,
 	}
@@ -267,6 +291,14 @@ func (c *grpcClient) GetCollectionStatistics(ctx context.Context, collName strin
 		return nil, ErrClientNotReady
 	}
 
+	has, err := c.HasCollection(ctx, collName)
+	if err != nil {
+		return nil, err
+	}
+	if !has {
+		return nil, collNotExistsErr(collName)
+	}
+
 	req := &server.GetCollectionStatisticsRequest{
 		CollectionName: collName,
 	}
@@ -304,7 +336,6 @@ func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async
 				continue
 			}
 			target[segment.ID] = segment
-
 		}
 		for len(target) > 0 {
 			current, err := c.GetQuerySegmentInfo(ctx, collName)
@@ -320,7 +351,6 @@ func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async
 			}
 			time.Sleep(time.Millisecond * 100)
 		}
-
 	}
 	return nil
 }
diff --git a/client/client_grpc_dm.go b/client/client_grpc_dm.go
index 095ff29..4b06953 100644
--- a/client/client_grpc_dm.go
+++ b/client/client_grpc_dm.go
@@ -85,7 +85,6 @@ func (c *grpcClient) Flush(ctx context.Context, collName string, async bool) err
 		if has {
 			waitingSet := make(map[int64]struct{})
 			for _, segmentID := range segmentIDs.GetData() {
-				fmt.Println("wait", segmentID)
 				waitingSet[segmentID] = struct{}{}
 			}
 			flushed := func() bool {
diff --git a/client/errors.go b/client/errors.go
new file mode 100644
index 0000000..9228197
--- /dev/null
+++ b/client/errors.go
@@ -0,0 +1,19 @@
+// Copyright (C) 2019-2021 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 client provides milvus client functions
+package client
+
+import "fmt"
+
+func collNotExistsErr(collName string) error {
+	return fmt.Errorf("collection %s does not exist", collName)
+}


Test patch:
diff --git a/client/client_grpc_dd_test.go b/client/client_grpc_dd_test.go
index b7d33ea..5c5b41e 100644
--- a/client/client_grpc_dd_test.go
+++ b/client/client_grpc_dd_test.go
@@ -172,6 +172,8 @@ func TestGrpcClientListCollections(t *testing.T) {
 func TestGrpcClientCreateCollection(t *testing.T) {
 	ctx := context.Background()
 	c := testClient(ctx, t)
+	// default, all collection name returns false
+	mock.delInjection(mHasCollection)
 	t.Run("Test normal creation", func(t *testing.T) {
 		ds := defaultSchema()
 		shardsNum := int32(1)
@@ -279,12 +281,57 @@ func TestGrpcClientCreateCollection(t *testing.T) {
 			mock.delInjection(mCreateCollection)
 		}
 	})
+
+	t.Run("test duplicated collection", func(t *testing.T) {
+		m := make(map[string]struct{})
+		mock.setInjection(mCreateCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.CreateCollectionRequest)
+			if !ok {
+				return badRequestStatus()
+			}
+			m[req.GetCollectionName()] = struct{}{}
+
+			return successStatus()
+		})
+		mock.setInjection(mHasCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.HasCollectionRequest)
+			resp := &server.BoolResponse{}
+			if !ok {
+				return badRequestStatus()
+			}
+
+			_, has := m[req.GetCollectionName()]
+			resp.Value = has
+			s, err := successStatus()
+			resp.Status = s
+			return resp, err
+		})
+
+		assert.Nil(t, c.CreateCollection(ctx, defaultSchema(), 1))
+		assert.NotNil(t, c.CreateCollection(ctx, defaultSchema(), 1))
+	})
+}
+
+// default HasCollection injection, returns true only when collection name is `testCollectionName`
+var hasCollectionDefault = func(_ context.Context, raw proto.Message) (proto.Message, error) {
+	req, ok := raw.(*server.HasCollectionRequest)
+	resp := &server.BoolResponse{}
+	if !ok {
+		s, err := badRequestStatus()
+		resp.Status = s
+		return s, err
+	}
+	resp.Value = req.GetCollectionName() == testCollectionName
+	s, err := successStatus()
+	resp.Status = s
+	return resp, err
 }
 
 func TestGrpcClientDropCollection(t *testing.T) {
 	ctx := context.Background()
 	c := testClient(ctx, t)
 
+	mock.setInjection(mHasCollection, hasCollectionDefault)
 	mock.setInjection(mDropCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := (raw).(*server.DropCollectionRequest)
 		if !ok {
diff --git a/client/client_grpc_dm_test.go b/client/client_grpc_dm_test.go
index 75d01ea..7376940 100644
--- a/client/client_grpc_dm_test.go
+++ b/client/client_grpc_dm_test.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/golang/protobuf/proto"
 	"github.com/milvus-io/milvus-sdk-go/entity"
+	"github.com/milvus-io/milvus-sdk-go/internal/proto/common"
 	"github.com/milvus-io/milvus-sdk-go/internal/proto/schema"
 	"github.com/milvus-io/milvus-sdk-go/internal/proto/server"
 	"github.com/stretchr/testify/assert"
@@ -52,8 +54,70 @@ func TestGrpcClientFlush(t *testing.T) {
 
 	c := testClient(ctx, t)
 
-	assert.Nil(t, c.Flush(ctx, testCollectionName, false))
-	assert.Nil(t, c.Flush(ctx, testCollectionName, true))
+	t.Run("test async flush", func(t *testing.T) {
+		assert.Nil(t, c.Flush(ctx, testCollectionName, true))
+	})
+
+	t.Run("test sync flush", func(t *testing.T) {
+		// 1~10 segments
+		segCount := rand.Intn(10) + 1
+		segments := make([]int64, 0, segCount)
+		for i := 0; i < segCount; i++ {
+			segments = append(segments, rand.Int63())
+		}
+		// 500ms ~ 2s
+		flushTime := 500 + rand.Intn(1500)
+		start := time.Now()
+		flag := false
+		mock.setInjection(mFlush, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.FlushRequest)
+			resp := &server.FlushResponse{}
+			if !ok {
+				s, err := badRequestStatus()
+				resp.Status = s
+				return resp, err
+			}
+			assert.ElementsMatch(t, []string{testCollectionName}, req.GetCollectionNames())
+
+			resp.CollSegIDs = make(map[string]*schema.LongArray)
+			resp.CollSegIDs[testCollectionName] = &schema.LongArray{
+				Data: segments,
+			}
+
+			s, err := successStatus()
+			resp.Status = s
+			return resp, err
+		})
+		mock.setInjection(mGetPersistentSegmentInfo, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.GetPersistentSegmentInfoRequest)
+			resp := &server.GetPersistentSegmentInfoResponse{}
+			if !ok {
+				s, err := badRequestStatus()
+				resp.Status = s
+				return resp, err
+			}
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+
+			state := common.SegmentState_Flushing
+			if time.Since(start) > time.Duration(flushTime)*time.Millisecond {
+				state = common.SegmentState_Flushed
+				flag = true
+			}
+
+			for _, segID := range segments {
+				resp.Infos = append(resp.Infos, &server.PersistentSegmentInfo{
+					SegmentID: segID,
+					State:     state,
+				})
+			}
+
+			s, err := successStatus()
+			resp.Status = s
+			return resp, err
+		})
+		assert.Nil(t, c.Flush(ctx, testCollectionName, false))
+		assert.True(t, flag)
+	})
 }
 
 func TestGrpcSearch(t *testing.T) {
@@ -64,7 +128,7 @@ func TestGrpcSearch(t *testing.T) {
 
 	vectors := generateFloatVector(4096, testVectorDim)
 
-	results, err := c.Search(context.Background(), testCollectionName, []string{}, "int64 > 0", []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])},
+	results, err := c.Search(ctx, testCollectionName, []string{}, "int64 > 0", []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])},
 		testVectorField, entity.L2, 10, map[string]string{
 			"nprobe": "10",
 		})