milvus/client/maintenance_options.go
congqixia 244d2c04f6
feat: Add milvusclient package and migrate GoSDK (#32907)
Related to #31293

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2024-05-10 18:01:30 +08:00

126 lines
3.2 KiB
Go

// 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 client
import (
"time"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
)
type LoadCollectionOption interface {
Request() *milvuspb.LoadCollectionRequest
CheckInterval() time.Duration
}
type loadCollectionOption struct {
collectionName string
interval time.Duration
replicaNum int
}
func (opt *loadCollectionOption) Request() *milvuspb.LoadCollectionRequest {
return &milvuspb.LoadCollectionRequest{
CollectionName: opt.collectionName,
ReplicaNumber: int32(opt.replicaNum),
}
}
func (opt *loadCollectionOption) CheckInterval() time.Duration {
return opt.interval
}
func (opt *loadCollectionOption) WithReplica(num int) *loadCollectionOption {
opt.replicaNum = num
return opt
}
func NewLoadCollectionOption(collectionName string) *loadCollectionOption {
return &loadCollectionOption{
collectionName: collectionName,
replicaNum: 1,
interval: time.Millisecond * 200,
}
}
type LoadPartitionsOption interface {
Request() *milvuspb.LoadPartitionsRequest
CheckInterval() time.Duration
}
var _ LoadPartitionsOption = (*loadPartitionsOption)(nil)
type loadPartitionsOption struct {
collectionName string
partitionNames []string
interval time.Duration
replicaNum int
}
func (opt *loadPartitionsOption) Request() *milvuspb.LoadPartitionsRequest {
return &milvuspb.LoadPartitionsRequest{
CollectionName: opt.collectionName,
PartitionNames: opt.partitionNames,
ReplicaNumber: int32(opt.replicaNum),
}
}
func (opt *loadPartitionsOption) CheckInterval() time.Duration {
return opt.interval
}
func NewLoadPartitionsOption(collectionName string, partitionsNames []string) *loadPartitionsOption {
return &loadPartitionsOption{
collectionName: collectionName,
partitionNames: partitionsNames,
replicaNum: 1,
interval: time.Millisecond * 200,
}
}
type FlushOption interface {
Request() *milvuspb.FlushRequest
CollectionName() string
CheckInterval() time.Duration
}
type flushOption struct {
collectionName string
interval time.Duration
}
func (opt *flushOption) Request() *milvuspb.FlushRequest {
return &milvuspb.FlushRequest{
CollectionNames: []string{opt.collectionName},
}
}
func (opt *flushOption) CollectionName() string {
return opt.collectionName
}
func (opt *flushOption) CheckInterval() time.Duration {
return opt.interval
}
func NewFlushOption(collName string) *flushOption {
return &flushOption{
collectionName: collName,
interval: time.Millisecond * 200,
}
}