mirror of https://github.com/milvus-io/milvus.git
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 <congqi.xia@zilliz.com>pull/39835/head
congqixia
GitHub
parent
0f4d3ef8b6
commit
7ca5e8309b
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
@ -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))
|
||||
}
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue