impl.go 108.7 KB
Newer Older
1 2 3 4 5 6
// 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
7 8
// with the License. You may obtain a copy of the License at
//
9
//     http://www.apache.org/licenses/LICENSE-2.0
10
//
11 12 13 14 15
// 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.
16

C
Cai Yudong 已提交
17
package proxy
18 19 20

import (
	"context"
21
	"errors"
22
	"fmt"
C
cai.zhang 已提交
23
	"os"
24 25
	"strconv"

26
	"github.com/milvus-io/milvus/internal/common"
27
	"github.com/milvus-io/milvus/internal/logutil"
28

29 30
	"github.com/milvus-io/milvus/internal/util/funcutil"

C
cai.zhang 已提交
31 32
	"github.com/milvus-io/milvus/internal/util/trace"

33 34
	"github.com/milvus-io/milvus/internal/util/metricsinfo"

35
	"go.uber.org/zap"
S
sunby 已提交
36

X
Xiangyu Wang 已提交
37 38 39 40 41 42 43 44
	"github.com/milvus-io/milvus/internal/log"
	"github.com/milvus-io/milvus/internal/msgstream"
	"github.com/milvus-io/milvus/internal/proto/commonpb"
	"github.com/milvus-io/milvus/internal/proto/datapb"
	"github.com/milvus-io/milvus/internal/proto/internalpb"
	"github.com/milvus-io/milvus/internal/proto/milvuspb"
	"github.com/milvus-io/milvus/internal/proto/proxypb"
	"github.com/milvus-io/milvus/internal/proto/querypb"
45
	"github.com/milvus-io/milvus/internal/proto/schemapb"
46
	"github.com/milvus-io/milvus/internal/util/distance"
X
Xiangyu Wang 已提交
47
	"github.com/milvus-io/milvus/internal/util/typeutil"
48 49
)

50 51
const moduleName = "Proxy"

52
// UpdateStateCode updates the state code of Proxy.
C
Cai Yudong 已提交
53
func (node *Proxy) UpdateStateCode(code internalpb.StateCode) {
54
	node.stateCode.Store(code)
Z
zhenshan.cao 已提交
55 56
}

57
// GetComponentStates get state of Proxy.
C
Cai Yudong 已提交
58
func (node *Proxy) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
G
godchen 已提交
59 60 61 62 63 64 65 66 67 68 69 70
	stats := &internalpb.ComponentStates{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		},
	}
	code, ok := node.stateCode.Load().(internalpb.StateCode)
	if !ok {
		errMsg := "unexpected error in type assertion"
		stats.Status = &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    errMsg,
		}
G
godchen 已提交
71
		return stats, nil
G
godchen 已提交
72
	}
73 74 75 76
	nodeID := common.NotRegisteredID
	if node.session != nil && node.session.Registered() {
		nodeID = node.session.ServerID
	}
G
godchen 已提交
77
	info := &internalpb.ComponentInfo{
78 79
		// NodeID:    Params.ProxyID, // will race with Proxy.Register()
		NodeID:    nodeID,
C
Cai Yudong 已提交
80
		Role:      typeutil.ProxyRole,
G
godchen 已提交
81 82 83 84 85 86
		StateCode: code,
	}
	stats.State = info
	return stats, nil
}

C
cxytz01 已提交
87
// GetStatisticsChannel gets statistics channel of Proxy.
C
Cai Yudong 已提交
88
func (node *Proxy) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
G
godchen 已提交
89 90 91 92 93 94 95 96 97
	return &milvuspb.StringResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
			Reason:    "",
		},
		Value: "",
	}, nil
}

98
// InvalidateCollectionMetaCache invalidate the meta cache of specific collection.
C
Cai Yudong 已提交
99
func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
100 101
	ctx = logutil.WithModule(ctx, moduleName)
	logutil.Logger(ctx).Debug("received request to invalidate collection meta cache",
102
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
103 104 105
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

106
	collectionName := request.CollectionName
N
neza2017 已提交
107 108 109
	if globalMetaCache != nil {
		globalMetaCache.RemoveCollection(ctx, collectionName) // no need to return error, though collection may be not cached
	}
110
	logutil.Logger(ctx).Debug("complete to invalidate collection meta cache",
111
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
112 113 114
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

115
	return &commonpb.Status{
116
		ErrorCode: commonpb.ErrorCode_Success,
117 118
		Reason:    "",
	}, nil
119 120
}

121
// ReleaseDQLMessageStream release the query message stream of specific collection.
C
Cai Yudong 已提交
122
func (node *Proxy) ReleaseDQLMessageStream(ctx context.Context, request *proxypb.ReleaseDQLMessageStreamRequest) (*commonpb.Status, error) {
123 124
	ctx = logutil.WithModule(ctx, moduleName)
	logutil.Logger(ctx).Debug("received request to release DQL message strem",
125
		zap.Any("role", typeutil.ProxyRole),
126 127 128
		zap.Any("db", request.DbID),
		zap.Any("collection", request.CollectionID))

129 130 131 132
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

133 134
	_ = node.chMgr.removeDQLStream(request.CollectionID)

135
	logutil.Logger(ctx).Debug("complete to release DQL message stream",
136
		zap.Any("role", typeutil.ProxyRole),
137 138 139 140 141 142 143 144 145
		zap.Any("db", request.DbID),
		zap.Any("collection", request.CollectionID))

	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_Success,
		Reason:    "",
	}, nil
}

146
// TODO(dragondriver): add more detailed ut for ConsistencyLevel, should we support multiple consistency level in Proxy?
147
// CreateCollection create a collection by the schema.
C
Cai Yudong 已提交
148
func (node *Proxy) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
149 150 151
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
152 153 154 155 156

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

157
	cct := &createCollectionTask{
S
sunby 已提交
158
		ctx:                     ctx,
159 160
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
161
		rootCoord:               node.rootCoord,
162 163
	}

164
	method := "CreateCollection"
165 166 167
	// avoid data race
	lenOfSchema := len(request.Schema)

168 169
	log.Debug(
		rpcReceived(method),
170
		zap.String("traceID", traceID),
171
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
172 173
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
174
		zap.Int("len(schema)", lenOfSchema),
175 176
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
177

178 179 180
	if err := node.sched.ddQueue.Enqueue(cct); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
181 182
			zap.Error(err),
			zap.String("traceID", traceID),
183
			zap.String("role", typeutil.ProxyRole),
184 185 186
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Int("len(schema)", lenOfSchema),
187 188
			zap.Int32("shards_num", request.ShardsNum),
			zap.String("consistency_level", request.ConsistencyLevel.String()))
189

190
		return &commonpb.Status{
191
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
192 193 194 195
			Reason:    err.Error(),
		}, nil
	}

196 197
	log.Debug(
		rpcEnqueued(method),
198
		zap.String("traceID", traceID),
199
		zap.String("role", typeutil.ProxyRole),
200 201 202
		zap.Int64("MsgID", cct.ID()),
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
D
dragondriver 已提交
203 204
		zap.Uint64("timestamp", request.Base.Timestamp),
		zap.String("db", request.DbName),
205 206
		zap.String("collection", request.CollectionName),
		zap.Int("len(schema)", lenOfSchema),
207 208
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
209

210 211 212
	if err := cct.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
213
			zap.Error(err),
214
			zap.String("traceID", traceID),
215
			zap.String("role", typeutil.ProxyRole),
216 217 218
			zap.Int64("MsgID", cct.ID()),
			zap.Uint64("BeginTs", cct.BeginTs()),
			zap.Uint64("EndTs", cct.EndTs()),
D
dragondriver 已提交
219 220
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
221
			zap.Int("len(schema)", lenOfSchema),
222 223
			zap.Int32("shards_num", request.ShardsNum),
			zap.String("consistency_level", request.ConsistencyLevel.String()))
D
dragondriver 已提交
224

225
		return &commonpb.Status{
226
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
227 228 229 230
			Reason:    err.Error(),
		}, nil
	}

231 232
	log.Debug(
		rpcDone(method),
233
		zap.String("traceID", traceID),
234
		zap.String("role", typeutil.ProxyRole),
235 236 237 238 239 240
		zap.Int64("MsgID", cct.ID()),
		zap.Uint64("BeginTs", cct.BeginTs()),
		zap.Uint64("EndTs", cct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Int("len(schema)", lenOfSchema),
241 242
		zap.Int32("shards_num", request.ShardsNum),
		zap.String("consistency_level", request.ConsistencyLevel.String()))
243

244 245 246
	return cct.result, nil
}

247
// DropCollection drop a collection.
C
Cai Yudong 已提交
248
func (node *Proxy) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
249 250 251
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
252 253 254 255 256

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

257
	dct := &dropCollectionTask{
S
sunby 已提交
258
		ctx:                   ctx,
259 260
		Condition:             NewTaskCondition(ctx),
		DropCollectionRequest: request,
261
		rootCoord:             node.rootCoord,
262
		chMgr:                 node.chMgr,
S
sunby 已提交
263
		chTicker:              node.chTicker,
264 265
	}

266 267
	log.Debug("DropCollection received",
		zap.String("traceID", traceID),
268
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
269 270
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
271 272 273 274 275

	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DropCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
276
			zap.String("role", typeutil.ProxyRole),
277 278 279
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

280
		return &commonpb.Status{
281
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
282 283 284 285
			Reason:    err.Error(),
		}, nil
	}

286 287
	log.Debug("DropCollection enqueued",
		zap.String("traceID", traceID),
288
		zap.String("role", typeutil.ProxyRole),
289 290 291
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTs", dct.BeginTs()),
		zap.Uint64("EndTs", dct.EndTs()),
D
dragondriver 已提交
292 293
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
294 295 296

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DropCollection failed to WaitToFinish",
D
dragondriver 已提交
297
			zap.Error(err),
298
			zap.String("traceID", traceID),
299
			zap.String("role", typeutil.ProxyRole),
300 301 302
			zap.Int64("MsgID", dct.ID()),
			zap.Uint64("BeginTs", dct.BeginTs()),
			zap.Uint64("EndTs", dct.EndTs()),
D
dragondriver 已提交
303 304 305
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

306
		return &commonpb.Status{
307
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
308 309 310 311
			Reason:    err.Error(),
		}, nil
	}

