Repository: milvus-io/milvus-sdk-go
Pull request #388 (2023-01-12 08:08:50)

Add GetLoadingProgress, GetLoadState, GetVersion interface
/kind improvement

related:#387 
Signed-off-by: wayblink <anyang.wang@zilliz.com>

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

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/388#" title="Author self-approved">wayblink</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **congqixia** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @congqixia` in a comment when ready.

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).

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

- **[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/v2.2.x/OWNERS)**

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":["congqixia"]} -->
codecov[bot]: # [Codecov](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/388?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> :exclamation: No coverage uploaded for pull request base (`v2.2.x@ec2c109`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/388/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/388?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff            @@
##             v2.2.x     #388   +/-   ##
=========================================
  Coverage          ?   83.67%           
=========================================
  Files             ?       29           
  Lines             ?     3038           
  Branches          ?        0           
=========================================
  Hits              ?     2542           
  Misses            ?      404           
  Partials          ?       92           
```



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/388#pullrequestreview-1246976886" title="LGTM">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/388#" title="Author self-approved">wayblink</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/v2.2.x/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":[]} -->

Patch:
diff --git a/client/client.go b/client/client.go
index 1bdbb8d..f9c0ff3 100644
--- a/client/client.go
+++ b/client/client.go
@@ -164,6 +164,14 @@ type Client interface {
 	Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
 	// Revoke removes privilege from role.
 	Revoke(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
+
+	// GetLoadingProgress get the collection or partitions loading progress
+	GetLoadingProgress(ctx context.Context, collectionName string, partitionNames []string) (int64, error)
+	// GetLoadState get the collection or partitions load state
+	GetLoadState(ctx context.Context, collectionName string, partitionNames []string) (entity.LoadState, error)
+
+	// GetVersion get milvus version
+	GetVersion(ctx context.Context) (string, error)
 }
 
 // SearchResult contains the result from Search api of client
diff --git a/client/client_grpc_admin.go b/client/client_grpc_admin.go
new file mode 100644
index 0000000..49527de
--- /dev/null
+++ b/client/client_grpc_admin.go
@@ -0,0 +1,29 @@
+// 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
+
+import (
+	"context"
+
+	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
+)
+
+func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
+	if c.Service == nil {
+		return "", ErrClientNotReady
+	}
+	resp, err := c.Service.GetVersion(ctx, &server.GetVersionRequest{})
+	if err != nil {
+		return "", err
+	}
+	return resp.GetVersion(), nil
+}
diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index f3ebd4f..32e9d70 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -448,3 +448,45 @@ func (c *GrpcClient) GetReplicas(ctx context.Context, collName string) ([]*entit
 	}
 	return groups, nil
 }
+
+// GetLoadingProgress get the collection or partitions loading progress
+func (c *GrpcClient) GetLoadingProgress(ctx context.Context, collName string, partitionNames []string) (int64, error) {
+	if c.Service == nil {
+		return 0, ErrClientNotReady
+	}
+	if err := c.checkCollectionExists(ctx, collName); err != nil {
+		return 0, err
+	}
+
+	req := &server.GetLoadingProgressRequest{
+		CollectionName: collName,
+		PartitionNames: partitionNames,
+	}
+	resp, err := c.Service.GetLoadingProgress(ctx, req)
+	if err != nil {
+		return 0, err
+	}
+
+	return resp.GetProgress(), nil
+}
+
+// GetLoadState get the collection or partitions load state
+func (c *GrpcClient) GetLoadState(ctx context.Context, collName string, partitionNames []string) (entity.LoadState, error) {
+	if c.Service == nil {
+		return 0, ErrClientNotReady
+	}
+	if err := c.checkCollectionExists(ctx, collName); err != nil {
+		return 0, err
+	}
+
+	req := &server.GetLoadStateRequest{
+		CollectionName: collName,
+		PartitionNames: partitionNames,
+	}
+	resp, err := c.Service.GetLoadState(ctx, req)
+	if err != nil {
+		return 0, err
+	}
+
+	return entity.LoadState(resp.GetState()), nil
+}
diff --git a/client/client_mock.go b/client/client_mock.go
index 9a24939..283e35a 100644
--- a/client/client_mock.go
+++ b/client/client_mock.go
@@ -30,6 +30,7 @@ const (
 	MGetCollectionStatistics ServiceMethod = 108
 	MAlterCollection         ServiceMethod = 109
 	MGetLoadingProgress      ServiceMethod = 110
+	MGetLoadState            ServiceMethod = 111
 
 	MCreatePartition   ServiceMethod = 201
 	MDropPartition     ServiceMethod = 202
@@ -62,6 +63,7 @@ const (
 	MGetFlushState ServiceMethod = 604
 	MDelete        ServiceMethod = 605
 	MQuery         ServiceMethod = 606
+	MUpsert        ServiceMethod = 607
 
 	MManualCompaction            ServiceMethod = 700
 	MGetCompactionState          ServiceMethod = 701
@@ -271,6 +273,16 @@ func (m *MockServer) GetLoadingProgress(ctx context.Context, req *server.GetLoad
 	return &server.GetLoadingProgressResponse{Status: s}, err
 }
 
+func (m *MockServer) GetLoadState(ctx context.Context, req *server.GetLoadStateRequest) (*server.GetLoadStateResponse, error) {
+	f := m.GetInjection(MGetLoadState)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.GetLoadStateResponse), err
+	}
+	s, err := SuccessStatus()
+	return &server.GetLoadStateResponse{Status: s}, err
+}
+
 func (m *MockServer) CreateIndex(ctx context.Context, req *server.CreateIndexRequest) (*common.Status, error) {
 	f := m.GetInjection(MCreateIndex)
 	if f != nil {
@@ -644,3 +656,13 @@ func BadStatus() (*common.Status, error) {
 		Reason:    "fail reason",
 	}, nil
 }
+
+func (m *MockServer) Upsert(ctx context.Context, req *server.UpsertRequest) (*server.MutationResult, error) {
+	f := m.GetInjection(MUpsert)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.MutationResult), err
+	}
+	s, err := SuccessStatus()
+	return &server.MutationResult{Status: s}, err
+}
diff --git a/entity/load.go b/entity/load.go
new file mode 100644
index 0000000..949f325
--- /dev/null
+++ b/entity/load.go
@@ -0,0 +1,21 @@
+// 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 entity
+
+type LoadState int32
+
+const (
+	LoadStateNotExist LoadState = 0
+	LoadStateNotLoad  LoadState = 1
+	LoadStateLoading  LoadState = 2
+	LoadStateLoaded   LoadState = 3
+)
diff --git a/go.mod b/go.mod
index 22874e6..0e9f1fd 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.17
 require (
 	github.com/golang/protobuf v1.5.2
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
-	github.com/milvus-io/milvus-proto/go-api v0.0.0-20221019080323-84e9fa2f9e45
+	github.com/milvus-io/milvus-proto/go-api v0.0.0-20230105121931-9f9303dcc729
 	github.com/stretchr/testify v1.8.1
 	google.golang.org/grpc v1.48.0
 	google.golang.org/grpc/examples v0.0.0-20220617181431-3e7b97febc7f
diff --git a/go.sum b/go.sum
index f6b330e..fbc7caf 100644
--- a/go.sum
+++ b/go.sum
@@ -59,6 +59,8 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxv
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20221019080323-84e9fa2f9e45 h1:QxGQqRtJbbdMf/jxodjYaTW8ZcqdyPDsuhIb9Y2jnwk=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20221019080323-84e9fa2f9e45/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230105121931-9f9303dcc729 h1:hsb1ifdNe3qlXi1YY5dWPPzWMNZmnqe5uunYPYK3gd0=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230105121931-9f9303dcc729/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
 github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -76,10 +78,12 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/vektra/mockery v1.1.2 h1:uc0Yn67rJpjt8U/mAZimdCKn9AeA97BOkjpmtBSlfP4=
+github.com/vektra/mockery v1.1.2/go.mod h1:VcfZjKaFOPO+MpN4ZvwPjs4c48lkq1o3Ym8yHZJu0jU=
+github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
@@ -94,6 +98,7 @@ golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTk
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/mod v0.3.0 h1:RM4zey1++hCTbCVQfnWeKs9/IEsaBLA8vTkd0WVtmH4=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -143,7 +148,9 @@ golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20200323144430-8dcfad9e016e/go.mod h1:Sl4aGygMT6LrqrWclx+PTx3U+LnKx/seiNR+3G19Ar8=
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
+golang.org/x/tools v0.0.0-20210106214847-113979e3529a h1:CB3a9Nez8M13wwlr/E2YtwoU+qYHKfC+JrDa45RXXoQ=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
diff --git a/mocks/MilvusServiceServer.go b/mocks/MilvusServiceServer.go
index 8eb0924..f47261e 100644
--- a/mocks/MilvusServiceServer.go
+++ b/mocks/MilvusServiceServer.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.14.0. DO NOT EDIT.
+// Code generated by mockery v2.15.0. DO NOT EDIT.
 
 package mocks
 
@@ -54,8 +54,8 @@ type MilvusServiceServer_AlterAlias_Call struct {
 }
 
 // AlterAlias is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.AlterAliasRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.AlterAliasRequest
 func (_e *MilvusServiceServer_Expecter) AlterAlias(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_AlterAlias_Call {
 	return &MilvusServiceServer_AlterAlias_Call{Call: _e.mock.On("AlterAlias", _a0, _a1)}
 }
@@ -101,8 +101,8 @@ type MilvusServiceServer_AlterCollection_Call struct {
 }
 
 // AlterCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.AlterCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.AlterCollectionRequest
 func (_e *MilvusServiceServer_Expecter) AlterCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_AlterCollection_Call {
 	return &MilvusServiceServer_AlterCollection_Call{Call: _e.mock.On("AlterCollection", _a0, _a1)}
 }
@@ -148,8 +148,8 @@ type MilvusServiceServer_CalcDistance_Call struct {
 }
 
 // CalcDistance is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CalcDistanceRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CalcDistanceRequest
 func (_e *MilvusServiceServer_Expecter) CalcDistance(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CalcDistance_Call {
 	return &MilvusServiceServer_CalcDistance_Call{Call: _e.mock.On("CalcDistance", _a0, _a1)}
 }
@@ -195,8 +195,8 @@ type MilvusServiceServer_CheckHealth_Call struct {
 }
 
 // CheckHealth is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CheckHealthRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CheckHealthRequest
 func (_e *MilvusServiceServer_Expecter) CheckHealth(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CheckHealth_Call {
 	return &MilvusServiceServer_CheckHealth_Call{Call: _e.mock.On("CheckHealth", _a0, _a1)}
 }
@@ -242,8 +242,8 @@ type MilvusServiceServer_CreateAlias_Call struct {
 }
 
 // CreateAlias is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreateAliasRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreateAliasRequest
 func (_e *MilvusServiceServer_Expecter) CreateAlias(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreateAlias_Call {
 	return &MilvusServiceServer_CreateAlias_Call{Call: _e.mock.On("CreateAlias", _a0, _a1)}
 }
@@ -289,8 +289,8 @@ type MilvusServiceServer_CreateCollection_Call struct {
 }
 
 // CreateCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreateCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreateCollectionRequest
 func (_e *MilvusServiceServer_Expecter) CreateCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreateCollection_Call {
 	return &MilvusServiceServer_CreateCollection_Call{Call: _e.mock.On("CreateCollection", _a0, _a1)}
 }
@@ -336,8 +336,8 @@ type MilvusServiceServer_CreateCredential_Call struct {
 }
 
 // CreateCredential is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreateCredentialRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreateCredentialRequest
 func (_e *MilvusServiceServer_Expecter) CreateCredential(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreateCredential_Call {
 	return &MilvusServiceServer_CreateCredential_Call{Call: _e.mock.On("CreateCredential", _a0, _a1)}
 }
@@ -383,8 +383,8 @@ type MilvusServiceServer_CreateIndex_Call struct {
 }
 
 // CreateIndex is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreateIndexRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreateIndexRequest
 func (_e *MilvusServiceServer_Expecter) CreateIndex(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreateIndex_Call {
 	return &MilvusServiceServer_CreateIndex_Call{Call: _e.mock.On("CreateIndex", _a0, _a1)}
 }
@@ -430,8 +430,8 @@ type MilvusServiceServer_CreatePartition_Call struct {
 }
 
 // CreatePartition is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreatePartitionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreatePartitionRequest
 func (_e *MilvusServiceServer_Expecter) CreatePartition(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreatePartition_Call {
 	return &MilvusServiceServer_CreatePartition_Call{Call: _e.mock.On("CreatePartition", _a0, _a1)}
 }
@@ -477,8 +477,8 @@ type MilvusServiceServer_CreateRole_Call struct {
 }
 
 // CreateRole is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.CreateRoleRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.CreateRoleRequest
 func (_e *MilvusServiceServer_Expecter) CreateRole(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_CreateRole_Call {
 	return &MilvusServiceServer_CreateRole_Call{Call: _e.mock.On("CreateRole", _a0, _a1)}
 }
@@ -524,8 +524,8 @@ type MilvusServiceServer_Delete_Call struct {
 }
 
 // Delete is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DeleteRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DeleteRequest
 func (_e *MilvusServiceServer_Expecter) Delete(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Delete_Call {
 	return &MilvusServiceServer_Delete_Call{Call: _e.mock.On("Delete", _a0, _a1)}
 }
@@ -571,8 +571,8 @@ type MilvusServiceServer_DeleteCredential_Call struct {
 }
 
 // DeleteCredential is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DeleteCredentialRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DeleteCredentialRequest
 func (_e *MilvusServiceServer_Expecter) DeleteCredential(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DeleteCredential_Call {
 	return &MilvusServiceServer_DeleteCredential_Call{Call: _e.mock.On("DeleteCredential", _a0, _a1)}
 }
@@ -618,8 +618,8 @@ type MilvusServiceServer_DescribeCollection_Call struct {
 }
 
 // DescribeCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DescribeCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DescribeCollectionRequest
 func (_e *MilvusServiceServer_Expecter) DescribeCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DescribeCollection_Call {
 	return &MilvusServiceServer_DescribeCollection_Call{Call: _e.mock.On("DescribeCollection", _a0, _a1)}
 }
@@ -665,8 +665,8 @@ type MilvusServiceServer_DescribeIndex_Call struct {
 }
 
 // DescribeIndex is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DescribeIndexRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DescribeIndexRequest
 func (_e *MilvusServiceServer_Expecter) DescribeIndex(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DescribeIndex_Call {
 	return &MilvusServiceServer_DescribeIndex_Call{Call: _e.mock.On("DescribeIndex", _a0, _a1)}
 }
@@ -712,8 +712,8 @@ type MilvusServiceServer_DropAlias_Call struct {
 }
 
 // DropAlias is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DropAliasRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DropAliasRequest
 func (_e *MilvusServiceServer_Expecter) DropAlias(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DropAlias_Call {
 	return &MilvusServiceServer_DropAlias_Call{Call: _e.mock.On("DropAlias", _a0, _a1)}
 }
@@ -759,8 +759,8 @@ type MilvusServiceServer_DropCollection_Call struct {
 }
 
 // DropCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DropCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DropCollectionRequest
 func (_e *MilvusServiceServer_Expecter) DropCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DropCollection_Call {
 	return &MilvusServiceServer_DropCollection_Call{Call: _e.mock.On("DropCollection", _a0, _a1)}
 }
@@ -806,8 +806,8 @@ type MilvusServiceServer_DropIndex_Call struct {
 }
 
 // DropIndex is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DropIndexRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DropIndexRequest
 func (_e *MilvusServiceServer_Expecter) DropIndex(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DropIndex_Call {
 	return &MilvusServiceServer_DropIndex_Call{Call: _e.mock.On("DropIndex", _a0, _a1)}
 }
@@ -853,8 +853,8 @@ type MilvusServiceServer_DropPartition_Call struct {
 }
 
 // DropPartition is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DropPartitionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DropPartitionRequest
 func (_e *MilvusServiceServer_Expecter) DropPartition(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DropPartition_Call {
 	return &MilvusServiceServer_DropPartition_Call{Call: _e.mock.On("DropPartition", _a0, _a1)}
 }
@@ -900,8 +900,8 @@ type MilvusServiceServer_DropRole_Call struct {
 }
 
 // DropRole is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DropRoleRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DropRoleRequest
 func (_e *MilvusServiceServer_Expecter) DropRole(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_DropRole_Call {
 	return &MilvusServiceServer_DropRole_Call{Call: _e.mock.On("DropRole", _a0, _a1)}
 }
@@ -947,8 +947,8 @@ type MilvusServiceServer_Dummy_Call struct {
 }
 
 // Dummy is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.DummyRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.DummyRequest
 func (_e *MilvusServiceServer_Expecter) Dummy(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Dummy_Call {
 	return &MilvusServiceServer_Dummy_Call{Call: _e.mock.On("Dummy", _a0, _a1)}
 }
@@ -994,8 +994,8 @@ type MilvusServiceServer_Flush_Call struct {
 }
 
 // Flush is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.FlushRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.FlushRequest
 func (_e *MilvusServiceServer_Expecter) Flush(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Flush_Call {
 	return &MilvusServiceServer_Flush_Call{Call: _e.mock.On("Flush", _a0, _a1)}
 }
@@ -1041,8 +1041,8 @@ type MilvusServiceServer_GetCollectionStatistics_Call struct {
 }
 
 // GetCollectionStatistics is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetCollectionStatisticsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetCollectionStatisticsRequest
 func (_e *MilvusServiceServer_Expecter) GetCollectionStatistics(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetCollectionStatistics_Call {
 	return &MilvusServiceServer_GetCollectionStatistics_Call{Call: _e.mock.On("GetCollectionStatistics", _a0, _a1)}
 }
@@ -1088,8 +1088,8 @@ type MilvusServiceServer_GetCompactionState_Call struct {
 }
 
 // GetCompactionState is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetCompactionStateRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetCompactionStateRequest
 func (_e *MilvusServiceServer_Expecter) GetCompactionState(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetCompactionState_Call {
 	return &MilvusServiceServer_GetCompactionState_Call{Call: _e.mock.On("GetCompactionState", _a0, _a1)}
 }
@@ -1135,8 +1135,8 @@ type MilvusServiceServer_GetCompactionStateWithPlans_Call struct {
 }
 
 // GetCompactionStateWithPlans is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetCompactionPlansRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetCompactionPlansRequest
 func (_e *MilvusServiceServer_Expecter) GetCompactionStateWithPlans(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetCompactionStateWithPlans_Call {
 	return &MilvusServiceServer_GetCompactionStateWithPlans_Call{Call: _e.mock.On("GetCompactionStateWithPlans", _a0, _a1)}
 }
@@ -1182,8 +1182,8 @@ type MilvusServiceServer_GetComponentStates_Call struct {
 }
 
 // GetComponentStates is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetComponentStatesRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetComponentStatesRequest
 func (_e *MilvusServiceServer_Expecter) GetComponentStates(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetComponentStates_Call {
 	return &MilvusServiceServer_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates", _a0, _a1)}
 }
@@ -1229,8 +1229,8 @@ type MilvusServiceServer_GetFlushState_Call struct {
 }
 
 // GetFlushState is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetFlushStateRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetFlushStateRequest
 func (_e *MilvusServiceServer_Expecter) GetFlushState(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetFlushState_Call {
 	return &MilvusServiceServer_GetFlushState_Call{Call: _e.mock.On("GetFlushState", _a0, _a1)}
 }
@@ -1276,8 +1276,8 @@ type MilvusServiceServer_GetImportState_Call struct {
 }
 
 // GetImportState is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetImportStateRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetImportStateRequest
 func (_e *MilvusServiceServer_Expecter) GetImportState(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetImportState_Call {
 	return &MilvusServiceServer_GetImportState_Call{Call: _e.mock.On("GetImportState", _a0, _a1)}
 }
@@ -1323,8 +1323,8 @@ type MilvusServiceServer_GetIndexBuildProgress_Call struct {
 }
 
 // GetIndexBuildProgress is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetIndexBuildProgressRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetIndexBuildProgressRequest
 func (_e *MilvusServiceServer_Expecter) GetIndexBuildProgress(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetIndexBuildProgress_Call {
 	return &MilvusServiceServer_GetIndexBuildProgress_Call{Call: _e.mock.On("GetIndexBuildProgress", _a0, _a1)}
 }
@@ -1370,8 +1370,8 @@ type MilvusServiceServer_GetIndexState_Call struct {
 }
 
 // GetIndexState is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetIndexStateRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetIndexStateRequest
 func (_e *MilvusServiceServer_Expecter) GetIndexState(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetIndexState_Call {
 	return &MilvusServiceServer_GetIndexState_Call{Call: _e.mock.On("GetIndexState", _a0, _a1)}
 }
@@ -1388,6 +1388,53 @@ func (_c *MilvusServiceServer_GetIndexState_Call) Return(_a0 *milvuspb.GetIndexS
 	return _c
 }
 
+// GetLoadState provides a mock function with given fields: _a0, _a1
+func (_m *MilvusServiceServer) GetLoadState(_a0 context.Context, _a1 *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *milvuspb.GetLoadStateResponse
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetLoadStateRequest) *milvuspb.GetLoadStateResponse); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*milvuspb.GetLoadStateResponse)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.GetLoadStateRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MilvusServiceServer_GetLoadState_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetLoadState'
+type MilvusServiceServer_GetLoadState_Call struct {
+	*mock.Call
+}
+
+// GetLoadState is a helper method to define mock.On call
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetLoadStateRequest
+func (_e *MilvusServiceServer_Expecter) GetLoadState(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetLoadState_Call {
+	return &MilvusServiceServer_GetLoadState_Call{Call: _e.mock.On("GetLoadState", _a0, _a1)}
+}
+
+func (_c *MilvusServiceServer_GetLoadState_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.GetLoadStateRequest)) *MilvusServiceServer_GetLoadState_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*milvuspb.GetLoadStateRequest))
+	})
+	return _c
+}
+
+func (_c *MilvusServiceServer_GetLoadState_Call) Return(_a0 *milvuspb.GetLoadStateResponse, _a1 error) *MilvusServiceServer_GetLoadState_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
 // GetLoadingProgress provides a mock function with given fields: _a0, _a1
 func (_m *MilvusServiceServer) GetLoadingProgress(_a0 context.Context, _a1 *milvuspb.GetLoadingProgressRequest) (*milvuspb.GetLoadingProgressResponse, error) {
 	ret := _m.Called(_a0, _a1)
@@ -1417,8 +1464,8 @@ type MilvusServiceServer_GetLoadingProgress_Call struct {
 }
 
 // GetLoadingProgress is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetLoadingProgressRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetLoadingProgressRequest
 func (_e *MilvusServiceServer_Expecter) GetLoadingProgress(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetLoadingProgress_Call {
 	return &MilvusServiceServer_GetLoadingProgress_Call{Call: _e.mock.On("GetLoadingProgress", _a0, _a1)}
 }
@@ -1464,8 +1511,8 @@ type MilvusServiceServer_GetMetrics_Call struct {
 }
 
 // GetMetrics is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetMetricsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetMetricsRequest
 func (_e *MilvusServiceServer_Expecter) GetMetrics(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetMetrics_Call {
 	return &MilvusServiceServer_GetMetrics_Call{Call: _e.mock.On("GetMetrics", _a0, _a1)}
 }
@@ -1511,8 +1558,8 @@ type MilvusServiceServer_GetPartitionStatistics_Call struct {
 }
 
 // GetPartitionStatistics is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetPartitionStatisticsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetPartitionStatisticsRequest
 func (_e *MilvusServiceServer_Expecter) GetPartitionStatistics(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetPartitionStatistics_Call {
 	return &MilvusServiceServer_GetPartitionStatistics_Call{Call: _e.mock.On("GetPartitionStatistics", _a0, _a1)}
 }
@@ -1558,8 +1605,8 @@ type MilvusServiceServer_GetPersistentSegmentInfo_Call struct {
 }
 
 // GetPersistentSegmentInfo is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetPersistentSegmentInfoRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetPersistentSegmentInfoRequest
 func (_e *MilvusServiceServer_Expecter) GetPersistentSegmentInfo(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetPersistentSegmentInfo_Call {
 	return &MilvusServiceServer_GetPersistentSegmentInfo_Call{Call: _e.mock.On("GetPersistentSegmentInfo", _a0, _a1)}
 }
@@ -1605,8 +1652,8 @@ type MilvusServiceServer_GetQuerySegmentInfo_Call struct {
 }
 
 // GetQuerySegmentInfo is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetQuerySegmentInfoRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetQuerySegmentInfoRequest
 func (_e *MilvusServiceServer_Expecter) GetQuerySegmentInfo(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetQuerySegmentInfo_Call {
 	return &MilvusServiceServer_GetQuerySegmentInfo_Call{Call: _e.mock.On("GetQuerySegmentInfo", _a0, _a1)}
 }
@@ -1652,8 +1699,8 @@ type MilvusServiceServer_GetReplicas_Call struct {
 }
 
 // GetReplicas is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetReplicasRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetReplicasRequest
 func (_e *MilvusServiceServer_Expecter) GetReplicas(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetReplicas_Call {
 	return &MilvusServiceServer_GetReplicas_Call{Call: _e.mock.On("GetReplicas", _a0, _a1)}
 }
@@ -1699,8 +1746,8 @@ type MilvusServiceServer_GetVersion_Call struct {
 }
 
 // GetVersion is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.GetVersionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.GetVersionRequest
 func (_e *MilvusServiceServer_Expecter) GetVersion(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetVersion_Call {
 	return &MilvusServiceServer_GetVersion_Call{Call: _e.mock.On("GetVersion", _a0, _a1)}
 }
@@ -1746,8 +1793,8 @@ type MilvusServiceServer_HasCollection_Call struct {
 }
 
 // HasCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.HasCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.HasCollectionRequest
 func (_e *MilvusServiceServer_Expecter) HasCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_HasCollection_Call {
 	return &MilvusServiceServer_HasCollection_Call{Call: _e.mock.On("HasCollection", _a0, _a1)}
 }
@@ -1793,8 +1840,8 @@ type MilvusServiceServer_HasPartition_Call struct {
 }
 
 // HasPartition is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.HasPartitionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.HasPartitionRequest
 func (_e *MilvusServiceServer_Expecter) HasPartition(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_HasPartition_Call {
 	return &MilvusServiceServer_HasPartition_Call{Call: _e.mock.On("HasPartition", _a0, _a1)}
 }
@@ -1840,8 +1887,8 @@ type MilvusServiceServer_Import_Call struct {
 }
 
 // Import is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ImportRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ImportRequest
 func (_e *MilvusServiceServer_Expecter) Import(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Import_Call {
 	return &MilvusServiceServer_Import_Call{Call: _e.mock.On("Import", _a0, _a1)}
 }
@@ -1887,8 +1934,8 @@ type MilvusServiceServer_Insert_Call struct {
 }
 
 // Insert is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.InsertRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.InsertRequest
 func (_e *MilvusServiceServer_Expecter) Insert(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Insert_Call {
 	return &MilvusServiceServer_Insert_Call{Call: _e.mock.On("Insert", _a0, _a1)}
 }
@@ -1934,8 +1981,8 @@ type MilvusServiceServer_ListCredUsers_Call struct {
 }
 
 // ListCredUsers is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ListCredUsersRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ListCredUsersRequest
 func (_e *MilvusServiceServer_Expecter) ListCredUsers(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ListCredUsers_Call {
 	return &MilvusServiceServer_ListCredUsers_Call{Call: _e.mock.On("ListCredUsers", _a0, _a1)}
 }
@@ -1981,8 +2028,8 @@ type MilvusServiceServer_ListImportTasks_Call struct {
 }
 
 // ListImportTasks is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ListImportTasksRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ListImportTasksRequest
 func (_e *MilvusServiceServer_Expecter) ListImportTasks(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ListImportTasks_Call {
 	return &MilvusServiceServer_ListImportTasks_Call{Call: _e.mock.On("ListImportTasks", _a0, _a1)}
 }
@@ -2028,8 +2075,8 @@ type MilvusServiceServer_LoadBalance_Call struct {
 }
 
 // LoadBalance is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.LoadBalanceRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.LoadBalanceRequest
 func (_e *MilvusServiceServer_Expecter) LoadBalance(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_LoadBalance_Call {
 	return &MilvusServiceServer_LoadBalance_Call{Call: _e.mock.On("LoadBalance", _a0, _a1)}
 }
@@ -2075,8 +2122,8 @@ type MilvusServiceServer_LoadCollection_Call struct {
 }
 
 // LoadCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.LoadCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.LoadCollectionRequest
 func (_e *MilvusServiceServer_Expecter) LoadCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_LoadCollection_Call {
 	return &MilvusServiceServer_LoadCollection_Call{Call: _e.mock.On("LoadCollection", _a0, _a1)}
 }
@@ -2122,8 +2169,8 @@ type MilvusServiceServer_LoadPartitions_Call struct {
 }
 
 // LoadPartitions is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.LoadPartitionsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.LoadPartitionsRequest
 func (_e *MilvusServiceServer_Expecter) LoadPartitions(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_LoadPartitions_Call {
 	return &MilvusServiceServer_LoadPartitions_Call{Call: _e.mock.On("LoadPartitions", _a0, _a1)}
 }
@@ -2169,8 +2216,8 @@ type MilvusServiceServer_ManualCompaction_Call struct {
 }
 
 // ManualCompaction is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ManualCompactionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ManualCompactionRequest
 func (_e *MilvusServiceServer_Expecter) ManualCompaction(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ManualCompaction_Call {
 	return &MilvusServiceServer_ManualCompaction_Call{Call: _e.mock.On("ManualCompaction", _a0, _a1)}
 }
@@ -2216,8 +2263,8 @@ type MilvusServiceServer_OperatePrivilege_Call struct {
 }
 
 // OperatePrivilege is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.OperatePrivilegeRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.OperatePrivilegeRequest
 func (_e *MilvusServiceServer_Expecter) OperatePrivilege(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_OperatePrivilege_Call {
 	return &MilvusServiceServer_OperatePrivilege_Call{Call: _e.mock.On("OperatePrivilege", _a0, _a1)}
 }
@@ -2263,8 +2310,8 @@ type MilvusServiceServer_OperateUserRole_Call struct {
 }
 
 // OperateUserRole is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.OperateUserRoleRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.OperateUserRoleRequest
 func (_e *MilvusServiceServer_Expecter) OperateUserRole(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_OperateUserRole_Call {
 	return &MilvusServiceServer_OperateUserRole_Call{Call: _e.mock.On("OperateUserRole", _a0, _a1)}
 }
@@ -2310,8 +2357,8 @@ type MilvusServiceServer_Query_Call struct {
 }
 
 // Query is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.QueryRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.QueryRequest
 func (_e *MilvusServiceServer_Expecter) Query(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Query_Call {
 	return &MilvusServiceServer_Query_Call{Call: _e.mock.On("Query", _a0, _a1)}
 }
@@ -2357,8 +2404,8 @@ type MilvusServiceServer_RegisterLink_Call struct {
 }
 
 // RegisterLink is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.RegisterLinkRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.RegisterLinkRequest
 func (_e *MilvusServiceServer_Expecter) RegisterLink(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_RegisterLink_Call {
 	return &MilvusServiceServer_RegisterLink_Call{Call: _e.mock.On("RegisterLink", _a0, _a1)}
 }
@@ -2404,8 +2451,8 @@ type MilvusServiceServer_ReleaseCollection_Call struct {
 }
 
 // ReleaseCollection is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ReleaseCollectionRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ReleaseCollectionRequest
 func (_e *MilvusServiceServer_Expecter) ReleaseCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ReleaseCollection_Call {
 	return &MilvusServiceServer_ReleaseCollection_Call{Call: _e.mock.On("ReleaseCollection", _a0, _a1)}
 }
@@ -2451,8 +2498,8 @@ type MilvusServiceServer_ReleasePartitions_Call struct {
 }
 
 // ReleasePartitions is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ReleasePartitionsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ReleasePartitionsRequest
 func (_e *MilvusServiceServer_Expecter) ReleasePartitions(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ReleasePartitions_Call {
 	return &MilvusServiceServer_ReleasePartitions_Call{Call: _e.mock.On("ReleasePartitions", _a0, _a1)}
 }
@@ -2498,8 +2545,8 @@ type MilvusServiceServer_Search_Call struct {
 }
 
 // Search is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.SearchRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.SearchRequest
 func (_e *MilvusServiceServer_Expecter) Search(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Search_Call {
 	return &MilvusServiceServer_Search_Call{Call: _e.mock.On("Search", _a0, _a1)}
 }
@@ -2545,8 +2592,8 @@ type MilvusServiceServer_SelectGrant_Call struct {
 }
 
 // SelectGrant is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.SelectGrantRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.SelectGrantRequest
 func (_e *MilvusServiceServer_Expecter) SelectGrant(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_SelectGrant_Call {
 	return &MilvusServiceServer_SelectGrant_Call{Call: _e.mock.On("SelectGrant", _a0, _a1)}
 }
@@ -2592,8 +2639,8 @@ type MilvusServiceServer_SelectRole_Call struct {
 }
 
 // SelectRole is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.SelectRoleRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.SelectRoleRequest
 func (_e *MilvusServiceServer_Expecter) SelectRole(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_SelectRole_Call {
 	return &MilvusServiceServer_SelectRole_Call{Call: _e.mock.On("SelectRole", _a0, _a1)}
 }
@@ -2639,8 +2686,8 @@ type MilvusServiceServer_SelectUser_Call struct {
 }
 
 // SelectUser is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.SelectUserRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.SelectUserRequest
 func (_e *MilvusServiceServer_Expecter) SelectUser(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_SelectUser_Call {
 	return &MilvusServiceServer_SelectUser_Call{Call: _e.mock.On("SelectUser", _a0, _a1)}
 }
@@ -2686,8 +2733,8 @@ type MilvusServiceServer_ShowCollections_Call struct {
 }
 
 // ShowCollections is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ShowCollectionsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ShowCollectionsRequest
 func (_e *MilvusServiceServer_Expecter) ShowCollections(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ShowCollections_Call {
 	return &MilvusServiceServer_ShowCollections_Call{Call: _e.mock.On("ShowCollections", _a0, _a1)}
 }
@@ -2733,8 +2780,8 @@ type MilvusServiceServer_ShowPartitions_Call struct {
 }
 
 // ShowPartitions is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.ShowPartitionsRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.ShowPartitionsRequest
 func (_e *MilvusServiceServer_Expecter) ShowPartitions(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ShowPartitions_Call {
 	return &MilvusServiceServer_ShowPartitions_Call{Call: _e.mock.On("ShowPartitions", _a0, _a1)}
 }
@@ -2780,8 +2827,8 @@ type MilvusServiceServer_UpdateCredential_Call struct {
 }
 
 // UpdateCredential is a helper method to define mock.On call
-//  - _a0 context.Context
-//  - _a1 *milvuspb.UpdateCredentialRequest
+//   - _a0 context.Context
+//   - _a1 *milvuspb.UpdateCredentialRequest
 func (_e *MilvusServiceServer_Expecter) UpdateCredential(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_UpdateCredential_Call {
 	return &MilvusServiceServer_UpdateCredential_Call{Call: _e.mock.On("UpdateCredential", _a0, _a1)}
 }
@@ -2798,6 +2845,53 @@ func (_c *MilvusServiceServer_UpdateCredential_Call) Return(_a0 *commonpb.Status
 	return _c
 }
 
+// Upsert provides a mock function with given fields: _a0, _a1
+func (_m *MilvusServiceServer) Upsert(_a0 context.Context, _a1 *milvuspb.UpsertRequest) (*milvuspb.MutationResult, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *milvuspb.MutationResult
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.UpsertRequest) *milvuspb.MutationResult); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*milvuspb.MutationResult)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.UpsertRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MilvusServiceServer_Upsert_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Upsert'
+type MilvusServiceServer_Upsert_Call struct {
+	*mock.Call
+}
+
+// Upsert is a helper method to define mock.On call
+//   - _a0 context.Context
+//   - _a1 *milvuspb.UpsertRequest
+func (_e *MilvusServiceServer_Expecter) Upsert(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_Upsert_Call {
+	return &MilvusServiceServer_Upsert_Call{Call: _e.mock.On("Upsert", _a0, _a1)}
+}
+
+func (_c *MilvusServiceServer_Upsert_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.UpsertRequest)) *MilvusServiceServer_Upsert_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*milvuspb.UpsertRequest))
+	})
+	return _c
+}
+
+func (_c *MilvusServiceServer_Upsert_Call) Return(_a0 *milvuspb.MutationResult, _a1 error) *MilvusServiceServer_Upsert_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
 type mockConstructorTestingTNewMilvusServiceServer interface {
 	mock.TestingT
 	Cleanup(func())


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index e322aa4..b99c71d 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -642,3 +642,61 @@ func TestGrpcClientGetReplicas(t *testing.T) {
 
 	mock.DelInjection(MGetReplicas)
 }
+
+func TestGrpcClientGetLoadingProgress(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(ctx, t)
+
+	mock.SetInjection(MHasCollection, hasCollectionDefault)
+
+	mock.SetInjection(MGetLoadingProgress, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.GetLoadingProgressRequest)
+		if !ok {
+			return BadRequestStatus()
+		}
+		resp := &server.GetLoadingProgressResponse{}
+		if !ok {
+			s, err := BadRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		s, err := SuccessStatus()
+		resp.Status, resp.Progress = s, 100
+		return resp, err
+	})
+
+	progress, err := c.GetLoadingProgress(ctx, testCollectionName, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, int64(100), progress)
+
+}
+
+func TestGrpcClientGetLoadState(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(ctx, t)
+
+	mock.SetInjection(MHasCollection, hasCollectionDefault)
+
+	mock.SetInjection(MGetLoadState, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.GetLoadStateRequest)
+		if !ok {
+			return BadRequestStatus()
+		}
+		resp := &server.GetLoadStateResponse{}
+		if !ok {
+			s, err := BadRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		s, err := SuccessStatus()
+		resp.Status, resp.State = s, common.LoadState_LoadStateLoaded
+		return resp, err
+	})
+
+	state, err := c.GetLoadState(ctx, testCollectionName, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, entity.LoadStateLoaded, state)
+
+}