From 7ca5e8309bc440d32c3b2291ad96b7b216471422 Mon Sep 17 00:00:00 2001 From: congqixia Date: Wed, 12 Feb 2025 18:46:47 +0800 Subject: [PATCH] enhance: [GoSDK][2.5] Sync gosdk commits from master branch (#39823) Cherry-pick from master pr: #38791 #38812 #38940 #39515 #39580 #39662 #39665 #39751 --------- Signed-off-by: Congqi Xia --- client/entity/database.go | 22 ++ client/entity/rbac.go | 28 +++ client/entity/resource_group.go | 54 ++++ client/entity/segment.go | 34 +++ client/go.mod | 4 +- client/go.sum | 8 +- client/index/ann_param.go | 8 + client/index/sparse.go | 16 ++ client/milvusclient/admin.go | 230 ++++++++++++++++++ client/milvusclient/admin_test.go | 228 +++++++++++++++++ client/milvusclient/client.go | 3 +- client/milvusclient/collection.go | 20 +- .../milvusclient/collection_example_test.go | 2 +- client/milvusclient/collection_options.go | 66 ++++- client/milvusclient/collection_test.go | 64 ++++- client/milvusclient/database.go | 40 +++ client/milvusclient/database_options.go | 78 +++++- client/milvusclient/database_test.go | 83 +++++++ client/milvusclient/maintenance.go | 29 ++- client/milvusclient/read.go | 5 + client/milvusclient/read_options.go | 15 +- client/milvusclient/resource_group.go | 79 ++++++ client/milvusclient/resource_group_option.go | 101 ++++++++ client/milvusclient/resource_group_test.go | 168 +++++++++++++ client/milvusclient/results.go | 1 + tests/go_client/go.mod | 2 +- tests/go_client/go.sum | 4 +- tests/go_client/testcases/partition_test.go | 3 +- tests/go_client/testcases/search_test.go | 9 +- 29 files changed, 1358 insertions(+), 46 deletions(-) create mode 100644 client/entity/database.go create mode 100644 client/entity/resource_group.go create mode 100644 client/entity/segment.go create mode 100644 client/milvusclient/admin.go create mode 100644 client/milvusclient/admin_test.go diff --git a/client/entity/database.go b/client/entity/database.go new file mode 100644 index 0000000000..87cfdbd580 --- /dev/null +++ b/client/entity/database.go @@ -0,0 +1,22 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 Database struct { + Name string + Properties map[string]string +} diff --git a/client/entity/rbac.go b/client/entity/rbac.go index b3c3fa11d2..9ff38961ac 100644 --- a/client/entity/rbac.go +++ b/client/entity/rbac.go @@ -33,3 +33,31 @@ type GrantItem struct { Grantor string Privilege string } + +type UserInfo struct { + UserDescription + Password string +} + +// UserDescription is the model for RBAC user description object. +type UserDescription struct { + Name string + Roles []string +} + +type RBACMeta struct { + Users []*UserInfo + Roles []*Role + RoleGrants []*RoleGrants + PrivilegeGroups []*PrivilegeGroup +} + +// RoleGrants is the model for RBAC role description object. +type RoleGrants struct { + Object string + ObjectName string + RoleName string + GrantorName string + PrivilegeName string + DbName string +} diff --git a/client/entity/resource_group.go b/client/entity/resource_group.go new file mode 100644 index 0000000000..09d5af6c36 --- /dev/null +++ b/client/entity/resource_group.go @@ -0,0 +1,54 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 ResourceGroup struct { + Name string + Capacity int32 + NumAvailableNode int32 + NumLoadedReplica map[string]int32 + NumOutgoingNode map[string]int32 + NumIncomingNode map[string]int32 + Config *ResourceGroupConfig + Nodes []NodeInfo +} + +type NodeInfo struct { + NodeID int64 + Address string + HostName string +} + +type ResourceGroupLimit struct { + NodeNum int32 +} + +type ResourceGroupTransfer struct { + ResourceGroup string +} + +type ResourceGroupNodeFilter struct { + NodeLabels map[string]string +} + +type ResourceGroupConfig struct { + Requests ResourceGroupLimit + Limits ResourceGroupLimit + TransferFrom []*ResourceGroupTransfer + TransferTo []*ResourceGroupTransfer + NodeFilter ResourceGroupNodeFilter +} diff --git a/client/entity/segment.go b/client/entity/segment.go new file mode 100644 index 0000000000..bc3cca4451 --- /dev/null +++ b/client/entity/segment.go @@ -0,0 +1,34 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 + +import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" + +// Segment represent segment in milvus +type Segment struct { + ID int64 + CollectionID int64 + ParititionID int64 + + NumRows int64 + State commonpb.SegmentState +} + +// Flushed indicates segment is flushed +func (s Segment) Flushed() bool { + return s.State == commonpb.SegmentState_Flushed +} diff --git a/client/go.mod b/client/go.mod index 1131d595a5..c36470a723 100644 --- a/client/go.mod +++ b/client/go.mod @@ -7,7 +7,7 @@ require ( github.com/cockroachdb/errors v1.9.1 github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 github.com/milvus-io/milvus-proto/go-api/v2 v2.5.5 - github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 + github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640 github.com/quasilyte/go-ruleguard/dsl v0.3.22 github.com/samber/lo v1.27.0 github.com/stretchr/testify v1.9.0 @@ -99,7 +99,7 @@ require ( go.uber.org/zap v1.27.0 // indirect golang.org/x/crypto v0.31.0 // indirect golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 // indirect - golang.org/x/net v0.27.0 // indirect + golang.org/x/net v0.33.0 // indirect golang.org/x/sync v0.10.0 // indirect golang.org/x/sys v0.28.0 // indirect golang.org/x/text v0.21.0 // indirect diff --git a/client/go.sum b/client/go.sum index 15cbd3a9b9..f0f58a95fc 100644 --- a/client/go.sum +++ b/client/go.sum @@ -320,8 +320,8 @@ github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/le github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/milvus-io/milvus-proto/go-api/v2 v2.5.5 h1:ci2uc+1Es669AM6KDFa/m1kmUKlHo/lidM/8/G6isG8= github.com/milvus-io/milvus-proto/go-api/v2 v2.5.5/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= -github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 h1:EAFxmxUVp5yYFDCrX1MQoSxkTO+ycy8NXEqEDEB3cRM= -github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84/go.mod h1:RATa0GS4jhkPpsYOvQ/QvcNz8rd+TlRPDiSyXQnMMxs= +github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640 h1:vsamolR4+Ru8dDCktU2Q5LCUXd6Soi9jShY9GeG8ZjY= +github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640/go.mod h1:NmwQrk0beEesiZJCwmL1ZrbQBv4DulSbm5mA9/jYPyo= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc= github.com/mitchellh/go-homedir v1.0.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= @@ -634,8 +634,8 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM= golang.org/x/net v0.0.0-20211008194852-3b03d305991f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= -golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys= -golang.org/x/net v0.27.0/go.mod h1:dDi0PyhWNoiUOrAS8uXv/vnScO4wnHQO4mj9fn/RytE= +golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I= +golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= diff --git a/client/index/ann_param.go b/client/index/ann_param.go index 766f84b901..bb6b4a2fb0 100644 --- a/client/index/ann_param.go +++ b/client/index/ann_param.go @@ -32,6 +32,14 @@ func (b baseAnnParam) Params() map[string]any { return b.params } +func (b baseAnnParam) WithRadius(radius float64) { + b.WithExtraParam("radius", radius) +} + +func (b baseAnnParam) WithRangeFilter(rangeFilter float64) { + b.WithExtraParam("range_filter", rangeFilter) +} + type CustomAnnParam struct { baseAnnParam } diff --git a/client/index/sparse.go b/client/index/sparse.go index e835c68bfb..57c87fc19e 100644 --- a/client/index/sparse.go +++ b/client/index/sparse.go @@ -61,3 +61,19 @@ func NewSparseWANDIndex(metricType MetricType, dropRatio float64) Index { dropRatio: dropRatio, } } + +type sparseAnnParam struct { + baseAnnParam +} + +func NewSparseAnnParam() sparseAnnParam { + return sparseAnnParam{ + baseAnnParam: baseAnnParam{ + params: make(map[string]any), + }, + } +} + +func (b sparseAnnParam) WithDropRatio(dropRatio float64) { + b.WithExtraParam("drop_ratio_search", dropRatio) +} diff --git a/client/milvusclient/admin.go b/client/milvusclient/admin.go new file mode 100644 index 0000000000..d7e15fa1ac --- /dev/null +++ b/client/milvusclient/admin.go @@ -0,0 +1,230 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 milvusclient + +import ( + "context" + + "github.com/samber/lo" + "google.golang.org/grpc" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" + "github.com/milvus-io/milvus/pkg/util/merr" +) + +// GetServerVersionOption is the interface for GetServerVersion request. +type GetServerVersionOption interface { + Request() *milvuspb.GetVersionRequest +} + +type getServerVersionOption struct{} + +func (opt *getServerVersionOption) Request() *milvuspb.GetVersionRequest { + return &milvuspb.GetVersionRequest{} +} + +func NewGetServerVersionOption() *getServerVersionOption { + return &getServerVersionOption{} +} + +// GetServerVersion returns connect Milvus instance version. +func (c *Client) GetServerVersion(ctx context.Context, option GetServerVersionOption, callOptions ...grpc.CallOption) (string, error) { + req := option.Request() + + var version string + + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.GetVersion(ctx, req, callOptions...) + version = resp.GetVersion() + return merr.CheckRPCCall(resp, err) + }) + return version, err +} + +type GetPersistentSegmentInfoOption interface { + Request() *milvuspb.GetPersistentSegmentInfoRequest +} + +type getPersistentSegmentInfoOption struct { + collectionName string +} + +func (opt *getPersistentSegmentInfoOption) Request() *milvuspb.GetPersistentSegmentInfoRequest { + return &milvuspb.GetPersistentSegmentInfoRequest{ + CollectionName: opt.collectionName, + } +} + +func NewGetPersistentSegmentInfoOption(collectionName string) GetPersistentSegmentInfoOption { + return &getPersistentSegmentInfoOption{ + collectionName: collectionName, + } +} + +func (c *Client) GetPersistentSegmentInfo(ctx context.Context, option GetPersistentSegmentInfoOption) ([]*entity.Segment, error) { + req := option.Request() + + var segments []*entity.Segment + + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.GetPersistentSegmentInfo(ctx, req) + if err = merr.CheckRPCCall(resp, err); err != nil { + return err + } + + segments = lo.Map(resp.GetInfos(), func(info *milvuspb.PersistentSegmentInfo, _ int) *entity.Segment { + return &entity.Segment{ + ID: info.GetSegmentID(), + CollectionID: info.GetCollectionID(), + ParititionID: info.GetPartitionID(), + NumRows: info.GetNumRows(), + State: info.GetState(), + } + }) + return nil + }) + + return segments, err +} + +type BackupRBACOption interface { + Request() *milvuspb.BackupRBACMetaRequest +} + +type backupRBACOption struct{} + +func (opt *backupRBACOption) Request() *milvuspb.BackupRBACMetaRequest { + return &milvuspb.BackupRBACMetaRequest{} +} + +func NewBackupRBACOption() BackupRBACOption { + return &backupRBACOption{} +} + +func (c *Client) BackupRBAC(ctx context.Context, option BackupRBACOption, callOptions ...grpc.CallOption) (*entity.RBACMeta, error) { + req := option.Request() + + var meta *entity.RBACMeta + + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.BackupRBAC(ctx, req, callOptions...) + if err = merr.CheckRPCCall(resp, err); err != nil { + return err + } + + rbacMeta := resp.GetRBACMeta() + + meta = &entity.RBACMeta{ + Users: lo.Map(rbacMeta.GetUsers(), func(user *milvuspb.UserInfo, _ int) *entity.UserInfo { + return &entity.UserInfo{ + UserDescription: entity.UserDescription{ + Name: user.GetUser(), + Roles: lo.Map(user.GetRoles(), func(role *milvuspb.RoleEntity, _ int) string { return role.GetName() }), + }, + Password: user.GetPassword(), + } + }), + Roles: lo.Map(rbacMeta.GetRoles(), func(role *milvuspb.RoleEntity, _ int) *entity.Role { + return &entity.Role{ + RoleName: role.GetName(), + } + }), + RoleGrants: lo.Map(rbacMeta.GetGrants(), func(grant *milvuspb.GrantEntity, _ int) *entity.RoleGrants { + return &entity.RoleGrants{ + Object: grant.GetObject().GetName(), + ObjectName: grant.GetObjectName(), + RoleName: grant.GetRole().GetName(), + GrantorName: grant.GetGrantor().GetUser().GetName(), + PrivilegeName: grant.GetGrantor().GetPrivilege().GetName(), + DbName: grant.GetDbName(), + } + }), + PrivilegeGroups: lo.Map(rbacMeta.GetPrivilegeGroups(), func(group *milvuspb.PrivilegeGroupInfo, _ int) *entity.PrivilegeGroup { + return &entity.PrivilegeGroup{ + GroupName: group.GetGroupName(), + Privileges: lo.Map(group.GetPrivileges(), func(privilege *milvuspb.PrivilegeEntity, _ int) string { return privilege.GetName() }), + } + }), + } + + return nil + }) + return meta, err +} + +type RestoreRBACOption interface { + Request() *milvuspb.RestoreRBACMetaRequest +} + +type restoreRBACOption struct { + meta *entity.RBACMeta +} + +func (opt *restoreRBACOption) Request() *milvuspb.RestoreRBACMetaRequest { + return &milvuspb.RestoreRBACMetaRequest{ + RBACMeta: &milvuspb.RBACMeta{ + Users: lo.Map(opt.meta.Users, func(user *entity.UserInfo, _ int) *milvuspb.UserInfo { + return &milvuspb.UserInfo{ + User: user.Name, + Roles: lo.Map(user.Roles, func(role string, _ int) *milvuspb.RoleEntity { return &milvuspb.RoleEntity{Name: role} }), + Password: user.Password, + } + }), + Roles: lo.Map(opt.meta.Roles, func(role *entity.Role, _ int) *milvuspb.RoleEntity { + return &milvuspb.RoleEntity{Name: role.RoleName} + }), + Grants: lo.Map(opt.meta.RoleGrants, func(grant *entity.RoleGrants, _ int) *milvuspb.GrantEntity { + return &milvuspb.GrantEntity{ + Object: &milvuspb.ObjectEntity{Name: grant.Object}, + ObjectName: grant.ObjectName, + Role: &milvuspb.RoleEntity{Name: grant.RoleName}, + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{ + Name: grant.GrantorName, + }, + Privilege: &milvuspb.PrivilegeEntity{ + Name: grant.PrivilegeName, + }, + }, + DbName: grant.DbName, + } + }), + PrivilegeGroups: lo.Map(opt.meta.PrivilegeGroups, func(group *entity.PrivilegeGroup, _ int) *milvuspb.PrivilegeGroupInfo { + return &milvuspb.PrivilegeGroupInfo{ + GroupName: group.GroupName, + Privileges: lo.Map(group.Privileges, func(privilege string, _ int) *milvuspb.PrivilegeEntity { + return &milvuspb.PrivilegeEntity{Name: privilege} + }), + } + }), + }, + } +} + +func NewRestoreRBACOption(meta *entity.RBACMeta) RestoreRBACOption { + return &restoreRBACOption{meta: meta} +} + +func (c *Client) RestoreRBAC(ctx context.Context, option RestoreRBACOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.RestoreRBAC(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} diff --git a/client/milvusclient/admin_test.go b/client/milvusclient/admin_test.go new file mode 100644 index 0000000000..fe0bd8f9da --- /dev/null +++ b/client/milvusclient/admin_test.go @@ -0,0 +1,228 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 milvusclient + +import ( + "context" + "fmt" + "math/rand" + "testing" + + "github.com/samber/lo" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/suite" + + "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" + "github.com/milvus-io/milvus/pkg/util/merr" +) + +type AdminSuite struct { + MockSuiteBase +} + +func (s *AdminSuite) TestGetServerVersion() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + version := fmt.Sprintf("v%s", s.randString(6)) + + s.mock.EXPECT().GetVersion(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, gvr *milvuspb.GetVersionRequest) (*milvuspb.GetVersionResponse, error) { + return &milvuspb.GetVersionResponse{ + Status: merr.Success(), + Version: version, + }, nil + }).Once() + + v, err := s.client.GetServerVersion(ctx, NewGetServerVersionOption()) + s.NoError(err) + s.Equal(version, v) + }) + + s.Run("failure", func() { + s.mock.EXPECT().GetVersion(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.GetServerVersion(ctx, NewGetServerVersionOption()) + s.Error(err) + }) +} + +func (s *AdminSuite) TestGetPersistentSegmentInfo() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + collectionName := fmt.Sprintf("coll_%s", s.randString(6)) + segments := []*entity.Segment{ + {ID: rand.Int63(), CollectionID: rand.Int63(), ParititionID: rand.Int63(), NumRows: rand.Int63(), State: commonpb.SegmentState_Flushed}, + {ID: rand.Int63(), CollectionID: rand.Int63(), ParititionID: rand.Int63(), NumRows: rand.Int63(), State: commonpb.SegmentState_Flushed}, + } + + s.mock.EXPECT().GetPersistentSegmentInfo(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, gpsi *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) { + return &milvuspb.GetPersistentSegmentInfoResponse{ + Status: merr.Success(), + Infos: lo.Map(segments, func(segment *entity.Segment, _ int) *milvuspb.PersistentSegmentInfo { + return &milvuspb.PersistentSegmentInfo{ + SegmentID: segment.ID, + CollectionID: segment.CollectionID, + PartitionID: segment.ParititionID, + NumRows: segment.NumRows, + State: segment.State, + } + }), + }, nil + }).Once() + + segments, err := s.client.GetPersistentSegmentInfo(ctx, NewGetPersistentSegmentInfoOption(collectionName)) + s.NoError(err) + s.Equal(segments, segments) + }) + + s.Run("failure", func() { + s.mock.EXPECT().GetPersistentSegmentInfo(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.GetPersistentSegmentInfo(ctx, &getPersistentSegmentInfoOption{collectionName: "coll"}) + s.Error(err) + }) +} + +func (s *AdminSuite) TestBackupRBAC() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + s.mock.EXPECT().BackupRBAC(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, brr *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) { + return &milvuspb.BackupRBACMetaResponse{ + Status: merr.Success(), + RBACMeta: &milvuspb.RBACMeta{ + Users: []*milvuspb.UserInfo{ + { + User: "user1", + Password: "passwd", + Roles: []*milvuspb.RoleEntity{ + {Name: "role1"}, + }, + }, + }, + Roles: []*milvuspb.RoleEntity{ + {Name: "role1"}, + }, + Grants: []*milvuspb.GrantEntity{ + { + Object: &milvuspb.ObjectEntity{ + Name: "testObject", + }, + ObjectName: "testObjectName", + Role: &milvuspb.RoleEntity{ + Name: "testRole", + }, + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{ + Name: "grantorUser", + }, + Privilege: &milvuspb.PrivilegeEntity{ + Name: "testPrivilege", + }, + }, + DbName: "testDB", + }, + }, + PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{ + { + GroupName: "testGroup", + Privileges: []*milvuspb.PrivilegeEntity{ + {Name: "testPrivilege"}, + }, + }, + }, + }, + }, nil + }).Once() + + meta, err := s.client.BackupRBAC(ctx, NewBackupRBACOption()) + s.NoError(err) + s.Len(meta.Users, 1) + s.Len(meta.Roles, 1) + s.Len(meta.RoleGrants, 1) + s.Len(meta.PrivilegeGroups, 1) + }) + + s.Run("failure", func() { + s.mock.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.BackupRBAC(ctx, NewBackupRBACOption()) + s.Error(err) + }) +} + +func (s *AdminSuite) TestRestoreRBAC() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + meta := &entity.RBACMeta{ + Users: []*entity.UserInfo{ + { + UserDescription: entity.UserDescription{ + Name: "user1", + Roles: []string{"role1"}, + }, + Password: "passwd", + }, + }, + Roles: []*entity.Role{ + {RoleName: "role1"}, + }, + RoleGrants: []*entity.RoleGrants{ + { + Object: "testObject", + ObjectName: "testObjectName", + RoleName: "testRole", + GrantorName: "grantorUser", + PrivilegeName: "testPrivilege", + DbName: "testDB", + }, + }, + PrivilegeGroups: []*entity.PrivilegeGroup{ + { + GroupName: "testGroup", + Privileges: []string{"testPrivilege"}, + }, + }, + } + + s.mock.EXPECT().RestoreRBAC(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, rrr *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) { + return merr.Success(), nil + }).Once() + + err := s.client.RestoreRBAC(ctx, NewRestoreRBACOption(meta)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().RestoreRBAC(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.RestoreRBAC(ctx, NewRestoreRBACOption(&entity.RBACMeta{})) + s.Error(err) + }) +} + +func TestAdminAPIs(t *testing.T) { + suite.Run(t, new(AdminSuite)) +} diff --git a/client/milvusclient/client.go b/client/milvusclient/client.go index 97aaa99934..7f4e4a7c52 100644 --- a/client/milvusclient/client.go +++ b/client/milvusclient/client.go @@ -59,7 +59,8 @@ func New(ctx context.Context, config *ClientConfig) (*Client, error) { } c := &Client{ - config: config, + config: config, + currentDB: config.DBName, } // Parse remote address. diff --git a/client/milvusclient/collection.go b/client/milvusclient/collection.go index 025942bcfa..f723916c2b 100644 --- a/client/milvusclient/collection.go +++ b/client/milvusclient/collection.go @@ -139,7 +139,7 @@ func (c *Client) RenameCollection(ctx context.Context, option RenameCollectionOp }) } -func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOption, callOptions ...grpc.CallOption) error { +func (c *Client) AlterCollectionProperties(ctx context.Context, option AlterCollectionPropertiesOption, callOptions ...grpc.CallOption) error { req := option.Request() return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { @@ -148,6 +148,24 @@ func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOpti }) } +func (c *Client) DropCollectionProperties(ctx context.Context, option DropCollectionPropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterCollection(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) AlterCollectionFieldProperty(ctx context.Context, option AlterCollectionFieldPropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterCollectionField(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + type GetCollectionOption interface { Request() *milvuspb.GetCollectionStatisticsRequest } diff --git a/client/milvusclient/collection_example_test.go b/client/milvusclient/collection_example_test.go index 612335cff8..1d573e27e4 100644 --- a/client/milvusclient/collection_example_test.go +++ b/client/milvusclient/collection_example_test.go @@ -260,7 +260,7 @@ func ExampleClient_AlterCollection_setTTL() { defer cli.Close(ctx) - err = cli.AlterCollection(ctx, milvusclient.NewAlterCollectionOption("my_collection").WithProperty(common.CollectionTTLConfigKey, 60)) + err = cli.AlterCollectionProperties(ctx, milvusclient.NewAlterCollectionPropertiesOption("my_collection").WithProperty(common.CollectionTTLConfigKey, 60)) if err != nil { // handle error } diff --git a/client/milvusclient/collection_options.go b/client/milvusclient/collection_options.go index 0907cb6c8f..1bd8029d63 100644 --- a/client/milvusclient/collection_options.go +++ b/client/milvusclient/collection_options.go @@ -286,29 +286,83 @@ func NewRenameCollectionOption(oldName, newName string) *renameCollectionOption } } -type AlterCollectionOption interface { +type AlterCollectionPropertiesOption interface { Request() *milvuspb.AlterCollectionRequest } -type alterCollectionOption struct { +type alterCollectionPropertiesOption struct { collectionName string properties map[string]string } -func (opt *alterCollectionOption) WithProperty(key string, value any) *alterCollectionOption { +func (opt *alterCollectionPropertiesOption) WithProperty(key string, value any) *alterCollectionPropertiesOption { opt.properties[key] = fmt.Sprintf("%v", value) return opt } -func (opt *alterCollectionOption) Request() *milvuspb.AlterCollectionRequest { +func (opt *alterCollectionPropertiesOption) Request() *milvuspb.AlterCollectionRequest { return &milvuspb.AlterCollectionRequest{ CollectionName: opt.collectionName, Properties: entity.MapKvPairs(opt.properties), } } -func NewAlterCollectionOption(collection string) *alterCollectionOption { - return &alterCollectionOption{collectionName: collection, properties: make(map[string]string)} +func NewAlterCollectionPropertiesOption(collection string) *alterCollectionPropertiesOption { + return &alterCollectionPropertiesOption{collectionName: collection, properties: make(map[string]string)} +} + +type DropCollectionPropertiesOption interface { + Request() *milvuspb.AlterCollectionRequest +} + +type dropCollectionPropertiesOption struct { + collectionName string + keys []string +} + +func (opt *dropCollectionPropertiesOption) Request() *milvuspb.AlterCollectionRequest { + return &milvuspb.AlterCollectionRequest{ + CollectionName: opt.collectionName, + DeleteKeys: opt.keys, + } +} + +func NewDropCollectionPropertiesOption(collection string, propertyKeys ...string) *dropCollectionPropertiesOption { + return &dropCollectionPropertiesOption{ + collectionName: collection, + keys: propertyKeys, + } +} + +type AlterCollectionFieldPropertiesOption interface { + Request() *milvuspb.AlterCollectionFieldRequest +} + +type alterCollectionFieldPropertiesOption struct { + collectionName string + fieldName string + properties map[string]string +} + +func (opt *alterCollectionFieldPropertiesOption) WithProperty(key string, value any) *alterCollectionFieldPropertiesOption { + opt.properties[key] = fmt.Sprintf("%v", value) + return opt +} + +func (opt *alterCollectionFieldPropertiesOption) Request() *milvuspb.AlterCollectionFieldRequest { + return &milvuspb.AlterCollectionFieldRequest{ + CollectionName: opt.collectionName, + FieldName: opt.fieldName, + Properties: entity.MapKvPairs(opt.properties), + } +} + +func NewAlterCollectionFieldPropertiesOption(collectionName string, fieldName string) *alterCollectionFieldPropertiesOption { + return &alterCollectionFieldPropertiesOption{ + collectionName: collectionName, + fieldName: fieldName, + properties: make(map[string]string), + } } type getCollectionStatsOption struct { diff --git a/client/milvusclient/collection_test.go b/client/milvusclient/collection_test.go index 4dc8e62d87..b905b24d58 100644 --- a/client/milvusclient/collection_test.go +++ b/client/milvusclient/collection_test.go @@ -285,7 +285,7 @@ func (s *CollectionSuite) TestRenameCollection() { }) } -func (s *CollectionSuite) TestAlterCollection() { +func (s *CollectionSuite) TestAlterCollectionProperties() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -304,14 +304,72 @@ func (s *CollectionSuite) TestAlterCollection() { return merr.Success(), nil }).Once() - err := s.client.AlterCollection(ctx, NewAlterCollectionOption(collName).WithProperty(key, value)) + err := s.client.AlterCollectionProperties(ctx, NewAlterCollectionPropertiesOption(collName).WithProperty(key, value)) s.NoError(err) }) s.Run("failure", func() { s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() - err := s.client.AlterCollection(ctx, NewAlterCollectionOption(collName).WithProperty(key, value)) + err := s.client.AlterCollectionProperties(ctx, NewAlterCollectionPropertiesOption(collName).WithProperty(key, value)) + s.Error(err) + }) +} + +func (s *CollectionSuite) TestDropCollectionProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) { + s.Equal([]string{key}, adr.GetDeleteKeys()) + return merr.Success(), nil + }).Once() + + err := s.client.DropCollectionProperties(ctx, NewDropCollectionPropertiesOption(dbName, key)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.DropCollectionProperties(ctx, NewDropCollectionPropertiesOption(dbName, "key")) + s.Error(err) + }) +} + +func (s *CollectionSuite) TestAlterCollectionFieldProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + collName := fmt.Sprintf("test_collection_%s", s.randString(6)) + fieldName := fmt.Sprintf("field_%s", s.randString(4)) + key := s.randString(6) + value := s.randString(6) + + s.Run("success", func() { + s.mock.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, acr *milvuspb.AlterCollectionFieldRequest) (*commonpb.Status, error) { + s.Equal(collName, acr.GetCollectionName()) + s.Equal(fieldName, acr.GetFieldName()) + if s.Len(acr.GetProperties(), 1) { + item := acr.GetProperties()[0] + s.Equal(key, item.GetKey()) + s.Equal(value, item.GetValue()) + } + return merr.Success(), nil + }).Once() + + err := s.client.AlterCollectionFieldProperty(ctx, NewAlterCollectionFieldPropertiesOption(collName, fieldName).WithProperty(key, value)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.AlterCollectionFieldProperty(ctx, NewAlterCollectionFieldPropertiesOption("coll", "field").WithProperty(key, value)) s.Error(err) }) } diff --git a/client/milvusclient/database.go b/client/milvusclient/database.go index eb5b352963..681dca15a1 100644 --- a/client/milvusclient/database.go +++ b/client/milvusclient/database.go @@ -22,6 +22,7 @@ import ( "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" "github.com/milvus-io/milvus/pkg/util/merr" ) @@ -64,3 +65,42 @@ func (c *Client) DropDatabase(ctx context.Context, option DropDatabaseOption, ca return merr.CheckRPCCall(resp, err) }) } + +func (c *Client) DescribeDatabase(ctx context.Context, option DescribeDatabaseOption, callOptions ...grpc.CallOption) (*entity.Database, error) { + req := option.Request() + + var db *entity.Database + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.DescribeDatabase(ctx, req, callOptions...) + err = merr.CheckRPCCall(resp, err) + if err != nil { + return err + } + // databaseInfo = resp + db = &entity.Database{ + Name: resp.GetDbName(), + Properties: entity.KvPairsMap(resp.GetProperties()), + } + return nil + }) + + return db, err +} + +func (c *Client) AlterDatabaseProperies(ctx context.Context, option AlterDatabasePropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterDatabase(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) DropDatabaseProperties(ctx context.Context, option DropDatabasePropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterDatabase(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} diff --git a/client/milvusclient/database_options.go b/client/milvusclient/database_options.go index 48542f9e58..4d644467ee 100644 --- a/client/milvusclient/database_options.go +++ b/client/milvusclient/database_options.go @@ -16,7 +16,12 @@ package milvusclient -import "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" +import ( + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" +) type UseDatabaseOption interface { DbName() string @@ -90,3 +95,74 @@ func NewDropDatabaseOption(dbName string) *dropDatabaseOption { dbName: dbName, } } + +type DescribeDatabaseOption interface { + Request() *milvuspb.DescribeDatabaseRequest +} + +type describeDatabaseOption struct { + dbName string +} + +func (opt *describeDatabaseOption) Request() *milvuspb.DescribeDatabaseRequest { + return &milvuspb.DescribeDatabaseRequest{ + DbName: opt.dbName, + } +} + +func NewDescribeDatabaseOption(dbName string) *describeDatabaseOption { + return &describeDatabaseOption{ + dbName: dbName, + } +} + +type AlterDatabasePropertiesOption interface { + Request() *milvuspb.AlterDatabaseRequest +} + +type alterDatabasePropertiesOption struct { + dbName string + properties map[string]string +} + +func (opt *alterDatabasePropertiesOption) Request() *milvuspb.AlterDatabaseRequest { + return &milvuspb.AlterDatabaseRequest{ + DbName: opt.dbName, + Properties: entity.MapKvPairs(opt.properties), + } +} + +func (opt *alterDatabasePropertiesOption) WithProperty(key string, value any) *alterDatabasePropertiesOption { + opt.properties[key] = fmt.Sprintf("%v", value) + return opt +} + +func NewAlterDatabasePropertiesOption(dbName string) *alterDatabasePropertiesOption { + return &alterDatabasePropertiesOption{ + dbName: dbName, + properties: make(map[string]string), + } +} + +type DropDatabasePropertiesOption interface { + Request() *milvuspb.AlterDatabaseRequest +} + +type dropDatabasePropertiesOption struct { + dbName string + keys []string +} + +func (opt *dropDatabasePropertiesOption) Request() *milvuspb.AlterDatabaseRequest { + return &milvuspb.AlterDatabaseRequest{ + DbName: opt.dbName, + DeleteKeys: opt.keys, + } +} + +func NewDropDatabasePropertiesOption(dbName string, propertyKeys ...string) *dropDatabasePropertiesOption { + return &dropDatabasePropertiesOption{ + dbName: dbName, + keys: propertyKeys, + } +} diff --git a/client/milvusclient/database_test.go b/client/milvusclient/database_test.go index 3cb3b7017f..541f095d22 100644 --- a/client/milvusclient/database_test.go +++ b/client/milvusclient/database_test.go @@ -108,6 +108,89 @@ func (s *DatabaseSuite) TestUseDatabase() { }) } +func (s *DatabaseSuite) TestDescribeDatabase() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + value := s.randString(6) + s.mock.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(&milvuspb.DescribeDatabaseResponse{ + Status: merr.Success(), + DbName: dbName, + Properties: []*commonpb.KeyValuePair{ + {Key: key, Value: value}, + }, + }, nil).Once() + + db, err := s.client.DescribeDatabase(ctx, NewDescribeDatabaseOption(dbName)) + s.NoError(err) + s.Equal(dbName, db.Name) + s.Equal(value, db.Properties[key]) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.DescribeDatabase(ctx, NewDescribeDatabaseOption(dbName)) + s.Error(err) + }) +} + +func (s *DatabaseSuite) TestAlterDatabaseProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + value := s.randString(6) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterDatabaseRequest) (*commonpb.Status, error) { + s.Equal(dbName, adr.GetDbName()) + s.Len(adr.GetProperties(), 1) + return merr.Success(), nil + }).Once() + + err := s.client.AlterDatabaseProperies(ctx, NewAlterDatabasePropertiesOption(dbName).WithProperty(key, value)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.AlterDatabaseProperies(ctx, NewAlterDatabasePropertiesOption(dbName).WithProperty("key", "value")) + s.Error(err) + }) +} + +func (s *DatabaseSuite) TestDropDatabaseProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterDatabaseRequest) (*commonpb.Status, error) { + s.Equal([]string{key}, adr.GetDeleteKeys()) + return merr.Success(), nil + }).Once() + + err := s.client.DropDatabaseProperties(ctx, NewDropDatabasePropertiesOption(dbName, key)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.DropDatabaseProperties(ctx, NewDropDatabasePropertiesOption(dbName, "key")) + s.Error(err) + }) +} + func TestDatabase(t *testing.T) { suite.Run(t, new(DatabaseSuite)) } diff --git a/client/milvusclient/maintenance.go b/client/milvusclient/maintenance.go index bbf7a636eb..98c70f23f9 100644 --- a/client/milvusclient/maintenance.go +++ b/client/milvusclient/maintenance.go @@ -23,6 +23,7 @@ import ( "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus-proto/go-api/v2/msgpb" "github.com/milvus-io/milvus/client/v2/entity" "github.com/milvus-io/milvus/pkg/util/merr" ) @@ -190,11 +191,13 @@ func (c *Client) RefreshLoad(ctx context.Context, option RefreshLoadOption, call } type FlushTask struct { - client *Client - collectionName string - segmentIDs []int64 - flushTs uint64 - interval time.Duration + client *Client + collectionName string + segmentIDs []int64 + flusheSegIDs []int64 + flushTs uint64 + channelCheckpoints map[string]*msgpb.MsgPosition + interval time.Duration } func (t *FlushTask) Await(ctx context.Context) error { @@ -237,6 +240,10 @@ func (t *FlushTask) Await(ctx context.Context) error { } } +func (t *FlushTask) GetFlushStats() (segIDs []int64, flushSegIDs []int64, flushTs uint64, channelCheckpoints map[string]*msgpb.MsgPosition) { + return t.segmentIDs, t.flusheSegIDs, t.flushTs, t.channelCheckpoints +} + func (c *Client) Flush(ctx context.Context, option FlushOption, callOptions ...grpc.CallOption) (*FlushTask, error) { req := option.Request() collectionName := option.CollectionName() @@ -250,11 +257,13 @@ func (c *Client) Flush(ctx context.Context, option FlushOption, callOptions ...g } task = &FlushTask{ - client: c, - collectionName: collectionName, - segmentIDs: resp.GetCollSegIDs()[collectionName].GetData(), - flushTs: resp.GetCollFlushTs()[collectionName], - interval: option.CheckInterval(), + client: c, + collectionName: collectionName, + segmentIDs: resp.GetCollSegIDs()[collectionName].GetData(), + flusheSegIDs: resp.GetFlushCollSegIDs()[collectionName].GetData(), + flushTs: resp.GetCollFlushTs()[collectionName], + channelCheckpoints: resp.GetChannelCps(), + interval: option.CheckInterval(), } return nil diff --git a/client/milvusclient/read.go b/client/milvusclient/read.go index ab87598c23..6d9a3dc981 100644 --- a/client/milvusclient/read.go +++ b/client/milvusclient/read.go @@ -71,6 +71,11 @@ func (c *Client) handleSearchResult(schema *entity.Schema, outputFields []string sch: schema, } + // set recall if returned + if i < len(results.Recalls) { + entry.Recall = results.Recalls[i] + } + entry.IDs, entry.Err = column.IDColumns(schema, results.GetIds(), offset, offset+rc) if entry.Err != nil { offset += rc diff --git a/client/milvusclient/read_options.go b/client/milvusclient/read_options.go index 89f3d4e3bb..f378cff215 100644 --- a/client/milvusclient/read_options.go +++ b/client/milvusclient/read_options.go @@ -85,6 +85,7 @@ func NewAnnRequest(annField string, limit int, vectors ...entity.Vector) *annReq annField: annField, vectors: vectors, topK: limit, + searchParam: make(map[string]string), templateParams: make(map[string]any), } } @@ -358,12 +359,7 @@ func (opt *searchOption) WithSearchParam(key, value string) *searchOption { func NewSearchOption(collectionName string, limit int, vectors []entity.Vector) *searchOption { return &searchOption{ - annRequest: &annRequest{ - vectors: vectors, - searchParam: make(map[string]string), - topK: limit, - templateParams: make(map[string]any), - }, + annRequest: NewAnnRequest("", limit, vectors...), collectionName: collectionName, useDefaultConsistencyLevel: true, consistencyLevel: entity.ClBounded, @@ -523,9 +519,10 @@ func (opt *queryOption) Request() (*milvuspb.QueryRequest, error) { PartitionNames: opt.partitionNames, OutputFields: opt.outputFields, - Expr: opt.expr, - QueryParams: entity.MapKvPairs(opt.queryParams), - ConsistencyLevel: opt.consistencyLevel.CommonConsistencyLevel(), + Expr: opt.expr, + QueryParams: entity.MapKvPairs(opt.queryParams), + ConsistencyLevel: opt.consistencyLevel.CommonConsistencyLevel(), + UseDefaultConsistency: opt.useDefaultConsistencyLevel, } req.ExprTemplateValues = make(map[string]*schemapb.TemplateValue) diff --git a/client/milvusclient/resource_group.go b/client/milvusclient/resource_group.go index e91654548d..153b3bcbfa 100644 --- a/client/milvusclient/resource_group.go +++ b/client/milvusclient/resource_group.go @@ -19,9 +19,13 @@ package milvusclient import ( "context" + "github.com/samber/lo" "google.golang.org/grpc" + "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus-proto/go-api/v2/rgpb" + "github.com/milvus-io/milvus/client/v2/entity" "github.com/milvus-io/milvus/pkg/util/merr" ) @@ -63,3 +67,78 @@ func (c *Client) DropResourceGroup(ctx context.Context, opt DropResourceGroupOpt return err } + +func (c *Client) DescribeResourceGroup(ctx context.Context, opt DescribeResourceGroupOption, callOptions ...grpc.CallOption) (*entity.ResourceGroup, error) { + req := opt.Request() + + var rg *entity.ResourceGroup + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.DescribeResourceGroup(ctx, req, callOptions...) + if err = merr.CheckRPCCall(resp, err); err != nil { + return err + } + + resultRg := resp.GetResourceGroup() + rg = &entity.ResourceGroup{ + Name: resultRg.GetName(), + Capacity: resultRg.GetCapacity(), + NumAvailableNode: resultRg.GetNumAvailableNode(), + NumLoadedReplica: resultRg.GetNumLoadedReplica(), + NumOutgoingNode: resultRg.GetNumOutgoingNode(), + NumIncomingNode: resultRg.GetNumIncomingNode(), + Config: &entity.ResourceGroupConfig{ + Requests: entity.ResourceGroupLimit{ + NodeNum: resultRg.GetConfig().GetRequests().GetNodeNum(), + }, + Limits: entity.ResourceGroupLimit{ + NodeNum: resultRg.GetConfig().GetLimits().GetNodeNum(), + }, + TransferFrom: lo.Map(resultRg.GetConfig().GetTransferFrom(), func(transfer *rgpb.ResourceGroupTransfer, i int) *entity.ResourceGroupTransfer { + return &entity.ResourceGroupTransfer{ + ResourceGroup: transfer.GetResourceGroup(), + } + }), + TransferTo: lo.Map(resultRg.GetConfig().GetTransferTo(), func(transfer *rgpb.ResourceGroupTransfer, i int) *entity.ResourceGroupTransfer { + return &entity.ResourceGroupTransfer{ + ResourceGroup: transfer.GetResourceGroup(), + } + }), + NodeFilter: entity.ResourceGroupNodeFilter{ + NodeLabels: entity.KvPairsMap(resultRg.GetConfig().GetNodeFilter().GetNodeLabels()), + }, + }, + Nodes: lo.Map(resultRg.GetNodes(), func(node *commonpb.NodeInfo, i int) entity.NodeInfo { + return entity.NodeInfo{ + NodeID: node.GetNodeId(), + Address: node.GetAddress(), + HostName: node.GetHostname(), + } + }), + } + + return nil + }) + return rg, err +} + +func (c *Client) UpdateResourceGroup(ctx context.Context, opt UpdateResourceGroupOption, callOptions ...grpc.CallOption) error { + req := opt.Request() + + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.UpdateResourceGroups(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) + + return err +} + +func (c *Client) TransferReplica(ctx context.Context, opt TransferReplicaOption, callOptions ...grpc.CallOption) error { + req := opt.Request() + + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.TransferReplica(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) + + return err +} diff --git a/client/milvusclient/resource_group_option.go b/client/milvusclient/resource_group_option.go index 5f70f69d0b..6c71405591 100644 --- a/client/milvusclient/resource_group_option.go +++ b/client/milvusclient/resource_group_option.go @@ -17,8 +17,11 @@ package milvusclient import ( + "github.com/samber/lo" + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus-proto/go-api/v2/rgpb" + "github.com/milvus-io/milvus/client/v2/entity" ) type ListResourceGroupsOption interface { @@ -90,3 +93,101 @@ func (opt *dropResourceGroupOption) Request() *milvuspb.DropResourceGroupRequest func NewDropResourceGroupOption(name string) *dropResourceGroupOption { return &dropResourceGroupOption{name: name} } + +type DescribeResourceGroupOption interface { + Request() *milvuspb.DescribeResourceGroupRequest +} + +type describeResourceGroupOption struct { + name string +} + +func (opt *describeResourceGroupOption) Request() *milvuspb.DescribeResourceGroupRequest { + return &milvuspb.DescribeResourceGroupRequest{ + ResourceGroup: opt.name, + } +} + +func NewDescribeResourceGroupOption(name string) *describeResourceGroupOption { + return &describeResourceGroupOption{name: name} +} + +type UpdateResourceGroupOption interface { + Request() *milvuspb.UpdateResourceGroupsRequest +} + +type updateResourceGroupOption struct { + name string + rgConfig *entity.ResourceGroupConfig +} + +func (opt *updateResourceGroupOption) Request() *milvuspb.UpdateResourceGroupsRequest { + return &milvuspb.UpdateResourceGroupsRequest{ + ResourceGroups: map[string]*rgpb.ResourceGroupConfig{ + opt.name: { + Requests: &rgpb.ResourceGroupLimit{ + NodeNum: opt.rgConfig.Requests.NodeNum, + }, + Limits: &rgpb.ResourceGroupLimit{ + NodeNum: opt.rgConfig.Limits.NodeNum, + }, + TransferFrom: lo.Map(opt.rgConfig.TransferFrom, func(transfer *entity.ResourceGroupTransfer, i int) *rgpb.ResourceGroupTransfer { + return &rgpb.ResourceGroupTransfer{ + ResourceGroup: transfer.ResourceGroup, + } + }), + TransferTo: lo.Map(opt.rgConfig.TransferTo, func(transfer *entity.ResourceGroupTransfer, i int) *rgpb.ResourceGroupTransfer { + return &rgpb.ResourceGroupTransfer{ + ResourceGroup: transfer.ResourceGroup, + } + }), + NodeFilter: &rgpb.ResourceGroupNodeFilter{ + NodeLabels: entity.MapKvPairs(opt.rgConfig.NodeFilter.NodeLabels), + }, + }, + }, + } +} + +func NewUpdateResourceGroupOption(name string, resourceGroupConfig *entity.ResourceGroupConfig) *updateResourceGroupOption { + return &updateResourceGroupOption{ + name: name, + rgConfig: resourceGroupConfig, + } +} + +type TransferReplicaOption interface { + Request() *milvuspb.TransferReplicaRequest +} + +type transferReplicaOption struct { + collectionName string + sourceRG string + targetRG string + replicaNum int64 + dbName string +} + +func (opt *transferReplicaOption) WithDBName(dbName string) *transferReplicaOption { + opt.dbName = dbName + return opt +} + +func (opt *transferReplicaOption) Request() *milvuspb.TransferReplicaRequest { + return &milvuspb.TransferReplicaRequest{ + CollectionName: opt.collectionName, + SourceResourceGroup: opt.sourceRG, + TargetResourceGroup: opt.targetRG, + NumReplica: opt.replicaNum, + DbName: opt.dbName, + } +} + +func NewTransferReplicaOption(collectionName, sourceGroup, targetGroup string, replicaNum int64) *transferReplicaOption { + return &transferReplicaOption{ + collectionName: collectionName, + sourceRG: sourceGroup, + targetRG: targetGroup, + replicaNum: replicaNum, + } +} diff --git a/client/milvusclient/resource_group_test.go b/client/milvusclient/resource_group_test.go index 2e87647bf0..757ecf305d 100644 --- a/client/milvusclient/resource_group_test.go +++ b/client/milvusclient/resource_group_test.go @@ -19,14 +19,19 @@ package milvusclient import ( "context" "fmt" + "math/rand" "testing" "github.com/cockroachdb/errors" + "github.com/samber/lo" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus-proto/go-api/v2/rgpb" + "github.com/milvus-io/milvus/client/v2/entity" + "github.com/milvus-io/milvus/pkg/util/merr" ) type ResourceGroupSuite struct { @@ -103,6 +108,169 @@ func (s *ResourceGroupSuite) TestDropResourceGroup() { }) } +func (s *ResourceGroupSuite) TestDescribeResourceGroup() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + limit := rand.Int31n(10) + 1 + request := rand.Int31n(10) + 1 + rgName := fmt.Sprintf("rg_%s", s.randString(6)) + transferFroms := []string{s.randString(6), s.randString(6)} + transferTos := []string{s.randString(6), s.randString(6)} + labels := map[string]string{ + "label1": s.randString(10), + } + node := entity.NodeInfo{ + NodeID: rand.Int63(), + Address: s.randString(6), + HostName: s.randString(10), + } + s.mock.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, drgr *milvuspb.DescribeResourceGroupRequest) (*milvuspb.DescribeResourceGroupResponse, error) { + s.Equal(rgName, drgr.GetResourceGroup()) + return &milvuspb.DescribeResourceGroupResponse{ + ResourceGroup: &milvuspb.ResourceGroup{ + Name: rgName, + Config: &rgpb.ResourceGroupConfig{ + Requests: &rgpb.ResourceGroupLimit{ + NodeNum: request, + }, + Limits: &rgpb.ResourceGroupLimit{ + NodeNum: limit, + }, + TransferFrom: lo.Map(transferFroms, func(transfer string, i int) *rgpb.ResourceGroupTransfer { + return &rgpb.ResourceGroupTransfer{ + ResourceGroup: transfer, + } + }), + TransferTo: lo.Map(transferTos, func(transfer string, i int) *rgpb.ResourceGroupTransfer { + return &rgpb.ResourceGroupTransfer{ + ResourceGroup: transfer, + } + }), + NodeFilter: &rgpb.ResourceGroupNodeFilter{ + NodeLabels: entity.MapKvPairs(labels), + }, + }, + Nodes: []*commonpb.NodeInfo{ + {NodeId: node.NodeID, Address: node.Address, Hostname: node.HostName}, + }, + }, + }, nil + }).Once() + opt := NewDescribeResourceGroupOption(rgName) + rg, err := s.client.DescribeResourceGroup(ctx, opt) + s.NoError(err) + s.Equal(rgName, rg.Name) + s.Equal(limit, rg.Config.Limits.NodeNum) + s.Equal(request, rg.Config.Requests.NodeNum) + s.ElementsMatch(lo.Map(transferFroms, func(transferFrom string, _ int) *entity.ResourceGroupTransfer { + return &entity.ResourceGroupTransfer{ResourceGroup: transferFrom} + }), rg.Config.TransferFrom) + s.ElementsMatch(lo.Map(transferTos, func(transferTo string, _ int) *entity.ResourceGroupTransfer { + return &entity.ResourceGroupTransfer{ResourceGroup: transferTo} + }), rg.Config.TransferTo) + s.Equal(labels, rg.Config.NodeFilter.NodeLabels) + s.ElementsMatch([]entity.NodeInfo{node}, rg.Nodes) + }) + + s.Run("failure", func() { + rgName := fmt.Sprintf("rg_%s", s.randString(6)) + s.mock.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).Return(nil, errors.New("mocked")).Once() + opt := NewDescribeResourceGroupOption(rgName) + _, err := s.client.DescribeResourceGroup(ctx, opt) + s.Error(err) + }) +} + +func (s *ResourceGroupSuite) TestUpdateResourceGroup() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + limit := rand.Int31n(10) + 1 + request := rand.Int31n(10) + 1 + rgName := fmt.Sprintf("rg_%s", s.randString(6)) + transferFroms := []string{s.randString(6), s.randString(6)} + transferTos := []string{s.randString(6), s.randString(6)} + labels := map[string]string{ + "label1": s.randString(10), + } + s.mock.EXPECT().UpdateResourceGroups(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, urgr *milvuspb.UpdateResourceGroupsRequest) (*commonpb.Status, error) { + config, ok := urgr.GetResourceGroups()[rgName] + s.Require().True(ok) + s.Equal(request, config.GetRequests().GetNodeNum()) + s.Equal(limit, config.GetLimits().GetNodeNum()) + s.ElementsMatch(transferFroms, lo.Map(config.GetTransferFrom(), func(transfer *rgpb.ResourceGroupTransfer, i int) string { + return transfer.GetResourceGroup() + })) + s.ElementsMatch(transferTos, lo.Map(config.GetTransferTo(), func(transfer *rgpb.ResourceGroupTransfer, i int) string { + return transfer.GetResourceGroup() + })) + s.Equal(labels, entity.KvPairsMap(config.GetNodeFilter().GetNodeLabels())) + return merr.Success(), nil + }).Once() + opt := NewUpdateResourceGroupOption(rgName, &entity.ResourceGroupConfig{ + Requests: entity.ResourceGroupLimit{NodeNum: request}, + Limits: entity.ResourceGroupLimit{NodeNum: limit}, + TransferFrom: []*entity.ResourceGroupTransfer{ + {ResourceGroup: transferFroms[0]}, + {ResourceGroup: transferFroms[1]}, + }, + TransferTo: []*entity.ResourceGroupTransfer{ + {ResourceGroup: transferTos[0]}, + {ResourceGroup: transferTos[1]}, + }, + NodeFilter: entity.ResourceGroupNodeFilter{ + NodeLabels: labels, + }, + }) + err := s.client.UpdateResourceGroup(ctx, opt) + s.NoError(err) + }) + + s.Run("failure", func() { + rgName := fmt.Sprintf("rg_%s", s.randString(6)) + s.mock.EXPECT().UpdateResourceGroups(mock.Anything, mock.Anything).Return(nil, errors.New("mocked")).Once() + opt := NewUpdateResourceGroupOption(rgName, &entity.ResourceGroupConfig{}) + err := s.client.UpdateResourceGroup(ctx, opt) + s.Error(err) + }) +} + +func (s *ResourceGroupSuite) TestTransferReplica() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + collName := fmt.Sprintf("rg_%s", s.randString(6)) + dbName := fmt.Sprintf("db_%s", s.randString(6)) + from := fmt.Sprintf("rg_%s", s.randString(6)) + to := fmt.Sprintf("rg_%s", s.randString(6)) + replicaNum := rand.Int63n(10) + 1 + s.mock.EXPECT().TransferReplica(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, tr *milvuspb.TransferReplicaRequest) (*commonpb.Status, error) { + s.Equal(collName, tr.GetCollectionName()) + s.Equal(dbName, tr.GetDbName()) + s.Equal(from, tr.GetSourceResourceGroup()) + s.Equal(to, tr.GetTargetResourceGroup()) + return merr.Success(), nil + }).Once() + opt := NewTransferReplicaOption(collName, from, to, replicaNum).WithDBName(dbName) + err := s.client.TransferReplica(ctx, opt) + s.NoError(err) + }) + + s.Run("failure", func() { + rgName := fmt.Sprintf("rg_%s", s.randString(6)) + from := fmt.Sprintf("rg_%s", s.randString(6)) + to := fmt.Sprintf("rg_%s", s.randString(6)) + s.mock.EXPECT().TransferReplica(mock.Anything, mock.Anything).Return(nil, errors.New("mocked")).Once() + opt := NewTransferReplicaOption(rgName, from, to, 1) + err := s.client.TransferReplica(ctx, opt) + s.Error(err) + }) +} + func TestResourceGroup(t *testing.T) { suite.Run(t, new(ResourceGroupSuite)) } diff --git a/client/milvusclient/results.go b/client/milvusclient/results.go index c5ef9e33f7..0f16fecd59 100644 --- a/client/milvusclient/results.go +++ b/client/milvusclient/results.go @@ -37,6 +37,7 @@ type ResultSet struct { IDs column.Column // auto generated id, can be mapped to the columns from `Insert` API Fields DataSet // output field data Scores []float32 // distance to the target vector + Recall float32 // recall of the query vector's search result (estimated by zilliz cloud) Err error // search error if any } diff --git a/tests/go_client/go.mod b/tests/go_client/go.mod index 8cf0df3868..8679c9cf6c 100644 --- a/tests/go_client/go.mod +++ b/tests/go_client/go.mod @@ -6,7 +6,7 @@ toolchain go1.21.11 require ( github.com/milvus-io/milvus/client/v2 v2.0.0-20241125024034-0b9edb62a92d - github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 + github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640 github.com/quasilyte/go-ruleguard/dsl v0.3.22 github.com/stretchr/testify v1.9.0 github.com/x448/float16 v0.8.4 diff --git a/tests/go_client/go.sum b/tests/go_client/go.sum index 317f947d5a..f0f58a95fc 100644 --- a/tests/go_client/go.sum +++ b/tests/go_client/go.sum @@ -320,8 +320,8 @@ github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/le github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/milvus-io/milvus-proto/go-api/v2 v2.5.5 h1:ci2uc+1Es669AM6KDFa/m1kmUKlHo/lidM/8/G6isG8= github.com/milvus-io/milvus-proto/go-api/v2 v2.5.5/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= -github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 h1:EAFxmxUVp5yYFDCrX1MQoSxkTO+ycy8NXEqEDEB3cRM= -github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84/go.mod h1:RATa0GS4jhkPpsYOvQ/QvcNz8rd+TlRPDiSyXQnMMxs= +github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640 h1:vsamolR4+Ru8dDCktU2Q5LCUXd6Soi9jShY9GeG8ZjY= +github.com/milvus-io/milvus/pkg v0.0.2-0.20250212080049-0f4d3ef8b640/go.mod h1:NmwQrk0beEesiZJCwmL1ZrbQBv4DulSbm5mA9/jYPyo= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc= github.com/mitchellh/go-homedir v1.0.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= diff --git a/tests/go_client/testcases/partition_test.go b/tests/go_client/testcases/partition_test.go index 61c1f9e256..de8e6d3818 100644 --- a/tests/go_client/testcases/partition_test.go +++ b/tests/go_client/testcases/partition_test.go @@ -106,7 +106,8 @@ func TestCreatePartitionInvalid(t *testing.T) { } func TestPartitionsNumExceedsMax(t *testing.T) { - ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout) + // 120 seconds may timeout for 1024 partitions + ctx := hp.CreateContext(t, time.Second*300) mc := createDefaultMilvusClient(ctx, t) // create collection diff --git a/tests/go_client/testcases/search_test.go b/tests/go_client/testcases/search_test.go index d38465f4f4..35961891cc 100644 --- a/tests/go_client/testcases/search_test.go +++ b/tests/go_client/testcases/search_test.go @@ -1090,7 +1090,6 @@ func TestSearchSparseVectorNotSupported(t *testing.T) { } func TestRangeSearchSparseVector(t *testing.T) { - t.Skipf("https://github.com/milvus-io/milvus/issues/38846") ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout*2) mc := createDefaultMilvusClient(ctx, t) @@ -1111,10 +1110,12 @@ func TestRangeSearchSparseVector(t *testing.T) { log.Info("default search", zap.Any("score", res.Scores)) } - radius := 10 - rangeFilter := 30 + annParams := index.NewSparseAnnParam() + annParams.WithRadius(10) + annParams.WithRangeFilter(30) + annParams.WithDropRatio(0.2) resRange, errSearch = mc.Search(ctx, client.NewSearchOption(schema.CollectionName, common.DefaultLimit, queryVec). - WithSearchParam("drop_ratio_search", "0.2").WithSearchParam("radius", strconv.Itoa(radius)).WithSearchParam("range_filter", strconv.Itoa(rangeFilter))) + WithAnnParam(annParams)) common.CheckErr(t, errSearch, true) common.CheckErr(t, errSearch, true) require.Len(t, resRange, common.DefaultNq)