312 313
	log.Debug("DropCollection done",
		zap.String("traceID", traceID),
314
		zap.String("role", typeutil.ProxyRole),
315 316 317 318 319 320
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTs", dct.BeginTs()),
		zap.Uint64("EndTs", dct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

321 322 323
	return dct.result, nil
}

324
// HasCollection check if the specific collection exists in Milvus.
C
Cai Yudong 已提交
325
func (node *Proxy) HasCollection(ctx context.Context, request *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) {
326 327 328 329 330
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
331 332 333 334 335 336 337

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

	log.Debug("HasCollection received",
		zap.String("traceID", traceID),
338
		zap.String("role", typeutil.ProxyRole),
339 340 341
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

342
	hct := &hasCollectionTask{
S
sunby 已提交
343
		ctx:                  ctx,
344 345
		Condition:            NewTaskCondition(ctx),
		HasCollectionRequest: request,
346
		rootCoord:            node.rootCoord,
347 348
	}

349 350 351 352
	if err := node.sched.ddQueue.Enqueue(hct); err != nil {
		log.Warn("HasCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
353
			zap.String("role", typeutil.ProxyRole),
354 355 356
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

357 358
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
359
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
360 361 362 363 364
				Reason:    err.Error(),
			},
		}, nil
	}

365 366
	log.Debug("HasCollection enqueued",
		zap.String("traceID", traceID),
367
		zap.String("role", typeutil.ProxyRole),
368 369 370
		zap.Int64("MsgID", hct.ID()),
		zap.Uint64("BeginTS", hct.BeginTs()),
		zap.Uint64("EndTS", hct.EndTs()),
D
dragondriver 已提交
371 372
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
373 374 375

	if err := hct.WaitToFinish(); err != nil {
		log.Warn("HasCollection failed to WaitToFinish",
D
dragondriver 已提交
376
			zap.Error(err),
377
			zap.String("traceID", traceID),
378
			zap.String("role", typeutil.ProxyRole),
379 380 381
			zap.Int64("MsgID", hct.ID()),
			zap.Uint64("BeginTS", hct.BeginTs()),
			zap.Uint64("EndTS", hct.EndTs()),
D
dragondriver 已提交
382 383 384
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

385 386
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
387
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
388 389 390 391 392
				Reason:    err.Error(),
			},
		}, nil
	}

393 394
	log.Debug("HasCollection done",
		zap.String("traceID", traceID),
395
		zap.String("role", typeutil.ProxyRole),
396 397 398 399 400 401
		zap.Int64("MsgID", hct.ID()),
		zap.Uint64("BeginTS", hct.BeginTs()),
		zap.Uint64("EndTS", hct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

402 403 404
	return hct.result, nil
}

405
// LoadCollection load a collection into query nodes.
C
Cai Yudong 已提交
406
func (node *Proxy) LoadCollection(ctx context.Context, request *milvuspb.LoadCollectionRequest) (*commonpb.Status, error) {
407 408 409
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
410 411 412 413 414

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadCollection")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

415
	lct := &loadCollectionTask{
S
sunby 已提交
416
		ctx:                   ctx,
417 418
		Condition:             NewTaskCondition(ctx),
		LoadCollectionRequest: request,
419
		queryCoord:            node.queryCoord,
420 421
	}

422 423
	log.Debug("LoadCollection received",
		zap.String("traceID", traceID),
424
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
425 426
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
427 428 429 430 431

	if err := node.sched.ddQueue.Enqueue(lct); err != nil {
		log.Warn("LoadCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
432
			zap.String("role", typeutil.ProxyRole),
433 434 435
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

436
		return &commonpb.Status{
437
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
438 439 440
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
441

442 443
	log.Debug("LoadCollection enqueued",
		zap.String("traceID", traceID),
444
		zap.String("role", typeutil.ProxyRole),
445 446 447
		zap.Int64("MsgID", lct.ID()),
		zap.Uint64("BeginTS", lct.BeginTs()),
		zap.Uint64("EndTS", lct.EndTs()),
D
dragondriver 已提交
448 449
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
450 451 452

	if err := lct.WaitToFinish(); err != nil {
		log.Warn("LoadCollection failed to WaitToFinish",
D
dragondriver 已提交
453
			zap.Error(err),
454
			zap.String("traceID", traceID),
455
			zap.String("role", typeutil.ProxyRole),
456 457 458
			zap.Int64("MsgID", lct.ID()),
			zap.Uint64("BeginTS", lct.BeginTs()),
			zap.Uint64("EndTS", lct.EndTs()),
D
dragondriver 已提交
459 460 461
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

462
		return &commonpb.Status{
463
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
464 465 466 467
			Reason:    err.Error(),
		}, nil
	}

468 469
	log.Debug("LoadCollection done",
		zap.String("traceID", traceID),
470
		zap.String("role", typeutil.ProxyRole),
471 472 473 474 475 476
		zap.Int64("MsgID", lct.ID()),
		zap.Uint64("BeginTS", lct.BeginTs()),
		zap.Uint64("EndTS", lct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

477
	return lct.result, nil
478 479
}

480
// ReleaseCollection remove the loaded collection from query nodes.
C
Cai Yudong 已提交
481
func (node *Proxy) ReleaseCollection(ctx context.Context, request *milvuspb.ReleaseCollectionRequest) (*commonpb.Status, error) {
482 483 484
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
485

486
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ReleaseCollection")
487 488 489
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

490
	rct := &releaseCollectionTask{
S
sunby 已提交
491
		ctx:                      ctx,
492 493
		Condition:                NewTaskCondition(ctx),
		ReleaseCollectionRequest: request,
494
		queryCoord:               node.queryCoord,
495
		chMgr:                    node.chMgr,
496 497
	}

498 499 500 501
	method := "ReleaseCollection"

	log.Debug(
		rpcReceived(method),
502
		zap.String("traceID", traceID),
503
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
504 505
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
506 507

	if err := node.sched.ddQueue.Enqueue(rct); err != nil {
508 509
		log.Warn(
			rpcFailedToEnqueue(method),
510 511
			zap.Error(err),
			zap.String("traceID", traceID),
512
			zap.String("role", typeutil.ProxyRole),
513 514 515
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

516
		return &commonpb.Status{
517
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
518 519 520 521
			Reason:    err.Error(),
		}, nil
	}

522 523
	log.Debug(
		rpcEnqueued(method),
524
		zap.String("traceID", traceID),
525
		zap.String("role", typeutil.ProxyRole),
526 527 528
		zap.Int64("MsgID", rct.ID()),
		zap.Uint64("BeginTS", rct.BeginTs()),
		zap.Uint64("EndTS", rct.EndTs()),
D
dragondriver 已提交
529 530
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
531 532

	if err := rct.WaitToFinish(); err != nil {
533 534
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
535
			zap.Error(err),
536
			zap.String("traceID", traceID),
537
			zap.String("role", typeutil.ProxyRole),
538 539 540
			zap.Int64("MsgID", rct.ID()),
			zap.Uint64("BeginTS", rct.BeginTs()),
			zap.Uint64("EndTS", rct.EndTs()),
D
dragondriver 已提交
541 542 543
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

544
		return &commonpb.Status{
545
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
546 547 548 549
			Reason:    err.Error(),
		}, nil
	}

550 551
	log.Debug(
		rpcDone(method),
552
		zap.String("traceID", traceID),
553
		zap.String("role", typeutil.ProxyRole),
554 555 556 557 558 559
		zap.Int64("MsgID", rct.ID()),
		zap.Uint64("BeginTS", rct.BeginTs()),
		zap.Uint64("EndTS", rct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

560
	return rct.result, nil
561 562
}

563
// DescribeCollection get the meta information of specific collection, such as schema, created timestamp and etc.
C
Cai Yudong 已提交
564
func (node *Proxy) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
565 566 567 568 569
	if !node.checkHealthy() {
		return &milvuspb.DescribeCollectionResponse{
			Status: unhealthyStatus(),
		}, nil
	}
570

571
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DescribeCollection")
572 573 574
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

575
	dct := &describeCollectionTask{
S
sunby 已提交
576
		ctx:                       ctx,
577 578
		Condition:                 NewTaskCondition(ctx),
		DescribeCollectionRequest: request,
579
		rootCoord:                 node.rootCoord,
580 581
	}

582 583
	log.Debug("DescribeCollection received",
		zap.String("traceID", traceID),
584
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
585 586
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
587 588 589 590 591

	if err := node.sched.ddQueue.Enqueue(dct); err != nil {
		log.Warn("DescribeCollection failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
592
			zap.String("role", typeutil.ProxyRole),
593 594 595
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

596 597
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
598
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
599 600 601 602 603
				Reason:    err.Error(),
			},
		}, nil
	}

604 605
	log.Debug("DescribeCollection enqueued",
		zap.String("traceID", traceID),
606
		zap.String("role", typeutil.ProxyRole),
607 608 609
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTS", dct.BeginTs()),
		zap.Uint64("EndTS", dct.EndTs()),
D
dragondriver 已提交
610 611
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
612 613 614

	if err := dct.WaitToFinish(); err != nil {
		log.Warn("DescribeCollection failed to WaitToFinish",
D
dragondriver 已提交
615
			zap.Error(err),
616
			zap.String("traceID", traceID),
617
			zap.String("role", typeutil.ProxyRole),
618 619 620
			zap.Int64("MsgID", dct.ID()),
			zap.Uint64("BeginTS", dct.BeginTs()),
			zap.Uint64("EndTS", dct.EndTs()),
D
dragondriver 已提交
621 622 623
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

624 625
		return &milvuspb.DescribeCollectionResponse{
			Status: &commonpb.Status{
626
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
627 628 629 630 631
				Reason:    err.Error(),
			},
		}, nil
	}

632 633
	log.Debug("DescribeCollection done",
		zap.String("traceID", traceID),
634
		zap.String("role", typeutil.ProxyRole),
635 636 637 638 639 640
		zap.Int64("MsgID", dct.ID()),
		zap.Uint64("BeginTS", dct.BeginTs()),
		zap.Uint64("EndTS", dct.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

641 642 643
	return dct.result, nil
}

644
// GetCollectionStatistics get the collection statistics, such as `num_rows`.
C
Cai Yudong 已提交
645
func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
646 647 648 649 650
	if !node.checkHealthy() {
		return &milvuspb.GetCollectionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
651 652 653 654 655

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetCollectionStatistics")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

656
	g := &getCollectionStatisticsTask{
G
godchen 已提交
657 658 659
		ctx:                            ctx,
		Condition:                      NewTaskCondition(ctx),
		GetCollectionStatisticsRequest: request,
660
		dataCoord:                      node.dataCoord,
661 662
	}

663 664
	log.Debug("GetCollectionStatistics received",
		zap.String("traceID", traceID),
665
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
666 667
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
668 669 670 671 672

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn("GetCollectionStatistics failed to enqueue",
			zap.Error(err),
			zap.String("traceID", traceID),
673
			zap.String("role", typeutil.ProxyRole),
674 675 676
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

G
godchen 已提交
677
		return &milvuspb.GetCollectionStatisticsResponse{
678
			Status: &commonpb.Status{
679
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
680 681 682 683 684
				Reason:    err.Error(),
			},
		}, nil
	}

685 686
	log.Debug("GetCollectionStatistics enqueued",
		zap.String("traceID", traceID),
687
		zap.String("role", typeutil.ProxyRole),
688 689 690
		zap.Int64("MsgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
D
dragondriver 已提交
691 692
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))
693 694 695

	if err := g.WaitToFinish(); err != nil {
		log.Warn("GetCollectionStatistics failed to WaitToFinish",
D
dragondriver 已提交
696
			zap.Error(err),
697
			zap.String("traceID", traceID),
698
			zap.String("role", typeutil.ProxyRole),
699 700 701
			zap.Int64("MsgID", g.ID()),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
D
dragondriver 已提交
702 703 704
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName))

G
godchen 已提交
705
		return &milvuspb.GetCollectionStatisticsResponse{
706
			Status: &commonpb.Status{
707
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
708 709 710 711 712
				Reason:    err.Error(),
			},
		}, nil
	}

713 714
	log.Debug("GetCollectionStatistics done",
		zap.String("traceID", traceID),
715
		zap.String("role", typeutil.ProxyRole),
716 717 718 719 720 721
		zap.Int64("MsgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName))

722
	return g.result, nil
723 724
}

725
// ShowCollections list all collections in Milvus.
C
Cai Yudong 已提交
726
func (node *Proxy) ShowCollections(ctx context.Context, request *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
727 728 729 730 731
	if !node.checkHealthy() {
		return &milvuspb.ShowCollectionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
732
	sct := &showCollectionsTask{
G
godchen 已提交
733 734 735
		ctx:                    ctx,
		Condition:              NewTaskCondition(ctx),
		ShowCollectionsRequest: request,
736
		queryCoord:             node.queryCoord,
737
		rootCoord:              node.rootCoord,
738 739
	}

740
	log.Debug("ShowCollections received",
741
		zap.String("role", typeutil.ProxyRole),
742 743 744 745 746 747
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
		zap.String("ShowType", request.Type.String()),
		zap.Any("CollectionNames", request.CollectionNames),
	)

748
	err := node.sched.ddQueue.Enqueue(sct)
749
	if err != nil {
750 751
		log.Warn("ShowCollections failed to enqueue",
			zap.Error(err),
752
			zap.String("role", typeutil.ProxyRole),
753 754 755 756 757 758
			zap.String("DbName", request.DbName),
			zap.Uint64("TimeStamp", request.TimeStamp),
			zap.String("ShowType", request.Type.String()),
			zap.Any("CollectionNames", request.CollectionNames),
		)

G
godchen 已提交
759
		return &milvuspb.ShowCollectionsResponse{
760
			Status: &commonpb.Status{
761
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
762 763 764 765 766
				Reason:    err.Error(),
			},
		}, nil
	}

767
	log.Debug("ShowCollections enqueued",
768
		zap.String("role", typeutil.ProxyRole),
769
		zap.Int64("MsgID", sct.ID()),
770
		zap.String("DbName", sct.ShowCollectionsRequest.DbName),
771
		zap.Uint64("TimeStamp", request.TimeStamp),
772 773 774
		zap.String("ShowType", sct.ShowCollectionsRequest.Type.String()),
		zap.Any("CollectionNames", sct.ShowCollectionsRequest.CollectionNames),
	)
D
dragondriver 已提交
775

776 777
	err = sct.WaitToFinish()
	if err != nil {
778 779
		log.Warn("ShowCollections failed to WaitToFinish",
			zap.Error(err),
780
			zap.String("role", typeutil.ProxyRole),
781 782 783 784 785 786 787
			zap.Int64("MsgID", sct.ID()),
			zap.String("DbName", request.DbName),
			zap.Uint64("TimeStamp", request.TimeStamp),
			zap.String("ShowType", request.Type.String()),
			zap.Any("CollectionNames", request.CollectionNames),
		)

G
godchen 已提交
788
		return &milvuspb.ShowCollectionsResponse{
789
			Status: &commonpb.Status{
790
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
791 792 793 794 795
				Reason:    err.Error(),
			},
		}, nil
	}

796
	log.Debug("ShowCollections Done",
797
		zap.String("role", typeutil.ProxyRole),
798 799 800 801 802 803 804 805
		zap.Int64("MsgID", sct.ID()),
		zap.String("DbName", request.DbName),
		zap.Uint64("TimeStamp", request.TimeStamp),
		zap.String("ShowType", request.Type.String()),
		zap.Any("CollectionNames", request.CollectionNames),
		zap.Any("result", sct.result),
	)

806 807 808
	return sct.result, nil
}

809
// CreatePartition create a partition in specific collection.
C
Cai Yudong 已提交
810
func (node *Proxy) CreatePartition(ctx context.Context, request *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) {
811 812 813
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
814

815
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreatePartition")
816 817 818
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

819
	cpt := &createPartitionTask{
S
sunby 已提交
820
		ctx:                    ctx,
821 822
		Condition:              NewTaskCondition(ctx),
		CreatePartitionRequest: request,
823
		rootCoord:              node.rootCoord,
824 825 826
		result:                 nil,
	}

827 828 829
	log.Debug(
		rpcReceived("CreatePartition"),
		zap.String("traceID", traceID),
830
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
831 832 833
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
834 835 836 837 838 839

	if err := node.sched.ddQueue.Enqueue(cpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue("CreatePartition"),
			zap.Error(err),
			zap.String("traceID", traceID),
840
			zap.String("role", typeutil.ProxyRole),
841 842 843 844
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

845
		return &commonpb.Status{
846
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
847 848 849
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
850

851 852 853
	log.Debug(
		rpcEnqueued("CreatePartition"),
		zap.String("traceID", traceID),
854
		zap.String("role", typeutil.ProxyRole),
855 856 857
		zap.Int64("MsgID", cpt.ID()),
		zap.Uint64("BeginTS", cpt.BeginTs()),
		zap.Uint64("EndTS", cpt.EndTs()),
D
dragondriver 已提交
858 859 860
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
861 862 863 864

	if err := cpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish("CreatePartition"),
D
dragondriver 已提交
865
			zap.Error(err),
866
			zap.String("traceID", traceID),
867
			zap.String("role", typeutil.ProxyRole),
868 869 870
			zap.Int64("MsgID", cpt.ID()),
			zap.Uint64("BeginTS", cpt.BeginTs()),
			zap.Uint64("EndTS", cpt.EndTs()),
D
dragondriver 已提交
871 872 873 874
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

875
		return &commonpb.Status{
876
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
877 878 879
			Reason:    err.Error(),
		}, nil
	}
880 881 882 883

	log.Debug(
		rpcDone("CreatePartition"),
		zap.String("traceID", traceID),
884
		zap.String("role", typeutil.ProxyRole),
885 886 887 888 889 890 891
		zap.Int64("MsgID", cpt.ID()),
		zap.Uint64("BeginTS", cpt.BeginTs()),
		zap.Uint64("EndTS", cpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

892 893 894
	return cpt.result, nil
}

895
// DropPartition drop a partition in specific collection.
C
Cai Yudong 已提交
896
func (node *Proxy) DropPartition(ctx context.Context, request *milvuspb.DropPartitionRequest) (*commonpb.Status, error) {
897 898 899
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
900

901
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropPartition")
902 903 904
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

905
	dpt := &dropPartitionTask{
S
sunby 已提交
906
		ctx:                  ctx,
907 908
		Condition:            NewTaskCondition(ctx),
		DropPartitionRequest: request,
909
		rootCoord:            node.rootCoord,
910 911 912
		result:               nil,
	}

913 914 915 916 917
	method := "DropPartition"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
918
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
919 920 921
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
922 923 924 925 926 927

	if err := node.sched.ddQueue.Enqueue(dpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
928
			zap.String("role", typeutil.ProxyRole),
929 930 931 932
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

933
		return &commonpb.Status{
934
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
935 936 937
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
938

939 940 941
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
942
		zap.String("role", typeutil.ProxyRole),
943 944 945
		zap.Int64("MsgID", dpt.ID()),
		zap.Uint64("BeginTS", dpt.BeginTs()),
		zap.Uint64("EndTS", dpt.EndTs()),
D
dragondriver 已提交
946 947 948
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
949 950 951 952

	if err := dpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
953
			zap.Error(err),
954
			zap.String("traceID", traceID),
955
			zap.String("role", typeutil.ProxyRole),
956 957 958
			zap.Int64("MsgID", dpt.ID()),
			zap.Uint64("BeginTS", dpt.BeginTs()),
			zap.Uint64("EndTS", dpt.EndTs()),
D
dragondriver 已提交
959 960 961 962
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

963
		return &commonpb.Status{
964
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
965 966 967
			Reason:    err.Error(),
		}, nil
	}
968 969 970 971

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
972
		zap.String("role", typeutil.ProxyRole),
973 974 975 976 977 978 979
		zap.Int64("MsgID", dpt.ID()),
		zap.Uint64("BeginTS", dpt.BeginTs()),
		zap.Uint64("EndTS", dpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

980 981 982
	return dpt.result, nil
}

983
// HasPartition check if partition exist.
C
Cai Yudong 已提交
984
func (node *Proxy) HasPartition(ctx context.Context, request *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) {
985 986 987 988 989
	if !node.checkHealthy() {
		return &milvuspb.BoolResponse{
			Status: unhealthyStatus(),
		}, nil
	}
D
dragondriver 已提交
990

D
dragondriver 已提交
991
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-HasPartition")
D
dragondriver 已提交
992 993 994
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

995
	hpt := &hasPartitionTask{
S
sunby 已提交
996
		ctx:                 ctx,
997 998
		Condition:           NewTaskCondition(ctx),
		HasPartitionRequest: request,
999
		rootCoord:           node.rootCoord,
1000 1001 1002
		result:              nil,
	}

D
dragondriver 已提交
1003 1004 1005 1006 1007
	method := "HasPartition"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1008
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1009 1010 1011
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1012 1013 1014 1015 1016 1017

	if err := node.sched.ddQueue.Enqueue(hpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1018
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1019 1020 1021 1022
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1023 1024
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1025
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1026 1027 1028 1029 1030
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1031

D
dragondriver 已提交
1032 1033 1034
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1035
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1036 1037 1038
		zap.Int64("MsgID", hpt.ID()),
		zap.Uint64("BeginTS", hpt.BeginTs()),
		zap.Uint64("EndTS", hpt.EndTs()),
D
dragondriver 已提交
1039 1040 1041
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
D
dragondriver 已提交
1042 1043 1044 1045

	if err := hpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1046
			zap.Error(err),
D
dragondriver 已提交
1047
			zap.String("traceID", traceID),
1048
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1049 1050 1051
			zap.Int64("MsgID", hpt.ID()),
			zap.Uint64("BeginTS", hpt.BeginTs()),
			zap.Uint64("EndTS", hpt.EndTs()),
D
dragondriver 已提交
1052 1053 1054 1055
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1056 1057
		return &milvuspb.BoolResponse{
			Status: &commonpb.Status{
1058
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1059 1060 1061 1062 1063
				Reason:    err.Error(),
			},
			Value: false,
		}, nil
	}
D
dragondriver 已提交
1064 1065 1066 1067

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1068
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1069 1070 1071 1072 1073 1074 1075
		zap.Int64("MsgID", hpt.ID()),
		zap.Uint64("BeginTS", hpt.BeginTs()),
		zap.Uint64("EndTS", hpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1076 1077 1078
	return hpt.result, nil
}

1079
// LoadPartitions load specific partitions into query nodes.
C
Cai Yudong 已提交
1080
func (node *Proxy) LoadPartitions(ctx context.Context, request *milvuspb.LoadPartitionsRequest) (*commonpb.Status, error) {
1081 1082 1083
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1084

D
dragondriver 已提交
1085
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-LoadPartitions")
1086 1087 1088
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1089
	lpt := &loadPartitionsTask{
G
godchen 已提交
1090 1091 1092
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		LoadPartitionsRequest: request,
1093
		queryCoord:            node.queryCoord,
1094 1095
	}

1096 1097 1098 1099 1100
	method := "LoadPartitions"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1101
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1102 1103 1104
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1105 1106 1107 1108 1109 1110

	if err := node.sched.ddQueue.Enqueue(lpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1111
			zap.String("role", typeutil.ProxyRole),
1112 1113 1114 1115
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1116
		return &commonpb.Status{
1117
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1118 1119 1120 1121
			Reason:    err.Error(),
		}, nil
	}

1122 1123 1124
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1125
		zap.String("role", typeutil.ProxyRole),
1126 1127 1128
		zap.Int64("MsgID", lpt.ID()),
		zap.Uint64("BeginTS", lpt.BeginTs()),
		zap.Uint64("EndTS", lpt.EndTs()),
D
dragondriver 已提交
1129 1130 1131
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1132 1133 1134 1135

	if err := lpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1136
			zap.Error(err),
1137
			zap.String("traceID", traceID),
1138
			zap.String("role", typeutil.ProxyRole),
1139 1140 1141
			zap.Int64("MsgID", lpt.ID()),
			zap.Uint64("BeginTS", lpt.BeginTs()),
			zap.Uint64("EndTS", lpt.EndTs()),
D
dragondriver 已提交
1142 1143 1144 1145
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1146
		return &commonpb.Status{
1147
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1148 1149 1150 1151
			Reason:    err.Error(),
		}, nil
	}

1152 1153 1154
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1155
		zap.String("role", typeutil.ProxyRole),
1156 1157 1158 1159 1160 1161 1162
		zap.Int64("MsgID", lpt.ID()),
		zap.Uint64("BeginTS", lpt.BeginTs()),
		zap.Uint64("EndTS", lpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))

1163
	return lpt.result, nil
1164 1165
}

1166
// ReleasePartitions release specific partitions from query nodes.
C
Cai Yudong 已提交
1167
func (node *Proxy) ReleasePartitions(ctx context.Context, request *milvuspb.ReleasePartitionsRequest) (*commonpb.Status, error) {
1168 1169 1170
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
1171 1172 1173 1174 1175

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ReleasePartitions")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1176
	rpt := &releasePartitionsTask{
G
godchen 已提交
1177 1178 1179
		ctx:                      ctx,
		Condition:                NewTaskCondition(ctx),
		ReleasePartitionsRequest: request,
1180
		queryCoord:               node.queryCoord,
1181 1182
	}

1183 1184 1185 1186 1187
	method := "ReleasePartitions"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1188
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1189 1190 1191
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1192 1193 1194 1195 1196 1197

	if err := node.sched.ddQueue.Enqueue(rpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1198
			zap.String("role", typeutil.ProxyRole),
1199 1200 1201 1202
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1203
		return &commonpb.Status{
1204
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1205 1206 1207 1208
			Reason:    err.Error(),
		}, nil
	}

1209 1210 1211
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1212
		zap.String("role", typeutil.ProxyRole),
1213 1214 1215
		zap.Int64("msgID", rpt.Base.MsgID),
		zap.Uint64("BeginTS", rpt.BeginTs()),
		zap.Uint64("EndTS", rpt.EndTs()),
D
dragondriver 已提交
1216 1217 1218
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
1219 1220 1221 1222

	if err := rpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1223
			zap.Error(err),
1224
			zap.String("traceID", traceID),
1225
			zap.String("role", typeutil.ProxyRole),
1226 1227 1228
			zap.Int64("msgID", rpt.Base.MsgID),
			zap.Uint64("BeginTS", rpt.BeginTs()),
			zap.Uint64("EndTS", rpt.EndTs()),
D
dragondriver 已提交
1229 1230 1231 1232
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))

1233
		return &commonpb.Status{
1234
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1235 1236 1237 1238
			Reason:    err.Error(),
		}, nil
	}

1239 1240 1241
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1242
		zap.String("role", typeutil.ProxyRole),
1243 1244 1245 1246 1247 1248 1249
		zap.Int64("msgID", rpt.Base.MsgID),
		zap.Uint64("BeginTS", rpt.BeginTs()),
		zap.Uint64("EndTS", rpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))

1250
	return rpt.result, nil
1251 1252
}

1253
// GetPartitionStatistics get the statistics of partition, such as num_rows.
C
Cai Yudong 已提交
1254
func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb.GetPartitionStatisticsRequest) (*milvuspb.GetPartitionStatisticsResponse, error) {
1255 1256 1257 1258 1259
	if !node.checkHealthy() {
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1260 1261 1262 1263 1264

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetPartitionStatistics")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1265
	g := &getPartitionStatisticsTask{
1266 1267 1268
		ctx:                           ctx,
		Condition:                     NewTaskCondition(ctx),
		GetPartitionStatisticsRequest: request,
1269
		dataCoord:                     node.dataCoord,
1270 1271
	}

1272 1273 1274 1275 1276
	method := "GetPartitionStatistics"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1277
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1278 1279 1280
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1281 1282 1283 1284 1285 1286

	if err := node.sched.ddQueue.Enqueue(g); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1287
			zap.String("role", typeutil.ProxyRole),
1288 1289 1290 1291
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

1292 1293 1294 1295 1296 1297 1298 1299
		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1300 1301 1302
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1303
		zap.String("role", typeutil.ProxyRole),
1304 1305 1306
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
1307 1308 1309
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))
1310 1311 1312 1313

	if err := g.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1314
			zap.Error(err),
1315
			zap.String("traceID", traceID),
1316
			zap.String("role", typeutil.ProxyRole),
1317 1318 1319
			zap.Int64("msgID", g.ID()),
			zap.Uint64("BeginTS", g.BeginTs()),
			zap.Uint64("EndTS", g.EndTs()),
1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("partition", request.PartitionName))

		return &milvuspb.GetPartitionStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1332 1333 1334
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1335
		zap.String("role", typeutil.ProxyRole),
1336 1337 1338 1339 1340 1341 1342
		zap.Int64("msgID", g.ID()),
		zap.Uint64("BeginTS", g.BeginTs()),
		zap.Uint64("EndTS", g.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName))

1343
	return g.result, nil
1344 1345
}

1346
// ShowPartitions list all partitions in the specific collection.
C
Cai Yudong 已提交
1347
func (node *Proxy) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
1348 1349 1350 1351 1352
	if !node.checkHealthy() {
		return &milvuspb.ShowPartitionsResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1353 1354 1355 1356 1357

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ShowPartitions")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1358
	spt := &showPartitionsTask{
G
godchen 已提交
1359 1360 1361
		ctx:                   ctx,
		Condition:             NewTaskCondition(ctx),
		ShowPartitionsRequest: request,
1362
		rootCoord:             node.rootCoord,
1363
		queryCoord:            node.queryCoord,
G
godchen 已提交
1364
		result:                nil,
1365 1366
	}

1367 1368 1369 1370 1371
	method := "ShowPartitions"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1372
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1373
		zap.Any("request", request))
1374 1375 1376 1377 1378 1379

	if err := node.sched.ddQueue.Enqueue(spt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1380
			zap.String("role", typeutil.ProxyRole),
1381 1382
			zap.Any("request", request))

G
godchen 已提交
1383
		return &milvuspb.ShowPartitionsResponse{
1384
			Status: &commonpb.Status{
1385
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1386 1387 1388 1389 1390
				Reason:    err.Error(),
			},
		}, nil
	}

1391 1392 1393
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1394
		zap.String("role", typeutil.ProxyRole),
1395 1396 1397
		zap.Int64("msgID", spt.ID()),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
1398 1399
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
1400 1401 1402 1403 1404
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

	if err := spt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1405
			zap.Error(err),
1406
			zap.String("traceID", traceID),
1407
			zap.String("role", typeutil.ProxyRole),
1408 1409 1410 1411 1412 1413
			zap.Int64("msgID", spt.ID()),
			zap.Uint64("BeginTS", spt.BeginTs()),
			zap.Uint64("EndTS", spt.EndTs()),
			zap.String("db", spt.ShowPartitionsRequest.DbName),
			zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
			zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))
D
dragondriver 已提交
1414

G
godchen 已提交
1415
		return &milvuspb.ShowPartitionsResponse{
1416
			Status: &commonpb.Status{
1417
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1418 1419 1420 1421
				Reason:    err.Error(),
			},
		}, nil
	}
1422 1423 1424 1425

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1426
		zap.String("role", typeutil.ProxyRole),
1427 1428 1429 1430 1431 1432 1433
		zap.Int64("msgID", spt.ID()),
		zap.Uint64("BeginTS", spt.BeginTs()),
		zap.Uint64("EndTS", spt.EndTs()),
		zap.String("db", spt.ShowPartitionsRequest.DbName),
		zap.String("collection", spt.ShowPartitionsRequest.CollectionName),
		zap.Any("partitions", spt.ShowPartitionsRequest.PartitionNames))

1434 1435 1436
	return spt.result, nil
}

1437
// CreateIndex create index for collection.
C
Cai Yudong 已提交
1438
func (node *Proxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
1439 1440 1441
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1442 1443 1444 1445 1446

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-ShowPartitions")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1447
	cit := &createIndexTask{
S
sunby 已提交
1448
		ctx:                ctx,
1449 1450
		Condition:          NewTaskCondition(ctx),
		CreateIndexRequest: request,
1451
		rootCoord:          node.rootCoord,
1452 1453
	}

D
dragondriver 已提交
1454 1455 1456 1457 1458
	method := "CreateIndex"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1459
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1460 1461 1462 1463
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1464 1465 1466 1467 1468 1469

	if err := node.sched.ddQueue.Enqueue(cit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1470
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1471 1472 1473 1474 1475
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.Any("extra_params", request.ExtraParams))

1476
		return &commonpb.Status{
1477
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1478 1479 1480 1481
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1482 1483 1484
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1485
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1486 1487 1488
		zap.Int64("MsgID", cit.ID()),
		zap.Uint64("BeginTs", cit.BeginTs()),
		zap.Uint64("EndTs", cit.EndTs()),
D
dragondriver 已提交
1489 1490 1491 1492
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))
D
dragondriver 已提交
1493 1494 1495 1496

	if err := cit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1497
			zap.Error(err),
D
dragondriver 已提交
1498
			zap.String("traceID", traceID),
1499
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1500 1501 1502
			zap.Int64("MsgID", cit.ID()),
			zap.Uint64("BeginTs", cit.BeginTs()),
			zap.Uint64("EndTs", cit.EndTs()),
D
dragondriver 已提交
1503 1504 1505 1506 1507
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.Any("extra_params", request.ExtraParams))

1508
		return &commonpb.Status{
1509
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
1510 1511 1512 1513
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
1514 1515 1516
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1517
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1518 1519 1520 1521 1522 1523 1524 1525
		zap.Int64("MsgID", cit.ID()),
		zap.Uint64("BeginTs", cit.BeginTs()),
		zap.Uint64("EndTs", cit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.Any("extra_params", request.ExtraParams))

1526 1527 1528
	return cit.result, nil
}

1529
// DescribeIndex get the meta information of index, such as index state, index id and etc.
C
Cai Yudong 已提交
1530
func (node *Proxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
1531 1532 1533 1534 1535
	if !node.checkHealthy() {
		return &milvuspb.DescribeIndexResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1536 1537 1538 1539 1540

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DescribeIndex")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1541
	dit := &describeIndexTask{
S
sunby 已提交
1542
		ctx:                  ctx,
1543 1544
		Condition:            NewTaskCondition(ctx),
		DescribeIndexRequest: request,
1545
		rootCoord:            node.rootCoord,
1546 1547
	}

1548 1549 1550 1551 1552 1553 1554
	method := "DescribeIndex"
	// avoid data race
	indexName := request.IndexName

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1555
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1556 1557 1558
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
1559 1560 1561 1562 1563 1564 1565
		zap.String("index name", indexName))

	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1566
			zap.String("role", typeutil.ProxyRole),
1567 1568 1569 1570 1571
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", indexName))

1572 1573
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
1574
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1575 1576 1577 1578 1579
				Reason:    err.Error(),
			},
		}, nil
	}

1580 1581 1582
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1583
		zap.String("role", typeutil.ProxyRole),
1584 1585 1586
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
1587 1588 1589
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
1590 1591 1592 1593 1594
		zap.String("index name", indexName))

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1595
			zap.Error(err),
1596
			zap.String("traceID", traceID),
1597
			zap.String("role", typeutil.ProxyRole),
1598 1599 1600
			zap.Int64("MsgID", dit.ID()),
			zap.Uint64("BeginTs", dit.BeginTs()),
			zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
1601 1602 1603
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
1604
			zap.String("index name", indexName))
D
dragondriver 已提交
1605

Z
zhenshan.cao 已提交
1606 1607 1608 1609
		errCode := commonpb.ErrorCode_UnexpectedError
		if dit.result != nil {
			errCode = dit.result.Status.GetErrorCode()
		}
1610 1611
		return &milvuspb.DescribeIndexResponse{
			Status: &commonpb.Status{
Z
zhenshan.cao 已提交
1612
				ErrorCode: errCode,
1613 1614 1615 1616 1617
				Reason:    err.Error(),
			},
		}, nil
	}

1618 1619 1620
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1621
		zap.String("role", typeutil.ProxyRole),
1622 1623 1624 1625 1626 1627 1628 1629
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", indexName))

1630 1631 1632
	return dit.result, nil
}

1633
// DropIndex drop the index of collection.
C
Cai Yudong 已提交
1634
func (node *Proxy) DropIndex(ctx context.Context, request *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
1635 1636 1637
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
1638 1639 1640 1641 1642

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropIndex")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1643
	dit := &dropIndexTask{
S
sunby 已提交
1644
		ctx:              ctx,
B
BossZou 已提交
1645 1646
		Condition:        NewTaskCondition(ctx),
		DropIndexRequest: request,
1647
		rootCoord:        node.rootCoord,
B
BossZou 已提交
1648
	}
G
godchen 已提交
1649

D
dragondriver 已提交
1650 1651 1652 1653 1654
	method := "DropIndex"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1655
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1656 1657 1658 1659 1660
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

D
dragondriver 已提交
1661 1662 1663 1664 1665
	if err := node.sched.ddQueue.Enqueue(dit); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1666
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1667 1668 1669 1670 1671
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

B
BossZou 已提交
1672
		return &commonpb.Status{
1673
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1674 1675 1676
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1677

D
dragondriver 已提交
1678 1679 1680
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1681
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1682 1683 1684
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
1685 1686 1687 1688
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
D
dragondriver 已提交
1689 1690 1691 1692

	if err := dit.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1693
			zap.Error(err),
D
dragondriver 已提交
1694
			zap.String("traceID", traceID),
1695
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1696 1697 1698
			zap.Int64("MsgID", dit.ID()),
			zap.Uint64("BeginTs", dit.BeginTs()),
			zap.Uint64("EndTs", dit.EndTs()),
D
dragondriver 已提交
1699 1700 1701 1702 1703
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

B
BossZou 已提交
1704
		return &commonpb.Status{
1705
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
B
BossZou 已提交
1706 1707 1708
			Reason:    err.Error(),
		}, nil
	}
D
dragondriver 已提交
1709 1710 1711 1712

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1713
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1714 1715 1716 1717 1718 1719 1720 1721
		zap.Int64("MsgID", dit.ID()),
		zap.Uint64("BeginTs", dit.BeginTs()),
		zap.Uint64("EndTs", dit.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

B
BossZou 已提交
1722 1723 1724
	return dit.result, nil
}

1725 1726
// GetIndexBuildProgress gets index build progress with filed_name and index_name.
// IndexRows is the num of indexed rows. And TotalRows is the total number of segment rows.
C
Cai Yudong 已提交
1727
func (node *Proxy) GetIndexBuildProgress(ctx context.Context, request *milvuspb.GetIndexBuildProgressRequest) (*milvuspb.GetIndexBuildProgressResponse, error) {
1728 1729 1730 1731 1732
	if !node.checkHealthy() {
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1733 1734 1735 1736 1737

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-GetIndexBuildProgress")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1738
	gibpt := &getIndexBuildProgressTask{
1739 1740 1741
		ctx:                          ctx,
		Condition:                    NewTaskCondition(ctx),
		GetIndexBuildProgressRequest: request,
1742 1743
		indexCoord:                   node.indexCoord,
		rootCoord:                    node.rootCoord,
1744
		dataCoord:                    node.dataCoord,
1745 1746
	}

1747 1748 1749 1750 1751
	method := "GetIndexBuildProgress"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1752
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1753 1754 1755 1756
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1757 1758 1759 1760 1761 1762

	if err := node.sched.ddQueue.Enqueue(gibpt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1763
			zap.String("role", typeutil.ProxyRole),
1764 1765 1766 1767 1768
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

1769 1770 1771 1772 1773 1774 1775 1776
		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

1777 1778 1779
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1780
		zap.String("role", typeutil.ProxyRole),
1781 1782 1783
		zap.Int64("MsgID", gibpt.ID()),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
		zap.Uint64("EndTs", gibpt.EndTs()),
1784 1785 1786 1787
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1788 1789 1790 1791

	if err := gibpt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
1792
			zap.Error(err),
1793
			zap.String("traceID", traceID),
1794
			zap.String("role", typeutil.ProxyRole),
1795 1796 1797
			zap.Int64("MsgID", gibpt.ID()),
			zap.Uint64("BeginTs", gibpt.BeginTs()),
			zap.Uint64("EndTs", gibpt.EndTs()),
1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

		return &milvuspb.GetIndexBuildProgressResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
1810 1811 1812 1813

	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1814
		zap.String("role", typeutil.ProxyRole),
1815 1816 1817 1818 1819 1820 1821 1822
		zap.Int64("MsgID", gibpt.ID()),
		zap.Uint64("BeginTs", gibpt.BeginTs()),
		zap.Uint64("EndTs", gibpt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName),
		zap.Any("result", gibpt.result))
1823 1824

	return gibpt.result, nil
1825 1826
}

1827
// GetIndexState get the build-state of index.
C
Cai Yudong 已提交
1828
func (node *Proxy) GetIndexState(ctx context.Context, request *milvuspb.GetIndexStateRequest) (*milvuspb.GetIndexStateResponse, error) {
1829 1830 1831 1832 1833
	if !node.checkHealthy() {
		return &milvuspb.GetIndexStateResponse{
			Status: unhealthyStatus(),
		}, nil
	}
1834 1835 1836 1837 1838

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Insert")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

1839
	dipt := &getIndexStateTask{
G
godchen 已提交
1840 1841 1842
		ctx:                  ctx,
		Condition:            NewTaskCondition(ctx),
		GetIndexStateRequest: request,
1843 1844
		indexCoord:           node.indexCoord,
		rootCoord:            node.rootCoord,
1845 1846
	}

1847 1848 1849 1850 1851
	method := "GetIndexState"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
1852
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
1853 1854 1855 1856
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1857 1858 1859 1860 1861 1862

	if err := node.sched.ddQueue.Enqueue(dipt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
1863
			zap.String("role", typeutil.ProxyRole),
1864 1865 1866 1867 1868
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

G
godchen 已提交
1869
		return &milvuspb.GetIndexStateResponse{
1870
			Status: &commonpb.Status{
1871
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1872 1873 1874 1875 1876
				Reason:    err.Error(),
			},
		}, nil
	}

1877 1878 1879
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
1880
		zap.String("role", typeutil.ProxyRole),
1881 1882 1883
		zap.Int64("MsgID", dipt.ID()),
		zap.Uint64("BeginTs", dipt.BeginTs()),
		zap.Uint64("EndTs", dipt.EndTs()),
D
dragondriver 已提交
1884 1885 1886 1887
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))
1888 1889 1890 1891

	if err := dipt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
1892
			zap.Error(err),
1893
			zap.String("traceID", traceID),
1894
			zap.String("role", typeutil.ProxyRole),
1895 1896 1897
			zap.Int64("MsgID", dipt.ID()),
			zap.Uint64("BeginTs", dipt.BeginTs()),
			zap.Uint64("EndTs", dipt.EndTs()),
D
dragondriver 已提交
1898 1899 1900 1901 1902
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.String("field", request.FieldName),
			zap.String("index name", request.IndexName))

G
godchen 已提交
1903
		return &milvuspb.GetIndexStateResponse{
1904
			Status: &commonpb.Status{
1905
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
1906 1907 1908 1909 1910
				Reason:    err.Error(),
			},
		}, nil
	}

1911 1912 1913
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
1914
		zap.String("role", typeutil.ProxyRole),
1915 1916 1917 1918 1919 1920 1921 1922
		zap.Int64("MsgID", dipt.ID()),
		zap.Uint64("BeginTs", dipt.BeginTs()),
		zap.Uint64("EndTs", dipt.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("field", request.FieldName),
		zap.String("index name", request.IndexName))

1923 1924 1925
	return dipt.result, nil
}

1926
// Insert insert records into collection.
C
Cai Yudong 已提交
1927
func (node *Proxy) Insert(ctx context.Context, request *milvuspb.InsertRequest) (*milvuspb.MutationResult, error) {
X
Xiangyu Wang 已提交
1928 1929 1930 1931 1932 1933
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Insert")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
	log.Info("Start processing insert request in Proxy", zap.String("traceID", traceID))
	defer log.Info("Finish processing insert request in Proxy", zap.String("traceID", traceID))

1934 1935 1936 1937 1938
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}
D
dragondriver 已提交
1939

1940
	it := &insertTask{
1941 1942 1943
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		req:       request,
1944 1945 1946 1947
		BaseInsertTask: BaseInsertTask{
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
1948
			InsertRequest: internalpb.InsertRequest{
1949
				Base: &commonpb.MsgBase{
1950
					MsgType: commonpb.MsgType_Insert,
1951 1952 1953 1954
					MsgID:   0,
				},
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
1955
				// RowData: transfer column based request to this
1956 1957
			},
		},
1958
		rowIDAllocator: node.idAllocator,
1959
		segIDAssigner:  node.segAssigner,
1960
		chMgr:          node.chMgr,
1961
		chTicker:       node.chTicker,
1962
	}
1963 1964

	if len(it.PartitionName) <= 0 {
1965
		it.PartitionName = Params.ProxyCfg.DefaultPartitionName
1966 1967
	}

X
Xiangyu Wang 已提交
1968
	constructFailedResponse := func(err error) *milvuspb.MutationResult {
1969 1970 1971 1972 1973
		numRows := it.req.NumRows
		errIndex := make([]uint32, numRows)
		for i := uint32(0); i < numRows; i++ {
			errIndex[i] = i
		}
X
Xiangyu Wang 已提交
1974 1975 1976 1977 1978 1979 1980
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			ErrIndex: errIndex,
		}
1981 1982
	}

X
Xiangyu Wang 已提交
1983
	log.Debug("Enqueue insert request in Proxy",
1984
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
1985 1986 1987 1988 1989
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.Int("len(FieldsData)", len(request.FieldsData)),
		zap.Int("len(HashKeys)", len(request.HashKeys)),
1990 1991
		zap.Uint32("NumRows", request.NumRows),
		zap.String("traceID", traceID))
D
dragondriver 已提交
1992

X
Xiangyu Wang 已提交
1993 1994 1995
	if err := node.sched.dmQueue.Enqueue(it); err != nil {
		log.Debug("Failed to enqueue insert task: " + err.Error())
		return constructFailedResponse(err), nil
1996
	}
D
dragondriver 已提交
1997

X
Xiangyu Wang 已提交
1998
	log.Debug("Detail of insert request in Proxy",
1999
		zap.String("role", typeutil.ProxyRole),
X
Xiangyu Wang 已提交
2000
		zap.Int64("msgID", it.Base.MsgID),
D
dragondriver 已提交
2001 2002 2003 2004 2005
		zap.Uint64("BeginTS", it.BeginTs()),
		zap.Uint64("EndTS", it.EndTs()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
X
Xiangyu Wang 已提交
2006 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028
		zap.Uint32("NumRows", request.NumRows),
		zap.String("traceID", traceID))

	if err := it.WaitToFinish(); err != nil {
		log.Debug("Failed to execute insert task in task scheduler: "+err.Error(), zap.String("traceID", traceID))
		return constructFailedResponse(err), nil
	}

	if it.result.Status.ErrorCode != commonpb.ErrorCode_Success {
		setErrorIndex := func() {
			numRows := it.req.NumRows
			errIndex := make([]uint32, numRows)
			for i := uint32(0); i < numRows; i++ {
				errIndex[i] = i
			}
			it.result.ErrIndex = errIndex
		}

		setErrorIndex()
	}

	// InsertCnt always equals to the number of entities in the request
	it.result.InsertCnt = int64(it.req.NumRows)
D
dragondriver 已提交
2029

2030 2031 2032
	return it.result, nil
}

2033
// Delete delete records from collection, then these records cannot be searched.
G
groot 已提交
2034
func (node *Proxy) Delete(ctx context.Context, request *milvuspb.DeleteRequest) (*milvuspb.MutationResult, error) {
2035 2036 2037
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Delete")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)
2038 2039
	log.Info("Start processing delete request in Proxy", zap.String("traceID", traceID))
	defer log.Info("Finish processing delete request in Proxy", zap.String("traceID", traceID))
2040

G
groot 已提交
2041 2042 2043 2044 2045 2046
	if !node.checkHealthy() {
		return &milvuspb.MutationResult{
			Status: unhealthyStatus(),
		}, nil
	}

C
Cai Yudong 已提交
2047 2048 2049 2050 2051 2052 2053
	deleteReq := &milvuspb.DeleteRequest{
		DbName:         request.DbName,
		CollectionName: request.CollectionName,
		PartitionName:  request.PartitionName,
		Expr:           request.Expr,
	}

2054
	dt := &deleteTask{
C
Cai Yudong 已提交
2055 2056 2057
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		req:       deleteReq,
G
godchen 已提交
2058
		BaseDeleteTask: BaseDeleteTask{
G
godchen 已提交
2059 2060 2061
			BaseMsg: msgstream.BaseMsg{
				HashValues: request.HashKeys,
			},
G
godchen 已提交
2062 2063 2064 2065 2066 2067 2068 2069
			DeleteRequest: internalpb.DeleteRequest{
				Base: &commonpb.MsgBase{
					MsgType: commonpb.MsgType_Delete,
					MsgID:   0,
				},
				CollectionName: request.CollectionName,
				PartitionName:  request.PartitionName,
				// RowData: transfer column based request to this
C
Cai Yudong 已提交
2070 2071 2072 2073
			},
		},
		chMgr:    node.chMgr,
		chTicker: node.chTicker,
G
groot 已提交
2074 2075
	}

2076
	log.Debug("Enqueue delete request in Proxy",
2077
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2078 2079 2080 2081
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
		zap.String("expr", request.Expr))
2082 2083 2084 2085

	// MsgID will be set by Enqueue()
	if err := node.sched.dmQueue.Enqueue(dt); err != nil {
		log.Error("Failed to enqueue delete task: "+err.Error(), zap.String("traceID", traceID))
G
groot 已提交
2086 2087 2088 2089 2090 2091 2092 2093
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

2094
	log.Debug("Detail of delete request in Proxy",
2095
		zap.String("role", typeutil.ProxyRole),
G
groot 已提交
2096 2097 2098 2099 2100
		zap.Int64("msgID", dt.Base.MsgID),
		zap.Uint64("timestamp", dt.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.String("partition", request.PartitionName),
2101 2102
		zap.String("expr", request.Expr),
		zap.String("traceID", traceID))
G
groot 已提交
2103

2104 2105
	if err := dt.WaitToFinish(); err != nil {
		log.Error("Failed to execute delete task in task scheduler: "+err.Error(), zap.String("traceID", traceID))
G
groot 已提交
2106 2107 2108 2109 2110 2111 2112 2113 2114 2115 2116
		return &milvuspb.MutationResult{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}

	return dt.result, nil
}

2117
// Search search the most similar records of requests.
C
Cai Yudong 已提交
2118
func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
2119 2120 2121 2122 2123
	if !node.checkHealthy() {
		return &milvuspb.SearchResults{
			Status: unhealthyStatus(),
		}, nil
	}
D
dragondriver 已提交
2124

C
cai.zhang 已提交
2125 2126
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Search")
	defer sp.Finish()
D
dragondriver 已提交
2127 2128
	traceID, _, _ := trace.InfoFromSpan(sp)

2129
	qt := &searchTask{
S
sunby 已提交
2130
		ctx:       ctx,
2131
		Condition: NewTaskCondition(ctx),
G
godchen 已提交
2132
		SearchRequest: &internalpb.SearchRequest{
2133
			Base: &commonpb.MsgBase{
2134
				MsgType:  commonpb.MsgType_Search,
2135
				SourceID: Params.ProxyCfg.ProxyID,
2136
			},
2137
			ResultChannelID: strconv.FormatInt(Params.ProxyCfg.ProxyID, 10),
2138
		},
2139 2140 2141
		resultBuf: make(chan []*internalpb.SearchResults),
		query:     request,
		chMgr:     node.chMgr,
2142
		qc:        node.queryCoord,
2143 2144
	}

2145 2146 2147 2148 2149 2150
	method := "Search"
	travelTs := request.TravelTimestamp
	guaranteeTs := request.GuaranteeTimestamp

	log.Debug(
		rpcReceived(method),
D
dragondriver 已提交
2151
		zap.String("traceID", traceID),
2152
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2153 2154 2155 2156 2157
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("dsl", request.Dsl),
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2158 2159 2160 2161
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2162

2163 2164 2165
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
D
dragondriver 已提交
2166 2167
			zap.Error(err),
			zap.String("traceID", traceID),
2168
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2169 2170 2171 2172 2173 2174
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames),
			zap.Any("dsl", request.Dsl),
			zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
			zap.Any("OutputFields", request.OutputFields),
2175 2176 2177
			zap.Any("search_params", request.SearchParams),
			zap.Uint64("travel_timestamp", travelTs),
			zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2178

2179 2180
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2181
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2182 2183 2184 2185 2186
				Reason:    err.Error(),
			},
		}, nil
	}

2187 2188
	log.Debug(
		rpcEnqueued(method),
D
dragondriver 已提交
2189
		zap.String("traceID", traceID),
2190
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2191
		zap.Int64("msgID", qt.ID()),
D
dragondriver 已提交
2192 2193 2194 2195 2196
		zap.Uint64("timestamp", qt.Base.Timestamp),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("dsl", request.Dsl),
2197
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2198 2199 2200 2201
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2202

2203 2204 2205
	if err := qt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2206
			zap.Error(err),
D
dragondriver 已提交
2207
			zap.String("traceID", traceID),
2208
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2209
			zap.Int64("msgID", qt.ID()),
D
dragondriver 已提交
2210 2211 2212 2213
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames),
			zap.Any("dsl", request.Dsl),
2214
			zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2215 2216 2217 2218
			zap.Any("OutputFields", request.OutputFields),
			zap.Any("search_params", request.SearchParams),
			zap.Uint64("travel_timestamp", travelTs),
			zap.Uint64("guarantee_timestamp", guaranteeTs))
2219

2220 2221
		return &milvuspb.SearchResults{
			Status: &commonpb.Status{
2222
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2223 2224 2225 2226 2227
				Reason:    err.Error(),
			},
		}, nil
	}

2228 2229
	log.Debug(
		rpcDone(method),
D
dragondriver 已提交
2230
		zap.String("traceID", traceID),
2231
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2232 2233 2234 2235 2236 2237
		zap.Int64("msgID", qt.ID()),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("dsl", request.Dsl),
		zap.Any("len(PlaceholderGroup)", len(request.PlaceholderGroup)),
2238 2239 2240 2241
		zap.Any("OutputFields", request.OutputFields),
		zap.Any("search_params", request.SearchParams),
		zap.Uint64("travel_timestamp", travelTs),
		zap.Uint64("guarantee_timestamp", guaranteeTs))
D
dragondriver 已提交
2242

2243 2244 2245
	return qt.result, nil
}

2246
// Flush notify data nodes to persist the data of collection.
2247 2248 2249 2250 2251 2252 2253
func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
	resp := &milvuspb.FlushResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    "",
		},
	}
2254
	if !node.checkHealthy() {
2255 2256
		resp.Status.Reason = "proxy is not healthy"
		return resp, nil
2257
	}
D
dragondriver 已提交
2258 2259 2260 2261 2262

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Flush")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2263
	ft := &flushTask{
T
ThreadDao 已提交
2264 2265 2266
		ctx:          ctx,
		Condition:    NewTaskCondition(ctx),
		FlushRequest: request,
2267
		dataCoord:    node.dataCoord,
2268 2269
	}

D
dragondriver 已提交
2270 2271 2272 2273 2274
	method := "Flush"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2275
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2276 2277
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2278 2279 2280 2281 2282 2283

	if err := node.sched.ddQueue.Enqueue(ft); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2284
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2285 2286 2287
			zap.String("db", request.DbName),
			zap.Any("collections", request.CollectionNames))

2288 2289
		resp.Status.Reason = err.Error()
		return resp, nil
2290 2291
	}

D
dragondriver 已提交
2292 2293 2294
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2295
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2296 2297 2298
		zap.Int64("MsgID", ft.ID()),
		zap.Uint64("BeginTs", ft.BeginTs()),
		zap.Uint64("EndTs", ft.EndTs()),
D
dragondriver 已提交
2299 2300
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))
D
dragondriver 已提交
2301 2302 2303 2304

	if err := ft.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
D
dragondriver 已提交
2305
			zap.Error(err),
D
dragondriver 已提交
2306
			zap.String("traceID", traceID),
2307
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2308 2309 2310
			zap.Int64("MsgID", ft.ID()),
			zap.Uint64("BeginTs", ft.BeginTs()),
			zap.Uint64("EndTs", ft.EndTs()),
D
dragondriver 已提交
2311 2312 2313
			zap.String("db", request.DbName),
			zap.Any("collections", request.CollectionNames))

D
dragondriver 已提交
2314
		resp.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
2315 2316
		resp.Status.Reason = err.Error()
		return resp, nil
2317 2318
	}

D
dragondriver 已提交
2319 2320 2321
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
2322
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2323 2324 2325 2326 2327 2328
		zap.Int64("MsgID", ft.ID()),
		zap.Uint64("BeginTs", ft.BeginTs()),
		zap.Uint64("EndTs", ft.EndTs()),
		zap.String("db", request.DbName),
		zap.Any("collections", request.CollectionNames))

2329
	return ft.result, nil
2330 2331
}

2332
// Query get the records by primary keys.
C
Cai Yudong 已提交
2333
func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
2334 2335 2336 2337 2338
	if !node.checkHealthy() {
		return &milvuspb.QueryResults{
			Status: unhealthyStatus(),
		}, nil
	}
2339

D
dragondriver 已提交
2340 2341 2342 2343
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-Query")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2344
	qt := &queryTask{
2345 2346 2347 2348 2349
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
			Base: &commonpb.MsgBase{
				MsgType:  commonpb.MsgType_Retrieve,
2350
				SourceID: Params.ProxyCfg.ProxyID,
2351
			},
2352
			ResultChannelID: strconv.FormatInt(Params.ProxyCfg.ProxyID, 10),
2353 2354
		},
		resultBuf: make(chan []*internalpb.RetrieveResults),
2355
		query:     request,
2356 2357
		chMgr:     node.chMgr,
		qc:        node.queryCoord,
2358 2359
	}

D
dragondriver 已提交
2360 2361 2362 2363 2364
	method := "Query"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2365
		zap.String("role", typeutil.ProxyRole),
2366 2367 2368
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
G
godchen 已提交
2369

D
dragondriver 已提交
2370 2371 2372 2373 2374 2375
	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
2376 2377 2378
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
D
dragondriver 已提交
2379

2380 2381 2382 2383 2384 2385
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2386 2387
	}

D
dragondriver 已提交
2388 2389 2390
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2391
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2392 2393 2394
		zap.Int64("MsgID", qt.ID()),
		zap.Uint64("BeginTs", qt.BeginTs()),
		zap.Uint64("EndTs", qt.EndTs()),
2395 2396 2397
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
D
dragondriver 已提交
2398 2399 2400 2401 2402 2403

	if err := qt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2404
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2405 2406 2407
			zap.Int64("MsgID", qt.ID()),
			zap.Uint64("BeginTs", qt.BeginTs()),
			zap.Uint64("EndTs", qt.EndTs()),
2408 2409 2410
			zap.String("db", request.DbName),
			zap.String("collection", request.CollectionName),
			zap.Any("partitions", request.PartitionNames))
2411

2412 2413 2414 2415 2416 2417 2418
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
2419

D
dragondriver 已提交
2420 2421 2422 2423 2424 2425 2426
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", qt.ID()),
		zap.Uint64("BeginTs", qt.BeginTs()),
		zap.Uint64("EndTs", qt.EndTs()),
2427 2428 2429
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames))
D
dragondriver 已提交
2430

2431 2432 2433 2434 2435
	return &milvuspb.QueryResults{
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
	}, nil
}
2436

2437
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
2438 2439 2440 2441
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2442 2443 2444 2445 2446

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CreateAlias")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

Y
Yusup 已提交
2447 2448 2449 2450 2451 2452 2453
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2454 2455 2456 2457 2458 2459 2460 2461 2462 2463 2464 2465 2466 2467 2468 2469 2470 2471 2472 2473
	method := "CreateAlias"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

	if err := node.sched.ddQueue.Enqueue(cat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

Y
Yusup 已提交
2474 2475 2476 2477 2478 2479
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2480 2481 2482
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2483
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2484 2485 2486 2487
		zap.Int64("MsgID", cat.ID()),
		zap.Uint64("BeginTs", cat.BeginTs()),
		zap.Uint64("EndTs", cat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2488 2489
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
2490 2491 2492 2493

	if err := cat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2494
			zap.Error(err),
D
dragondriver 已提交
2495
			zap.String("traceID", traceID),
2496
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2497 2498 2499 2500
			zap.Int64("MsgID", cat.ID()),
			zap.Uint64("BeginTs", cat.BeginTs()),
			zap.Uint64("EndTs", cat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2501 2502 2503 2504 2505 2506 2507 2508 2509
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2510 2511 2512 2513 2514 2515 2516 2517 2518 2519 2520
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", cat.ID()),
		zap.Uint64("BeginTs", cat.BeginTs()),
		zap.Uint64("EndTs", cat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

Y
Yusup 已提交
2521 2522 2523
	return cat.result, nil
}

2524
// DropAlias alter the alias of collection.
Y
Yusup 已提交
2525 2526 2527 2528
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2529 2530 2531 2532 2533

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-DropAlias")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

Y
Yusup 已提交
2534 2535 2536 2537 2538 2539 2540
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2541 2542 2543 2544 2545 2546 2547 2548 2549 2550 2551 2552 2553 2554 2555 2556 2557 2558
	method := "DropAlias"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

	if err := node.sched.ddQueue.Enqueue(dat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias))

Y
Yusup 已提交
2559 2560 2561 2562 2563 2564
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2565 2566 2567
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2568
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2569 2570 2571 2572
		zap.Int64("MsgID", dat.ID()),
		zap.Uint64("BeginTs", dat.BeginTs()),
		zap.Uint64("EndTs", dat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2573
		zap.String("alias", request.Alias))
D
dragondriver 已提交
2574 2575 2576 2577

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2578
			zap.Error(err),
D
dragondriver 已提交
2579
			zap.String("traceID", traceID),
2580
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2581 2582 2583 2584
			zap.Int64("MsgID", dat.ID()),
			zap.Uint64("BeginTs", dat.BeginTs()),
			zap.Uint64("EndTs", dat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2585 2586 2587 2588 2589 2590 2591 2592
			zap.String("alias", request.Alias))

		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2593 2594 2595 2596 2597 2598 2599 2600 2601 2602
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", dat.ID()),
		zap.Uint64("BeginTs", dat.BeginTs()),
		zap.Uint64("EndTs", dat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias))

Y
Yusup 已提交
2603 2604 2605
	return dat.result, nil
}

2606
// AlterAlias alter alias of collection.
Y
Yusup 已提交
2607 2608 2609 2610
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2611 2612 2613 2614 2615

	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-AlterAlias")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

Y
Yusup 已提交
2616 2617 2618 2619 2620 2621 2622
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2623 2624 2625 2626 2627 2628 2629 2630 2631 2632 2633 2634 2635 2636 2637 2638 2639 2640 2641 2642
	method := "AlterAlias"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

	if err := node.sched.ddQueue.Enqueue(aat); err != nil {
		log.Warn(
			rpcFailedToEnqueue(method),
			zap.Error(err),
			zap.String("traceID", traceID),
			zap.String("role", typeutil.ProxyRole),
			zap.String("db", request.DbName),
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

Y
Yusup 已提交
2643 2644 2645 2646 2647 2648
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2649 2650 2651
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2652
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2653 2654 2655 2656
		zap.Int64("MsgID", aat.ID()),
		zap.Uint64("BeginTs", aat.BeginTs()),
		zap.Uint64("EndTs", aat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2657 2658
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
2659 2660 2661 2662

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2663
			zap.Error(err),
D
dragondriver 已提交
2664
			zap.String("traceID", traceID),
2665
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2666 2667 2668 2669
			zap.Int64("MsgID", aat.ID()),
			zap.Uint64("BeginTs", aat.BeginTs()),
			zap.Uint64("EndTs", aat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2670 2671 2672 2673 2674 2675 2676 2677 2678
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2679 2680 2681 2682 2683 2684 2685 2686 2687 2688 2689
	log.Debug(
		rpcDone(method),
		zap.String("traceID", traceID),
		zap.String("role", typeutil.ProxyRole),
		zap.Int64("MsgID", aat.ID()),
		zap.Uint64("BeginTs", aat.BeginTs()),
		zap.Uint64("EndTs", aat.EndTs()),
		zap.String("db", request.DbName),
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))

Y
Yusup 已提交
2690 2691 2692
	return aat.result, nil
}

2693
// CalcDistance calculates the distances between vectors.
2694
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
2695 2696 2697 2698 2699 2700
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}

2701
	param, _ := funcutil.GetAttrByKeyFromRepeatedKV("metric", request.GetParams())
2702 2703 2704 2705 2706 2707 2708 2709
	metric, err := distance.ValidateMetricType(param)
	if err != nil {
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2710 2711
	}

2712 2713 2714 2715
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CalcDistance")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2716 2717
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2718

2719 2720 2721 2722 2723
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2724 2725
		}

2726
		qt := &queryTask{
2727 2728 2729 2730 2731
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
				Base: &commonpb.MsgBase{
					MsgType:  commonpb.MsgType_Retrieve,
2732
					SourceID: Params.ProxyCfg.ProxyID,
2733
				},
2734
				ResultChannelID: strconv.FormatInt(Params.ProxyCfg.ProxyID, 10),
2735
			},
2736
			resultBuf: make(chan []*internalpb.RetrieveResults),
2737
			query:     queryRequest,
2738
			chMgr:     node.chMgr,
2739
			qc:        node.queryCoord,
Y
yukun 已提交
2740
			ids:       ids.IdArray,
2741 2742
		}

2743
		err := node.sched.dqQueue.Enqueue(qt)
2744
		if err != nil {
2745 2746 2747
			log.Debug("CalcDistance queryTask failed to enqueue",
				zap.Error(err),
				zap.String("traceID", traceID),
2748
				zap.String("role", typeutil.ProxyRole),
2749 2750 2751 2752
				zap.String("db", queryRequest.DbName),
				zap.String("collection", queryRequest.CollectionName),
				zap.Any("partitions", queryRequest.PartitionNames))

2753 2754 2755 2756 2757
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2758
			}, err
2759
		}
2760 2761 2762

		log.Debug("CalcDistance queryTask enqueued",
			zap.String("traceID", traceID),
2763
			zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2764 2765 2766 2767 2768 2769
			zap.Int64("msgID", qt.Base.MsgID),
			zap.Uint64("timestamp", qt.Base.Timestamp),
			zap.String("db", queryRequest.DbName),
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
			zap.Any("OutputFields", queryRequest.OutputFields))
2770 2771 2772 2773

		err = qt.WaitToFinish()
		if err != nil {
			log.Debug("CalcDistance queryTask failed to WaitToFinish",
G
godchen 已提交
2774
				zap.Error(err),
2775
				zap.String("traceID", traceID),
2776
				zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2777 2778 2779 2780 2781 2782
				zap.Int64("msgID", qt.Base.MsgID),
				zap.Uint64("timestamp", qt.Base.Timestamp),
				zap.String("db", queryRequest.DbName),
				zap.String("collection", queryRequest.CollectionName),
				zap.Any("partitions", queryRequest.PartitionNames),
				zap.Any("OutputFields", queryRequest.OutputFields))
2783 2784 2785 2786 2787 2788

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2789
			}, err
2790
		}
2791 2792 2793

		log.Debug("CalcDistance queryTask Done",
			zap.String("traceID", traceID),
2794
			zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2795 2796 2797 2798 2799 2800
			zap.Int64("msgID", qt.Base.MsgID),
			zap.Uint64("timestamp", qt.Base.Timestamp),
			zap.String("db", queryRequest.DbName),
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames),
			zap.Any("OutputFields", queryRequest.OutputFields))
2801 2802

		return &milvuspb.QueryResults{
2803 2804
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2805 2806 2807
		}, nil
	}

2808 2809 2810 2811 2812 2813 2814 2815 2816 2817 2818 2819 2820 2821
	// the vectors retrieved are random order, we need re-arrange the vectors by the order of input ids
	arrangeFunc := func(ids *milvuspb.VectorIDs, retrievedFields []*schemapb.FieldData) (*schemapb.VectorField, error) {
		var retrievedIds *schemapb.ScalarField
		var retrievedVectors *schemapb.VectorField
		for _, fieldData := range retrievedFields {
			if fieldData.FieldName == ids.FieldName {
				retrievedVectors = fieldData.GetVectors()
			}
			if fieldData.Type == schemapb.DataType_Int64 {
				retrievedIds = fieldData.GetScalars()
			}
		}

		if retrievedIds == nil || retrievedVectors == nil {
2822
			return nil, errors.New("failed to fetch vectors")
2823 2824 2825 2826 2827 2828 2829 2830 2831 2832 2833 2834 2835 2836 2837 2838
		}

		dict := make(map[int64]int)
		for index, id := range retrievedIds.GetLongData().Data {
			dict[id] = index
		}

		inputIds := ids.IdArray.GetIntId().Data
		if retrievedVectors.GetFloatVector() != nil {
			floatArr := retrievedVectors.GetFloatVector().Data
			element := retrievedVectors.GetDim()
			result := make([]float32, 0, int64(len(inputIds))*element)
			for _, id := range inputIds {
				index, ok := dict[id]
				if !ok {
					log.Error("id not found in CalcDistance", zap.Int64("id", id))
2839
					return nil, errors.New("failed to fetch vectors by id: " + fmt.Sprintln(id))
2840 2841 2842 2843 2844 2845 2846 2847 2848 2849 2850 2851 2852 2853 2854 2855 2856 2857 2858 2859 2860 2861 2862 2863 2864 2865
				}
				result = append(result, floatArr[int64(index)*element:int64(index+1)*element]...)
			}

			return &schemapb.VectorField{
				Dim: element,
				Data: &schemapb.VectorField_FloatVector{
					FloatVector: &schemapb.FloatArray{
						Data: result,
					},
				},
			}, nil
		}

		if retrievedVectors.GetBinaryVector() != nil {
			binaryArr := retrievedVectors.GetBinaryVector()
			element := retrievedVectors.GetDim()
			if element%8 != 0 {
				element = element + 8 - element%8
			}

			result := make([]byte, 0, int64(len(inputIds))*element)
			for _, id := range inputIds {
				index, ok := dict[id]
				if !ok {
					log.Error("id not found in CalcDistance", zap.Int64("id", id))
2866
					return nil, errors.New("failed to fetch vectors by id: " + fmt.Sprintln(id))
2867 2868 2869 2870 2871 2872 2873 2874 2875 2876 2877 2878
				}
				result = append(result, binaryArr[int64(index)*element:int64(index+1)*element]...)
			}

			return &schemapb.VectorField{
				Dim: element * 8,
				Data: &schemapb.VectorField_BinaryVector{
					BinaryVector: result,
				},
			}, nil
		}

2879
		return nil, errors.New("failed to fetch vectors")
2880 2881
	}

2882 2883
	log.Debug("CalcDistance received",
		zap.String("traceID", traceID),
2884
		zap.String("role", typeutil.ProxyRole),
2885
		zap.String("metric", metric))
G
godchen 已提交
2886

2887 2888 2889
	vectorsLeft := request.GetOpLeft().GetDataArray()
	opLeft := request.GetOpLeft().GetIdArray()
	if opLeft != nil {
2890 2891
		log.Debug("OpLeft IdArray not empty, Get vectors by id",
			zap.String("traceID", traceID),
2892
			zap.String("role", typeutil.ProxyRole))
2893

2894
		result, err := query(opLeft)
2895
		if err != nil {
2896 2897 2898
			log.Debug("Failed to get left vectors by id",
				zap.Error(err),
				zap.String("traceID", traceID),
2899
				zap.String("role", typeutil.ProxyRole))
2900

2901 2902 2903 2904 2905 2906 2907 2908
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

2909 2910
		log.Debug("OpLeft IdArray not empty, Get vectors by id done",
			zap.String("traceID", traceID),
2911
			zap.String("role", typeutil.ProxyRole))
2912

2913 2914
		vectorsLeft, err = arrangeFunc(opLeft, result.FieldsData)
		if err != nil {
2915 2916 2917
			log.Debug("Failed to re-arrange left vectors",
				zap.Error(err),
				zap.String("traceID", traceID),
2918
				zap.String("role", typeutil.ProxyRole))
2919

2920 2921 2922 2923 2924 2925
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
2926
		}
2927 2928 2929

		log.Debug("Re-arrange left vectors done",
			zap.String("traceID", traceID),
2930
			zap.String("role", typeutil.ProxyRole))
2931 2932
	}

G
groot 已提交
2933
	if vectorsLeft == nil {
2934 2935 2936
		msg := "Left vectors array is empty"
		log.Debug(msg,
			zap.String("traceID", traceID),
2937
			zap.String("role", typeutil.ProxyRole))
2938

G
groot 已提交
2939 2940 2941
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2942
				Reason:    msg,
G
groot 已提交
2943 2944 2945 2946
			},
		}, nil
	}

2947 2948 2949
	vectorsRight := request.GetOpRight().GetDataArray()
	opRight := request.GetOpRight().GetIdArray()
	if opRight != nil {
2950 2951
		log.Debug("OpRight IdArray not empty, Get vectors by id",
			zap.String("traceID", traceID),
2952
			zap.String("role", typeutil.ProxyRole))
2953

2954
		result, err := query(opRight)
2955
		if err != nil {
2956 2957 2958
			log.Debug("Failed to get right vectors by id",
				zap.Error(err),
				zap.String("traceID", traceID),
2959
				zap.String("role", typeutil.ProxyRole))
2960

2961 2962 2963 2964 2965 2966 2967 2968
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

2969 2970
		log.Debug("OpRight IdArray not empty, Get vectors by id done",
			zap.String("traceID", traceID),
2971
			zap.String("role", typeutil.ProxyRole))
2972

2973 2974
		vectorsRight, err = arrangeFunc(opRight, result.FieldsData)
		if err != nil {
2975 2976 2977
			log.Debug("Failed to re-arrange right vectors",
				zap.Error(err),
				zap.String("traceID", traceID),
2978
				zap.String("role", typeutil.ProxyRole))
2979

2980 2981 2982 2983 2984 2985
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
2986
		}
2987 2988 2989

		log.Debug("Re-arrange right vectors done",
			zap.String("traceID", traceID),
2990
			zap.String("role", typeutil.ProxyRole))
2991 2992
	}

G
groot 已提交
2993
	if vectorsRight == nil {
2994 2995 2996
		msg := "Right vectors array is empty"
		log.Debug(msg,
			zap.String("traceID", traceID),
2997
			zap.String("role", typeutil.ProxyRole))
2998

G
groot 已提交
2999 3000 3001
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3002
				Reason:    msg,
G
groot 已提交
3003 3004 3005 3006
			},
		}, nil
	}

3007
	if vectorsLeft.Dim != vectorsRight.Dim {
3008 3009 3010
		msg := "Vectors dimension is not equal"
		log.Debug(msg,
			zap.String("traceID", traceID),
3011
			zap.String("role", typeutil.ProxyRole))
3012

3013 3014 3015
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3016
				Reason:    msg,
3017 3018 3019 3020 3021
			},
		}, nil
	}

	if vectorsLeft.GetFloatVector() != nil && vectorsRight.GetFloatVector() != nil {
3022 3023
		distances, err := distance.CalcFloatDistance(vectorsLeft.Dim, vectorsLeft.GetFloatVector().Data, vectorsRight.GetFloatVector().Data, metric)
		if err != nil {
3024 3025 3026
			log.Debug("Failed to CalcFloatDistance",
				zap.Error(err),
				zap.String("traceID", traceID),
3027
				zap.String("role", typeutil.ProxyRole))
3028

3029 3030 3031 3032 3033 3034 3035 3036
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

3037 3038 3039
		log.Debug("CalcFloatDistance done",
			zap.Error(err),
			zap.String("traceID", traceID),
3040
			zap.String("role", typeutil.ProxyRole))
3041

3042 3043 3044 3045 3046 3047 3048 3049 3050 3051
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
			Array: &milvuspb.CalcDistanceResults_FloatDist{
				FloatDist: &schemapb.FloatArray{
					Data: distances,
				},
			},
		}, nil
	}

3052
	if vectorsLeft.GetBinaryVector() != nil && vectorsRight.GetBinaryVector() != nil {
G
groot 已提交
3053
		hamming, err := distance.CalcHammingDistance(vectorsLeft.Dim, vectorsLeft.GetBinaryVector(), vectorsRight.GetBinaryVector())
3054
		if err != nil {
3055 3056 3057
			log.Debug("Failed to CalcHammingDistance",
				zap.Error(err),
				zap.String("traceID", traceID),
3058
				zap.String("role", typeutil.ProxyRole))
3059

3060 3061 3062 3063 3064 3065 3066 3067 3068
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

		if metric == distance.HAMMING {
3069 3070
			log.Debug("CalcHammingDistance done",
				zap.String("traceID", traceID),
3071
				zap.String("role", typeutil.ProxyRole))
3072

3073 3074 3075 3076
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
				Array: &milvuspb.CalcDistanceResults_IntDist{
					IntDist: &schemapb.IntArray{
G
groot 已提交
3077
						Data: hamming,
3078 3079 3080 3081 3082 3083
					},
				},
			}, nil
		}

		if metric == distance.TANIMOTO {
G
groot 已提交
3084
			tanimoto, err := distance.CalcTanimotoCoefficient(vectorsLeft.Dim, hamming)
3085
			if err != nil {
3086 3087 3088
				log.Debug("Failed to CalcTanimotoCoefficient",
					zap.Error(err),
					zap.String("traceID", traceID),
3089
					zap.String("role", typeutil.ProxyRole))
3090

3091 3092 3093 3094 3095 3096 3097 3098
				return &milvuspb.CalcDistanceResults{
					Status: &commonpb.Status{
						ErrorCode: commonpb.ErrorCode_UnexpectedError,
						Reason:    err.Error(),
					},
				}, nil
			}

3099 3100
			log.Debug("CalcTanimotoCoefficient done",
				zap.String("traceID", traceID),
3101
				zap.String("role", typeutil.ProxyRole))
3102

3103 3104 3105 3106 3107 3108 3109 3110 3111 3112 3113
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
				Array: &milvuspb.CalcDistanceResults_FloatDist{
					FloatDist: &schemapb.FloatArray{
						Data: tanimoto,
					},
				},
			}, nil
		}
	}

3114
	err = errors.New("unexpected error")
3115
	if (vectorsLeft.GetBinaryVector() != nil && vectorsRight.GetFloatVector() != nil) || (vectorsLeft.GetFloatVector() != nil && vectorsRight.GetBinaryVector() != nil) {
3116
		err = errors.New("cannot calculate distance between binary vectors and float vectors")
3117 3118
	}

3119 3120 3121
	log.Debug("Failed to CalcDistance",
		zap.Error(err),
		zap.String("traceID", traceID),
3122
		zap.String("role", typeutil.ProxyRole))
3123

3124 3125 3126 3127 3128 3129 3130 3131
	return &milvuspb.CalcDistanceResults{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		},
	}, nil
}

3132
// GetDdChannel returns the used channel for dd operations.
C
Cai Yudong 已提交
3133
func (node *Proxy) GetDdChannel(ctx context.Context, request *internalpb.GetDdChannelRequest) (*milvuspb.StringResponse, error) {
3134 3135
	panic("implement me")
}
X
XuanYang-cn 已提交
3136

3137
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
3138
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
D
dragondriver 已提交
3139
	log.Debug("GetPersistentSegmentInfo",
3140
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3141 3142 3143
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3144
	resp := &milvuspb.GetPersistentSegmentInfoResponse{
X
XuanYang-cn 已提交
3145
		Status: &commonpb.Status{
3146
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
X
XuanYang-cn 已提交
3147 3148
		},
	}
3149 3150 3151 3152
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
G
godchen 已提交
3153
	segments, err := node.getSegmentsOfCollection(ctx, req.DbName, req.CollectionName)
X
XuanYang-cn 已提交
3154
	if err != nil {
3155
		resp.Status.Reason = fmt.Errorf("getSegmentsOfCollection, err:%w", err).Error()
X
XuanYang-cn 已提交
3156 3157
		return resp, nil
	}
3158
	infoResp, err := node.dataCoord.GetSegmentInfo(ctx, &datapb.GetSegmentInfoRequest{
X
XuanYang-cn 已提交
3159
		Base: &commonpb.MsgBase{
3160
			MsgType:   commonpb.MsgType_SegmentInfo,
X
XuanYang-cn 已提交
3161 3162
			MsgID:     0,
			Timestamp: 0,
3163
			SourceID:  Params.ProxyCfg.ProxyID,
X
XuanYang-cn 已提交
3164 3165 3166 3167
		},
		SegmentIDs: segments,
	})
	if err != nil {
3168
		log.Debug("GetPersistentSegmentInfo fail", zap.Error(err))
3169
		resp.Status.Reason = fmt.Errorf("dataCoord:GetSegmentInfo, err:%w", err).Error()
X
XuanYang-cn 已提交
3170 3171
		return resp, nil
	}
3172
	log.Debug("GetPersistentSegmentInfo ", zap.Int("len(infos)", len(infoResp.Infos)), zap.Any("status", infoResp.Status))
3173
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
3174 3175 3176 3177 3178 3179
		resp.Status.Reason = infoResp.Status.Reason
		return resp, nil
	}
	persistentInfos := make([]*milvuspb.PersistentSegmentInfo, len(infoResp.Infos))
	for i, info := range infoResp.Infos {
		persistentInfos[i] = &milvuspb.PersistentSegmentInfo{
S
sunby 已提交
3180
			SegmentID:    info.ID,
X
XuanYang-cn 已提交
3181 3182
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
S
sunby 已提交
3183
			NumRows:      info.NumOfRows,
X
XuanYang-cn 已提交
3184 3185 3186
			State:        info.State,
		}
	}
3187
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
X
XuanYang-cn 已提交
3188 3189 3190 3191
	resp.Infos = persistentInfos
	return resp, nil
}

J
jingkl 已提交
3192
// GetQuerySegmentInfo gets segment information from QueryCoord.
C
Cai Yudong 已提交
3193
func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
D
dragondriver 已提交
3194
	log.Debug("GetQuerySegmentInfo",
3195
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3196 3197 3198
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3199
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
3200
		Status: &commonpb.Status{
3201
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
3202 3203
		},
	}
3204 3205 3206 3207
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
G
godchen 已提交
3208
	segments, err := node.getSegmentsOfCollection(ctx, req.DbName, req.CollectionName)
Z
zhenshan.cao 已提交
3209 3210 3211 3212
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3213 3214 3215 3216 3217
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3218
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
Z
zhenshan.cao 已提交
3219
		Base: &commonpb.MsgBase{
3220
			MsgType:   commonpb.MsgType_SegmentInfo,
Z
zhenshan.cao 已提交
3221 3222
			MsgID:     0,
			Timestamp: 0,
3223
			SourceID:  Params.ProxyCfg.ProxyID,
Z
zhenshan.cao 已提交
3224
		},
3225 3226
		CollectionID: collID,
		SegmentIDs:   segments,
Z
zhenshan.cao 已提交
3227 3228
	})
	if err != nil {
3229
		log.Error("Failed to get segment info from QueryCoord",
3230
			zap.Int64s("segmentIDs", segments), zap.Error(err))
Z
zhenshan.cao 已提交
3231 3232 3233
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3234
	log.Debug("GetQuerySegmentInfo ", zap.Any("infos", infoResp.Infos), zap.Any("status", infoResp.Status))
3235
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
3236
		log.Error("Failed to get segment info from QueryCoord", zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
3237 3238 3239 3240 3241 3242 3243 3244 3245 3246 3247 3248 3249
		resp.Status.Reason = infoResp.Status.Reason
		return resp, nil
	}
	queryInfos := make([]*milvuspb.QuerySegmentInfo, len(infoResp.Infos))
	for i, info := range infoResp.Infos {
		queryInfos[i] = &milvuspb.QuerySegmentInfo{
			SegmentID:    info.SegmentID,
			CollectionID: info.CollectionID,
			PartitionID:  info.PartitionID,
			NumRows:      info.NumRows,
			MemSize:      info.MemSize,
			IndexName:    info.IndexName,
			IndexID:      info.IndexID,
3250
			NodeID:       info.NodeID,
X
xige-16 已提交
3251
			State:        info.SegmentState,
Z
zhenshan.cao 已提交
3252 3253
		}
	}
3254
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
3255 3256 3257 3258
	resp.Infos = queryInfos
	return resp, nil
}

C
Cai Yudong 已提交
3259
func (node *Proxy) getSegmentsOfCollection(ctx context.Context, dbName string, collectionName string) ([]UniqueID, error) {
3260
	describeCollectionResponse, err := node.rootCoord.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
X
XuanYang-cn 已提交
3261
		Base: &commonpb.MsgBase{
3262
			MsgType:   commonpb.MsgType_DescribeCollection,
X
XuanYang-cn 已提交
3263 3264
			MsgID:     0,
			Timestamp: 0,
3265
			SourceID:  Params.ProxyCfg.ProxyID,
X
XuanYang-cn 已提交
3266 3267 3268 3269 3270 3271 3272
		},
		DbName:         dbName,
		CollectionName: collectionName,
	})
	if err != nil {
		return nil, err
	}
3273
	if describeCollectionResponse.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
3274 3275 3276
		return nil, errors.New(describeCollectionResponse.Status.Reason)
	}
	collectionID := describeCollectionResponse.CollectionID
3277
	showPartitionsResp, err := node.rootCoord.ShowPartitions(ctx, &milvuspb.ShowPartitionsRequest{
X
XuanYang-cn 已提交
3278
		Base: &commonpb.MsgBase{
3279
			MsgType:   commonpb.MsgType_ShowPartitions,
X
XuanYang-cn 已提交
3280 3281
			MsgID:     0,
			Timestamp: 0,
3282
			SourceID:  Params.ProxyCfg.ProxyID,
X
XuanYang-cn 已提交
3283 3284 3285 3286 3287 3288 3289 3290
		},
		DbName:         dbName,
		CollectionName: collectionName,
		CollectionID:   collectionID,
	})
	if err != nil {
		return nil, err
	}
3291
	if showPartitionsResp.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
3292 3293 3294 3295 3296
		return nil, errors.New(showPartitionsResp.Status.Reason)
	}

	ret := make([]UniqueID, 0)
	for _, partitionID := range showPartitionsResp.PartitionIDs {
3297
		showSegmentResponse, err := node.rootCoord.ShowSegments(ctx, &milvuspb.ShowSegmentsRequest{
X
XuanYang-cn 已提交
3298
			Base: &commonpb.MsgBase{
3299
				MsgType:   commonpb.MsgType_ShowSegments,
X
XuanYang-cn 已提交
3300 3301
				MsgID:     0,
				Timestamp: 0,
3302
				SourceID:  Params.ProxyCfg.ProxyID,
X
XuanYang-cn 已提交
3303 3304 3305 3306 3307 3308 3309
			},
			CollectionID: collectionID,
			PartitionID:  partitionID,
		})
		if err != nil {
			return nil, err
		}
3310
		if showSegmentResponse.Status.ErrorCode != commonpb.ErrorCode_Success {
X
XuanYang-cn 已提交
3311 3312 3313 3314 3315 3316
			return nil, errors.New(showSegmentResponse.Status.Reason)
		}
		ret = append(ret, showSegmentResponse.SegmentIDs...)
	}
	return ret, nil
}
3317

J
jingkl 已提交
3318
// Dummy handles dummy request
C
Cai Yudong 已提交
3319
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
3320 3321 3322 3323 3324 3325 3326 3327 3328 3329 3330
	failedResponse := &milvuspb.DummyResponse{
		Response: `{"status": "fail"}`,
	}

	// TODO(wxyu): change name RequestType to Request
	drt, err := parseDummyRequestType(req.RequestType)
	if err != nil {
		log.Debug("Failed to parse dummy request type")
		return failedResponse, nil
	}

3331 3332
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3333
		if err != nil {
3334
			log.Debug("Failed to parse dummy query request")
3335 3336 3337
			return failedResponse, nil
		}

3338
		request := &milvuspb.QueryRequest{
3339 3340 3341
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3342
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3343 3344
		}

3345
		_, err = node.Query(ctx, request)
3346
		if err != nil {
3347
			log.Debug("Failed to execute dummy query")
3348 3349
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3350 3351 3352 3353 3354 3355

		return &milvuspb.DummyResponse{
			Response: `{"status": "success"}`,
		}, nil
	}

3356 3357
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3358 3359
}

J
jingkl 已提交
3360
// RegisterLink registers a link
C
Cai Yudong 已提交
3361
func (node *Proxy) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
G
godchen 已提交
3362
	code := node.stateCode.Load().(internalpb.StateCode)
D
dragondriver 已提交
3363
	log.Debug("RegisterLink",
3364
		zap.String("role", typeutil.ProxyRole),
C
Cai Yudong 已提交
3365
		zap.Any("state code of proxy", code))
D
dragondriver 已提交
3366

G
godchen 已提交
3367
	if code != internalpb.StateCode_Healthy {
3368 3369 3370
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
3371
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3372
				Reason:    "proxy not healthy",
3373 3374 3375 3376 3377 3378
			},
		}, nil
	}
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
3379
			ErrorCode: commonpb.ErrorCode_Success,
3380
			Reason:    os.Getenv(metricsinfo.DeployModeEnvKey),
3381 3382 3383
		},
	}, nil
}
3384

3385 3386 3387
// TODO(dragondriver): cache the Metrics and set a retention to the cache
func (node *Proxy) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
	log.Debug("Proxy.GetMetrics",
3388
		zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3389 3390 3391 3392
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetMetrics failed",
3393
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3394
			zap.String("req", req.Request),
3395
			zap.Error(errProxyIsUnhealthy(Params.ProxyCfg.ProxyID)))
3396 3397 3398 3399

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3400
				Reason:    msgProxyIsUnhealthy(Params.ProxyCfg.ProxyID),
3401 3402 3403 3404 3405 3406 3407 3408
			},
			Response: "",
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
3409
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3410 3411 3412 3413 3414 3415 3416 3417 3418 3419 3420 3421 3422 3423 3424
			zap.String("req", req.Request),
			zap.Error(err))

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
			Response: "",
		}, nil
	}

	log.Debug("Proxy.GetMetrics",
		zap.String("metric_type", metricType))

D
dragondriver 已提交
3425 3426 3427 3428 3429 3430 3431 3432 3433 3434
	msgID := UniqueID(0)
	msgID, err = node.idAllocator.AllocOne()
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to allocate id",
			zap.Error(err))
	}
	req.Base = &commonpb.MsgBase{
		MsgType:   commonpb.MsgType_SystemInfo,
		MsgID:     msgID,
		Timestamp: 0,
3435
		SourceID:  Params.ProxyCfg.ProxyID,
D
dragondriver 已提交
3436 3437
	}

3438
	if metricType == metricsinfo.SystemInfoMetrics {
3439 3440 3441 3442 3443 3444 3445
		ret, err := node.metricsCacheManager.GetSystemInfoMetrics()
		if err == nil && ret != nil {
			return ret, nil
		}
		log.Debug("failed to get system info metrics from cache, recompute instead",
			zap.Error(err))

3446
		metrics, err := getSystemInfoMetrics(ctx, req, node)
3447 3448

		log.Debug("Proxy.GetMetrics",
3449
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3450 3451 3452 3453 3454
			zap.String("req", req.Request),
			zap.String("metric_type", metricType),
			zap.Any("metrics", metrics), // TODO(dragondriver): necessary? may be very large
			zap.Error(err))

3455 3456
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3457
		return metrics, nil
3458 3459 3460
	}

	log.Debug("Proxy.GetMetrics failed, request metric type is not implemented yet",
3461
		zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3462 3463 3464 3465 3466 3467 3468 3469 3470 3471 3472 3473
		zap.String("req", req.Request),
		zap.String("metric_type", metricType))

	return &milvuspb.GetMetricsResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    metricsinfo.MsgUnimplementedMetric,
		},
		Response: "",
	}, nil
}

B
bigsheeper 已提交
3474 3475 3476
// LoadBalance would do a load balancing operation between query nodes
func (node *Proxy) LoadBalance(ctx context.Context, req *milvuspb.LoadBalanceRequest) (*commonpb.Status, error) {
	log.Debug("Proxy.LoadBalance",
3477
		zap.Int64("proxy_id", Params.ProxyCfg.ProxyID),
B
bigsheeper 已提交
3478 3479 3480 3481 3482 3483 3484 3485 3486 3487 3488 3489 3490 3491
		zap.Any("req", req))

	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}

	status := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
	}
	infoResp, err := node.queryCoord.LoadBalance(ctx, &querypb.LoadBalanceRequest{
		Base: &commonpb.MsgBase{
			MsgType:   commonpb.MsgType_LoadBalanceSegments,
			MsgID:     0,
			Timestamp: 0,
3492
			SourceID:  Params.ProxyCfg.ProxyID,
B
bigsheeper 已提交
3493 3494 3495
		},
		SourceNodeIDs:    []int64{req.SrcNodeID},
		DstNodeIDs:       req.DstNodeIDs,
X
xige-16 已提交
3496
		BalanceReason:    querypb.TriggerCondition_GrpcRequest,
B
bigsheeper 已提交
3497 3498 3499 3500 3501 3502 3503 3504 3505 3506 3507 3508 3509 3510 3511 3512 3513 3514
		SealedSegmentIDs: req.SealedSegmentIDs,
	})
	if err != nil {
		log.Error("Failed to LoadBalance from Query Coordinator",
			zap.Any("req", req), zap.Error(err))
		status.Reason = err.Error()
		return status, nil
	}
	if infoResp.ErrorCode != commonpb.ErrorCode_Success {
		log.Error("Failed to LoadBalance from Query Coordinator", zap.String("errMsg", infoResp.Reason))
		status.Reason = infoResp.Reason
		return status, nil
	}
	log.Debug("LoadBalance Done", zap.Any("req", req), zap.Any("status", infoResp))
	status.ErrorCode = commonpb.ErrorCode_Success
	return status, nil
}

J
jingkl 已提交
3515
//GetCompactionState gets the compaction state of multiple segments
3516 3517 3518 3519 3520 3521 3522 3523 3524 3525 3526 3527 3528
func (node *Proxy) GetCompactionState(ctx context.Context, req *milvuspb.GetCompactionStateRequest) (*milvuspb.GetCompactionStateResponse, error) {
	log.Info("received GetCompactionState request", zap.Int64("compactionID", req.GetCompactionID()))
	resp := &milvuspb.GetCompactionStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionState(ctx, req)
	log.Info("received GetCompactionState response", zap.Int64("compactionID", req.GetCompactionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

3529
// ManualCompaction invokes compaction on specified collection
3530 3531 3532 3533 3534 3535 3536 3537 3538 3539 3540 3541 3542 3543 3544 3545 3546 3547 3548 3549 3550 3551 3552 3553 3554 3555
func (node *Proxy) ManualCompaction(ctx context.Context, req *milvuspb.ManualCompactionRequest) (*milvuspb.ManualCompactionResponse, error) {
	log.Info("received ManualCompaction request", zap.Int64("collectionID", req.GetCollectionID()))
	resp := &milvuspb.ManualCompactionResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.ManualCompaction(ctx, req)
	log.Info("received ManualCompaction response", zap.Int64("collectionID", req.GetCollectionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

func (node *Proxy) GetCompactionStateWithPlans(ctx context.Context, req *milvuspb.GetCompactionPlansRequest) (*milvuspb.GetCompactionPlansResponse, error) {
	log.Info("received GetCompactionStateWithPlans request", zap.Int64("compactionID", req.GetCompactionID()))
	resp := &milvuspb.GetCompactionPlansResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}

	resp, err := node.dataCoord.GetCompactionStateWithPlans(ctx, req)
	log.Info("received GetCompactionStateWithPlans response", zap.Int64("compactionID", req.GetCompactionID()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

B
Bingyi Sun 已提交
3556 3557 3558
// GetFlushState gets the flush state of multiple segments
func (node *Proxy) GetFlushState(ctx context.Context, req *milvuspb.GetFlushStateRequest) (*milvuspb.GetFlushStateResponse, error) {
	log.Info("received get flush state request", zap.Any("request", req))
3559
	var err error
B
Bingyi Sun 已提交
3560 3561 3562 3563 3564 3565 3566
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		log.Info("unable to get flush state because of closed server")
		return resp, nil
	}

3567
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3568 3569 3570 3571
	if err != nil {
		log.Info("failed to get flush state response", zap.Error(err))
		return nil, err
	}
B
Bingyi Sun 已提交
3572 3573 3574 3575
	log.Info("received get flush state response", zap.Any("response", resp))
	return resp, err
}

C
Cai Yudong 已提交
3576 3577
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3578 3579 3580 3581
	code := node.stateCode.Load().(internalpb.StateCode)
	return code == internalpb.StateCode_Healthy
}

J
jingkl 已提交
3582
//unhealthyStatus returns the proxy not healthy status
3583 3584 3585
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3586
		Reason:    "proxy not healthy",
3587 3588
	}
}