-
Notifications
You must be signed in to change notification settings - Fork 105
/
client.go
330 lines (287 loc) · 16.8 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
// Copyright (C) 2019-2021 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
// Package client provides milvus client functions
package client
import (
"context"
"time"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus-sdk-go/v2/entity"
)
// Client is the interface used to communicate with Milvus
type Client interface {
// -- client --
// Close close the remaining connection resources
Close() error
// UsingDatabase for database operation after this function call.
// All request in any goroutine will be applied to new database on the same client. e.g.
// 1. goroutine A access DB1.
// 2. goroutine B call UsingDatabase(ctx, "DB2").
// 3. goroutine A access DB2 after 2.
UsingDatabase(ctx context.Context, dbName string) error
// -- database --
// ListDatabases list all database in milvus cluster.
ListDatabases(ctx context.Context) ([]entity.Database, error)
// CreateDatabase create database with the given name.
CreateDatabase(ctx context.Context, dbName string, opts ...CreateDatabaseOption) error
// DropDatabase drop database with the given db name.
DropDatabase(ctx context.Context, dbName string, opts ...DropDatabaseOption) error
// AlterDatabase alter database props with given db name.
AlterDatabase(ctx context.Context, dbName string, attrs ...entity.DatabaseAttribute) error
DescribeDatabase(ctx context.Context, dbName string) (*entity.Database, error)
// -- collection --
// NewCollection intializeds a new collection with pre defined attributes
NewCollection(ctx context.Context, collName string, dimension int64, opts ...CreateCollectionOption) error
// ListCollections list collections from connection
ListCollections(ctx context.Context, opts ...ListCollectionOption) ([]*entity.Collection, error)
// CreateCollection create collection using provided schema
CreateCollection(ctx context.Context, schema *entity.Schema, shardsNum int32, opts ...CreateCollectionOption) error
// DescribeCollection describe collection meta
DescribeCollection(ctx context.Context, collName string) (*entity.Collection, error)
// DropCollection drop the specified collection
DropCollection(ctx context.Context, collName string, opts ...DropCollectionOption) error
// GetCollectionStatistics get collection statistics
GetCollectionStatistics(ctx context.Context, collName string) (map[string]string, error)
// LoadCollection load collection into memory
LoadCollection(ctx context.Context, collName string, async bool, opts ...LoadCollectionOption) error
// ReleaseCollection release loaded collection
ReleaseCollection(ctx context.Context, collName string, opts ...ReleaseCollectionOption) error
// HasCollection check whether collection exists
HasCollection(ctx context.Context, collName string) (bool, error)
// RenameCollection performs renaming for provided collection.
RenameCollection(ctx context.Context, collName, newName string) error
// AlterCollection changes collection attributes.
AlterCollection(ctx context.Context, collName string, attrs ...entity.CollectionAttribute) error
// CreateAlias creates an alias for collection
CreateAlias(ctx context.Context, collName string, alias string) error
// DropAlias drops the specified Alias
DropAlias(ctx context.Context, alias string) error
// AlterAlias changes collection alias to provided alias
AlterAlias(ctx context.Context, collName string, alias string) error
// GetReplicas gets the replica groups as well as their querynodes and shards information
GetReplicas(ctx context.Context, collName string) ([]*entity.ReplicaGroup, error)
// BackupRBAC back up rbac meta info, include user/role/grants
BackupRBAC(ctx context.Context) (*entity.RBACMeta, error)
RestoreRBAC(ctx context.Context, meta *entity.RBACMeta) error
// -- authentication --
// CreateCredential create new user and password
CreateCredential(ctx context.Context, username string, password string) error
// UpdateCredential update password for a user
UpdateCredential(ctx context.Context, username string, oldPassword string, newPassword string) error
// DeleteCredential delete a user
DeleteCredential(ctx context.Context, username string) error
// ListCredUsers list all usernames
ListCredUsers(ctx context.Context) ([]string, error)
// -- partition --
// CreatePartition create partition for collection
CreatePartition(ctx context.Context, collName string, partitionName string, opts ...CreatePartitionOption) error
// DropPartition drop partition from collection
DropPartition(ctx context.Context, collName string, partitionName string, opts ...DropPartitionOption) error
// ShowPartitions list all partitions from collection
ShowPartitions(ctx context.Context, collName string) ([]*entity.Partition, error)
// HasPartition check whether partition exists in collection
HasPartition(ctx context.Context, collName string, partitionName string) (bool, error)
// LoadPartitions load partitions into memory
LoadPartitions(ctx context.Context, collName string, partitionNames []string, async bool, opts ...LoadPartitionsOption) error
// ReleasePartitions release partitions
ReleasePartitions(ctx context.Context, collName string, partitionNames []string, opts ...ReleasePartitionsOption) error
// -- segment --
GetPersistentSegmentInfo(ctx context.Context, collName string) ([]*entity.Segment, error)
// -- index --
// CreateIndex create index for field of specified collection
// currently index naming is not supported, so only one index on vector field is supported
CreateIndex(ctx context.Context, collName string, fieldName string, idx entity.Index, async bool, opts ...IndexOption) error
// DescribeIndex describe index on collection
// currently index naming is not supported, so only one index on vector field is supported
DescribeIndex(ctx context.Context, collName string, fieldName string, opts ...IndexOption) ([]entity.Index, error)
// DropIndex drop index from collection with specified field name
DropIndex(ctx context.Context, collName string, fieldName string, opts ...IndexOption) error
// GetIndexState get index state with specified collection and field name
// index naming is not supported yet
GetIndexState(ctx context.Context, collName string, fieldName string, opts ...IndexOption) (entity.IndexState, error)
// AlterIndex modifies the index params.
AlterIndex(ctx context.Context, collName, indexName string, opts ...IndexOption) error
// -- basic operation --
// Insert column-based data into collection, returns id column values
Insert(ctx context.Context, collName string, partitionName string, columns ...entity.Column) (entity.Column, error)
// Flush collection, specified
Flush(ctx context.Context, collName string, async bool, opts ...FlushOption) error
// FlushV2 flush collection, specified, return newly sealed segmentIds, all flushed segmentIds of the collection, seal time and error
// currently it is only used in milvus-backup(https://github.com/zilliztech/milvus-backup)
FlushV2(ctx context.Context, collName string, async bool, opts ...FlushOption) ([]int64, []int64, int64, map[string]msgpb.MsgPosition, error)
// DeleteByPks deletes entries related to provided primary keys
DeleteByPks(ctx context.Context, collName string, partitionName string, ids entity.Column) error
// Delete deletes entries match expression
Delete(ctx context.Context, collName string, partitionName string, expr string) error
// Upsert column-based data of collection, returns id column values
Upsert(ctx context.Context, collName string, partitionName string, columns ...entity.Column) (entity.Column, error)
// Search with bool expression
Search(ctx context.Context, collName string, partitions []string,
expr string, outputFields []string, vectors []entity.Vector, vectorField string, metricType entity.MetricType, topK int, sp entity.SearchParam, opts ...SearchQueryOptionFunc) ([]SearchResult, error)
// QueryByPks query record by specified primary key(s).
QueryByPks(ctx context.Context, collectionName string, partitionNames []string, ids entity.Column, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
// Query performs query records with boolean expression.
Query(ctx context.Context, collectionName string, partitionNames []string, expr string, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
// Get grabs the inserted entities using the primary key from the Collection.
Get(ctx context.Context, collectionName string, ids entity.Column, opts ...GetOption) (ResultSet, error)
// SearchIterator returns knn result in iterator mode.
SearchIterator(ctx context.Context, opt *SearchIteratorOption) (*SearchIterator, error)
// QueryIterator returns data matches provided criterion in iterator mode.
QueryIterator(ctx context.Context, opt *QueryIteratorOption) (*QueryIterator, error)
// CalcDistance calculate the distance between vectors specified by ids or provided
CalcDistance(ctx context.Context, collName string, partitions []string,
metricType entity.MetricType, opLeft, opRight entity.Column) (entity.Column, error)
// -- row based apis --
// CreateCollectionByRow create collection by row
CreateCollectionByRow(ctx context.Context, row entity.Row, shardNum int32) error
// DEPRECATED
// InsertByRows insert by rows
InsertByRows(ctx context.Context, collName string, paritionName string, rows []entity.Row) (entity.Column, error)
// InsertRows insert with row base data.
InsertRows(ctx context.Context, collName string, partitionName string, rows []interface{}) (entity.Column, error)
// ManualCompaction triggers a compaction on provided collection
ManualCompaction(ctx context.Context, collName string, toleranceDuration time.Duration) (int64, error)
// GetCompactionState get compaction state of provided compaction id
GetCompactionState(ctx context.Context, id int64) (entity.CompactionState, error)
// GetCompactionStateWithPlans get compaction state with plans of provided compaction id
GetCompactionStateWithPlans(ctx context.Context, id int64) (entity.CompactionState, []entity.CompactionPlan, error)
// BulkInsert import data files(json, numpy, etc.) on MinIO/S3 storage, read and parse them into sealed segments
BulkInsert(ctx context.Context, collName string, partitionName string, files []string, opts ...BulkInsertOption) (int64, error)
// GetBulkInsertState checks import task state
GetBulkInsertState(ctx context.Context, taskID int64) (*entity.BulkInsertTaskState, error)
// ListBulkInsertTasks list state of all import tasks
ListBulkInsertTasks(ctx context.Context, collName string, limit int64) ([]*entity.BulkInsertTaskState, error)
// CreateRole creates a role entity in Milvus.
CreateRole(ctx context.Context, name string) error
// DropRole drops a role entity in Milvus.
DropRole(ctx context.Context, name string) error
// AddUserRole adds one role for user.
AddUserRole(ctx context.Context, username string, role string) error
// RemoveUserRole removes one role from user.
RemoveUserRole(ctx context.Context, username string, role string) error
// ListRoles lists the role objects in system.
ListRoles(ctx context.Context) ([]entity.Role, error)
// ListUsers lists the user objects in system.
ListUsers(ctx context.Context) ([]entity.User, error)
// DescribeUser describes specific user attributes in the system
DescribeUser(ctx context.Context, username string) (entity.UserDescription, error)
// DescribeUsers describe all users attributes in the system
DescribeUsers(ctx context.Context) ([]entity.UserDescription, error)
// ListGrant lists a grant info for the role and the specific object
ListGrant(ctx context.Context, role string, object string, objectName string, dbName string) ([]entity.RoleGrants, error)
// ListGrants lists all assigned privileges and objects for the role.
ListGrants(ctx context.Context, role string, dbName string) ([]entity.RoleGrants, error)
// Grant adds privilege for role.
Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string, options ...entity.OperatePrivilegeOption) error
// Revoke removes privilege from role.
Revoke(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string, options ...entity.OperatePrivilegeOption) error
// GetLoadingProgress get the collection or partitions loading progress
GetLoadingProgress(ctx context.Context, collectionName string, partitionNames []string) (int64, error)
// GetLoadState get the collection or partitions load state
GetLoadState(ctx context.Context, collectionName string, partitionNames []string) (entity.LoadState, error)
// ListResourceGroups returns list of resource group names in current Milvus instance.
ListResourceGroups(ctx context.Context) ([]string, error)
// CreateResourceGroup creates a resource group with provided name.
CreateResourceGroup(ctx context.Context, rgName string, opts ...CreateResourceGroupOption) error
// UpdateResourceGroups updates resource groups with provided options.
UpdateResourceGroups(ctx context.Context, opts ...UpdateResourceGroupsOption) error
// DescribeResourceGroup returns resource groups information.
DescribeResourceGroup(ctx context.Context, rgName string) (*entity.ResourceGroup, error)
// DropResourceGroup drops the resource group with provided name.
DropResourceGroup(ctx context.Context, rgName string) error
// TransferNode transfers querynodes between resource groups.
TransferNode(ctx context.Context, sourceRg, targetRg string, nodesNum int32) error
// TransferReplica transfer collection replicas between source,target resource group.
TransferReplica(ctx context.Context, sourceRg, targetRg string, collectionName string, replicaNum int64) error
// GetVersion get milvus version
GetVersion(ctx context.Context) (string, error)
// CheckHealth returns milvus state
CheckHealth(ctx context.Context) (*entity.MilvusState, error)
ReplicateMessage(ctx context.Context,
channelName string, beginTs, endTs uint64,
msgsBytes [][]byte, startPositions, endPositions []*msgpb.MsgPosition,
opts ...ReplicateMessageOption,
) (*entity.MessageInfo, error)
HybridSearch(ctx context.Context, collName string, partitions []string, limit int, outputFields []string, reranker Reranker, subRequests []*ANNSearchRequest, opts ...SearchQueryOptionFunc) ([]SearchResult, error)
}
// NewClient create a client connected to remote milvus cluster.
// More connect option can be modified by Config.
func NewClient(ctx context.Context, config Config) (Client, error) {
if err := config.parse(); err != nil {
return nil, err
}
c := &GrpcClient{
config: &config,
}
// Parse remote address.
addr := c.config.getParsedAddress()
// Parse grpc options
options := c.config.getDialOption()
// Connect the grpc server.
if err := c.connect(ctx, addr, options...); err != nil {
return nil, err
}
return c, nil
}
// NewGrpcClient create client with grpc addr
// the `Connect` API will be called for you
// dialOptions contains the dial option(s) that control the grpc dialing process
// !!!Deprecated in future, use `NewClient` first.
func NewGrpcClient(ctx context.Context, addr string, dialOptions ...grpc.DialOption) (Client, error) {
return NewClient(ctx, Config{
Address: addr,
DialOptions: dialOptions,
})
}
// NewDefaultGrpcClient creates a new gRPC client.
// !!!Deprecated in future, use `NewClient` first.
func NewDefaultGrpcClient(ctx context.Context, addr string) (Client, error) {
return NewClient(
ctx, Config{
Address: addr,
},
)
}
// NewDefaultGrpcClientWithURI creates a new gRPC client with URI.
// !!!Deprecated in future, use `NewClient` first.
func NewDefaultGrpcClientWithURI(ctx context.Context, uri, username, password string) (Client, error) {
return NewClient(ctx, Config{
Address: uri,
Username: username,
Password: password,
})
}
// NewDefaultGrpcClientWithTLSAuth creates a new gRPC client with TLS authentication.
// !!!Deprecated in future, use `NewClient` first.
func NewDefaultGrpcClientWithTLSAuth(ctx context.Context, addr, username, password string) (Client, error) {
return NewClient(
ctx,
Config{
Address: addr,
Username: username,
Password: password,
EnableTLSAuth: true,
},
)
}
// NewDefaultGrpcClientWithAuth creates a new gRPC client with authentication.
// !!!Deprecated in future, use `NewClient` first.
func NewDefaultGrpcClientWithAuth(ctx context.Context, addr, username, password string) (Client, error) {
return NewClient(
ctx,
Config{
Address: addr,
Username: username,
Password: password,
},
)
}