impl.go 108.9 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 2345 2346 2347 2348 2349
	queryRequest := &milvuspb.QueryRequest{
		DbName:         request.DbName,
		CollectionName: request.CollectionName,
		PartitionNames: request.PartitionNames,
		Expr:           request.Expr,
		OutputFields:   request.OutputFields,
2350
	}
2351

2352
	qt := &queryTask{
2353 2354 2355 2356 2357
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
			Base: &commonpb.MsgBase{
				MsgType:  commonpb.MsgType_Retrieve,
2358
				SourceID: Params.ProxyCfg.ProxyID,
2359
			},
2360
			ResultChannelID: strconv.FormatInt(Params.ProxyCfg.ProxyID, 10),
2361 2362 2363 2364 2365
		},
		resultBuf: make(chan []*internalpb.RetrieveResults),
		query:     queryRequest,
		chMgr:     node.chMgr,
		qc:        node.queryCoord,
2366 2367
	}

D
dragondriver 已提交
2368 2369 2370 2371 2372
	method := "Query"

	log.Debug(
		rpcReceived(method),
		zap.String("traceID", traceID),
2373
		zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2374 2375 2376 2377
		zap.String("db", queryRequest.DbName),
		zap.String("collection", queryRequest.CollectionName),
		zap.Any("partitions", queryRequest.PartitionNames))

D
dragondriver 已提交
2378 2379 2380 2381 2382 2383 2384 2385 2386 2387
	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),
			zap.String("db", queryRequest.DbName),
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames))

2388 2389 2390 2391 2392 2393
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2394 2395
	}

D
dragondriver 已提交
2396 2397 2398
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2399
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2400 2401 2402
		zap.Int64("MsgID", qt.ID()),
		zap.Uint64("BeginTs", qt.BeginTs()),
		zap.Uint64("EndTs", qt.EndTs()),
2403 2404 2405
		zap.String("db", queryRequest.DbName),
		zap.String("collection", queryRequest.CollectionName),
		zap.Any("partitions", queryRequest.PartitionNames))
D
dragondriver 已提交
2406 2407 2408 2409 2410 2411

	if err := qt.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
			zap.Error(err),
			zap.String("traceID", traceID),
2412
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2413 2414 2415
			zap.Int64("MsgID", qt.ID()),
			zap.Uint64("BeginTs", qt.BeginTs()),
			zap.Uint64("EndTs", qt.EndTs()),
2416 2417 2418
			zap.String("db", queryRequest.DbName),
			zap.String("collection", queryRequest.CollectionName),
			zap.Any("partitions", queryRequest.PartitionNames))
2419

2420 2421 2422 2423 2424 2425 2426
		return &milvuspb.QueryResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
	}
2427

D
dragondriver 已提交
2428 2429 2430 2431 2432 2433 2434 2435 2436 2437 2438
	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()),
		zap.String("db", queryRequest.DbName),
		zap.String("collection", queryRequest.CollectionName),
		zap.Any("partitions", queryRequest.PartitionNames))

2439 2440 2441 2442 2443
	return &milvuspb.QueryResults{
		Status:     qt.result.Status,
		FieldsData: qt.result.FieldsData,
	}, nil
}
2444

2445
// CreateAlias create alias for collection, then you can search the collection with alias.
Y
Yusup 已提交
2446 2447 2448 2449
func (node *Proxy) CreateAlias(ctx context.Context, request *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2450 2451 2452 2453 2454

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

Y
Yusup 已提交
2455 2456 2457 2458 2459 2460 2461
	cat := &CreateAliasTask{
		ctx:                ctx,
		Condition:          NewTaskCondition(ctx),
		CreateAliasRequest: request,
		rootCoord:          node.rootCoord,
	}

D
dragondriver 已提交
2462 2463 2464 2465 2466 2467 2468 2469 2470 2471 2472 2473 2474 2475 2476 2477 2478 2479 2480 2481
	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 已提交
2482 2483 2484 2485 2486 2487
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2488 2489 2490
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2491
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2492 2493 2494 2495
		zap.Int64("MsgID", cat.ID()),
		zap.Uint64("BeginTs", cat.BeginTs()),
		zap.Uint64("EndTs", cat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2496 2497
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
2498 2499 2500 2501

	if err := cat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2502
			zap.Error(err),
D
dragondriver 已提交
2503
			zap.String("traceID", traceID),
2504
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2505 2506 2507 2508
			zap.Int64("MsgID", cat.ID()),
			zap.Uint64("BeginTs", cat.BeginTs()),
			zap.Uint64("EndTs", cat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2509 2510 2511 2512 2513 2514 2515 2516 2517
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

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

D
dragondriver 已提交
2518 2519 2520 2521 2522 2523 2524 2525 2526 2527 2528
	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 已提交
2529 2530 2531
	return cat.result, nil
}

2532
// DropAlias alter the alias of collection.
Y
Yusup 已提交
2533 2534 2535 2536
func (node *Proxy) DropAlias(ctx context.Context, request *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2537 2538 2539 2540 2541

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

Y
Yusup 已提交
2542 2543 2544 2545 2546 2547 2548
	dat := &DropAliasTask{
		ctx:              ctx,
		Condition:        NewTaskCondition(ctx),
		DropAliasRequest: request,
		rootCoord:        node.rootCoord,
	}

D
dragondriver 已提交
2549 2550 2551 2552 2553 2554 2555 2556 2557 2558 2559 2560 2561 2562 2563 2564 2565 2566
	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 已提交
2567 2568 2569 2570 2571 2572
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2573 2574 2575
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2576
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2577 2578 2579 2580
		zap.Int64("MsgID", dat.ID()),
		zap.Uint64("BeginTs", dat.BeginTs()),
		zap.Uint64("EndTs", dat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2581
		zap.String("alias", request.Alias))
D
dragondriver 已提交
2582 2583 2584 2585

	if err := dat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2586
			zap.Error(err),
D
dragondriver 已提交
2587
			zap.String("traceID", traceID),
2588
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2589 2590 2591 2592
			zap.Int64("MsgID", dat.ID()),
			zap.Uint64("BeginTs", dat.BeginTs()),
			zap.Uint64("EndTs", dat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2593 2594 2595 2596 2597 2598 2599 2600
			zap.String("alias", request.Alias))

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

D
dragondriver 已提交
2601 2602 2603 2604 2605 2606 2607 2608 2609 2610
	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 已提交
2611 2612 2613
	return dat.result, nil
}

2614
// AlterAlias alter alias of collection.
Y
Yusup 已提交
2615 2616 2617 2618
func (node *Proxy) AlterAlias(ctx context.Context, request *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
	if !node.checkHealthy() {
		return unhealthyStatus(), nil
	}
D
dragondriver 已提交
2619 2620 2621 2622 2623

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

Y
Yusup 已提交
2624 2625 2626 2627 2628 2629 2630
	aat := &AlterAliasTask{
		ctx:               ctx,
		Condition:         NewTaskCondition(ctx),
		AlterAliasRequest: request,
		rootCoord:         node.rootCoord,
	}

D
dragondriver 已提交
2631 2632 2633 2634 2635 2636 2637 2638 2639 2640 2641 2642 2643 2644 2645 2646 2647 2648 2649 2650
	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 已提交
2651 2652 2653 2654 2655 2656
		return &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		}, nil
	}

D
dragondriver 已提交
2657 2658 2659
	log.Debug(
		rpcEnqueued(method),
		zap.String("traceID", traceID),
2660
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2661 2662 2663 2664
		zap.Int64("MsgID", aat.ID()),
		zap.Uint64("BeginTs", aat.BeginTs()),
		zap.Uint64("EndTs", aat.EndTs()),
		zap.String("db", request.DbName),
Y
Yusup 已提交
2665 2666
		zap.String("alias", request.Alias),
		zap.String("collection", request.CollectionName))
D
dragondriver 已提交
2667 2668 2669 2670

	if err := aat.WaitToFinish(); err != nil {
		log.Warn(
			rpcFailedToWaitToFinish(method),
Y
Yusup 已提交
2671
			zap.Error(err),
D
dragondriver 已提交
2672
			zap.String("traceID", traceID),
2673
			zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
2674 2675 2676 2677
			zap.Int64("MsgID", aat.ID()),
			zap.Uint64("BeginTs", aat.BeginTs()),
			zap.Uint64("EndTs", aat.EndTs()),
			zap.String("db", request.DbName),
Y
Yusup 已提交
2678 2679 2680 2681 2682 2683 2684 2685 2686
			zap.String("alias", request.Alias),
			zap.String("collection", request.CollectionName))

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

D
dragondriver 已提交
2687 2688 2689 2690 2691 2692 2693 2694 2695 2696 2697
	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 已提交
2698 2699 2700
	return aat.result, nil
}

2701
// CalcDistance calculates the distances between vectors.
2702
func (node *Proxy) CalcDistance(ctx context.Context, request *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
2703 2704 2705 2706 2707 2708
	if !node.checkHealthy() {
		return &milvuspb.CalcDistanceResults{
			Status: unhealthyStatus(),
		}, nil
	}

2709
	param, _ := funcutil.GetAttrByKeyFromRepeatedKV("metric", request.GetParams())
2710 2711 2712 2713 2714 2715 2716 2717
	metric, err := distance.ValidateMetricType(param)
	if err != nil {
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    err.Error(),
			},
		}, nil
2718 2719
	}

2720 2721 2722 2723
	sp, ctx := trace.StartSpanFromContextWithOperationName(ctx, "Proxy-CalcDistance")
	defer sp.Finish()
	traceID, _, _ := trace.InfoFromSpan(sp)

2724 2725
	query := func(ids *milvuspb.VectorIDs) (*milvuspb.QueryResults, error) {
		outputFields := []string{ids.FieldName}
2726

2727 2728 2729 2730 2731
		queryRequest := &milvuspb.QueryRequest{
			DbName:         "",
			CollectionName: ids.CollectionName,
			PartitionNames: ids.PartitionNames,
			OutputFields:   outputFields,
2732 2733
		}

2734
		qt := &queryTask{
2735 2736 2737 2738 2739
			ctx:       ctx,
			Condition: NewTaskCondition(ctx),
			RetrieveRequest: &internalpb.RetrieveRequest{
				Base: &commonpb.MsgBase{
					MsgType:  commonpb.MsgType_Retrieve,
2740
					SourceID: Params.ProxyCfg.ProxyID,
2741
				},
2742
				ResultChannelID: strconv.FormatInt(Params.ProxyCfg.ProxyID, 10),
2743
			},
2744
			resultBuf: make(chan []*internalpb.RetrieveResults),
2745
			query:     queryRequest,
2746
			chMgr:     node.chMgr,
2747
			qc:        node.queryCoord,
Y
yukun 已提交
2748
			ids:       ids.IdArray,
2749 2750
		}

2751
		err := node.sched.dqQueue.Enqueue(qt)
2752
		if err != nil {
2753 2754 2755
			log.Debug("CalcDistance queryTask failed to enqueue",
				zap.Error(err),
				zap.String("traceID", traceID),
2756
				zap.String("role", typeutil.ProxyRole),
2757 2758 2759 2760
				zap.String("db", queryRequest.DbName),
				zap.String("collection", queryRequest.CollectionName),
				zap.Any("partitions", queryRequest.PartitionNames))

2761 2762 2763 2764 2765
			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2766
			}, err
2767
		}
2768 2769 2770

		log.Debug("CalcDistance queryTask enqueued",
			zap.String("traceID", traceID),
2771
			zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2772 2773 2774 2775 2776 2777
			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))
2778 2779 2780 2781

		err = qt.WaitToFinish()
		if err != nil {
			log.Debug("CalcDistance queryTask failed to WaitToFinish",
G
godchen 已提交
2782
				zap.Error(err),
2783
				zap.String("traceID", traceID),
2784
				zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2785 2786 2787 2788 2789 2790
				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))
2791 2792 2793 2794 2795 2796

			return &milvuspb.QueryResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
2797
			}, err
2798
		}
2799 2800 2801

		log.Debug("CalcDistance queryTask Done",
			zap.String("traceID", traceID),
2802
			zap.String("role", typeutil.ProxyRole),
G
godchen 已提交
2803 2804 2805 2806 2807 2808
			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))
2809 2810

		return &milvuspb.QueryResults{
2811 2812
			Status:     qt.result.Status,
			FieldsData: qt.result.FieldsData,
2813 2814 2815
		}, nil
	}

2816 2817 2818 2819 2820 2821 2822 2823 2824 2825 2826 2827 2828 2829
	// 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 {
2830
			return nil, errors.New("failed to fetch vectors")
2831 2832 2833 2834 2835 2836 2837 2838 2839 2840 2841 2842 2843 2844 2845 2846
		}

		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))
2847
					return nil, errors.New("failed to fetch vectors by id: " + fmt.Sprintln(id))
2848 2849 2850 2851 2852 2853 2854 2855 2856 2857 2858 2859 2860 2861 2862 2863 2864 2865 2866 2867 2868 2869 2870 2871 2872 2873
				}
				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))
2874
					return nil, errors.New("failed to fetch vectors by id: " + fmt.Sprintln(id))
2875 2876 2877 2878 2879 2880 2881 2882 2883 2884 2885 2886
				}
				result = append(result, binaryArr[int64(index)*element:int64(index+1)*element]...)
			}

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

2887
		return nil, errors.New("failed to fetch vectors")
2888 2889
	}

2890 2891
	log.Debug("CalcDistance received",
		zap.String("traceID", traceID),
2892
		zap.String("role", typeutil.ProxyRole),
2893
		zap.String("metric", metric))
G
godchen 已提交
2894

2895 2896 2897
	vectorsLeft := request.GetOpLeft().GetDataArray()
	opLeft := request.GetOpLeft().GetIdArray()
	if opLeft != nil {
2898 2899
		log.Debug("OpLeft IdArray not empty, Get vectors by id",
			zap.String("traceID", traceID),
2900
			zap.String("role", typeutil.ProxyRole))
2901

2902
		result, err := query(opLeft)
2903
		if err != nil {
2904 2905 2906
			log.Debug("Failed to get left vectors by id",
				zap.Error(err),
				zap.String("traceID", traceID),
2907
				zap.String("role", typeutil.ProxyRole))
2908

2909 2910 2911 2912 2913 2914 2915 2916
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

2917 2918
		log.Debug("OpLeft IdArray not empty, Get vectors by id done",
			zap.String("traceID", traceID),
2919
			zap.String("role", typeutil.ProxyRole))
2920

2921 2922
		vectorsLeft, err = arrangeFunc(opLeft, result.FieldsData)
		if err != nil {
2923 2924 2925
			log.Debug("Failed to re-arrange left vectors",
				zap.Error(err),
				zap.String("traceID", traceID),
2926
				zap.String("role", typeutil.ProxyRole))
2927

2928 2929 2930 2931 2932 2933
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
2934
		}
2935 2936 2937

		log.Debug("Re-arrange left vectors done",
			zap.String("traceID", traceID),
2938
			zap.String("role", typeutil.ProxyRole))
2939 2940
	}

G
groot 已提交
2941
	if vectorsLeft == nil {
2942 2943 2944
		msg := "Left vectors array is empty"
		log.Debug(msg,
			zap.String("traceID", traceID),
2945
			zap.String("role", typeutil.ProxyRole))
2946

G
groot 已提交
2947 2948 2949
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
2950
				Reason:    msg,
G
groot 已提交
2951 2952 2953 2954
			},
		}, nil
	}

2955 2956 2957
	vectorsRight := request.GetOpRight().GetDataArray()
	opRight := request.GetOpRight().GetIdArray()
	if opRight != nil {
2958 2959
		log.Debug("OpRight IdArray not empty, Get vectors by id",
			zap.String("traceID", traceID),
2960
			zap.String("role", typeutil.ProxyRole))
2961

2962
		result, err := query(opRight)
2963
		if err != nil {
2964 2965 2966
			log.Debug("Failed to get right vectors by id",
				zap.Error(err),
				zap.String("traceID", traceID),
2967
				zap.String("role", typeutil.ProxyRole))
2968

2969 2970 2971 2972 2973 2974 2975 2976
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

2977 2978
		log.Debug("OpRight IdArray not empty, Get vectors by id done",
			zap.String("traceID", traceID),
2979
			zap.String("role", typeutil.ProxyRole))
2980

2981 2982
		vectorsRight, err = arrangeFunc(opRight, result.FieldsData)
		if err != nil {
2983 2984 2985
			log.Debug("Failed to re-arrange right vectors",
				zap.Error(err),
				zap.String("traceID", traceID),
2986
				zap.String("role", typeutil.ProxyRole))
2987

2988 2989 2990 2991 2992 2993
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
2994
		}
2995 2996 2997

		log.Debug("Re-arrange right vectors done",
			zap.String("traceID", traceID),
2998
			zap.String("role", typeutil.ProxyRole))
2999 3000
	}

G
groot 已提交
3001
	if vectorsRight == nil {
3002 3003 3004
		msg := "Right vectors array is empty"
		log.Debug(msg,
			zap.String("traceID", traceID),
3005
			zap.String("role", typeutil.ProxyRole))
3006

G
groot 已提交
3007 3008 3009
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3010
				Reason:    msg,
G
groot 已提交
3011 3012 3013 3014
			},
		}, nil
	}

3015
	if vectorsLeft.Dim != vectorsRight.Dim {
3016 3017 3018
		msg := "Vectors dimension is not equal"
		log.Debug(msg,
			zap.String("traceID", traceID),
3019
			zap.String("role", typeutil.ProxyRole))
3020

3021 3022 3023
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3024
				Reason:    msg,
3025 3026 3027 3028 3029
			},
		}, nil
	}

	if vectorsLeft.GetFloatVector() != nil && vectorsRight.GetFloatVector() != nil {
3030 3031
		distances, err := distance.CalcFloatDistance(vectorsLeft.Dim, vectorsLeft.GetFloatVector().Data, vectorsRight.GetFloatVector().Data, metric)
		if err != nil {
3032 3033 3034
			log.Debug("Failed to CalcFloatDistance",
				zap.Error(err),
				zap.String("traceID", traceID),
3035
				zap.String("role", typeutil.ProxyRole))
3036

3037 3038 3039 3040 3041 3042 3043 3044
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

3045 3046 3047
		log.Debug("CalcFloatDistance done",
			zap.Error(err),
			zap.String("traceID", traceID),
3048
			zap.String("role", typeutil.ProxyRole))
3049

3050 3051 3052 3053 3054 3055 3056 3057 3058 3059
		return &milvuspb.CalcDistanceResults{
			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
			Array: &milvuspb.CalcDistanceResults_FloatDist{
				FloatDist: &schemapb.FloatArray{
					Data: distances,
				},
			},
		}, nil
	}

3060
	if vectorsLeft.GetBinaryVector() != nil && vectorsRight.GetBinaryVector() != nil {
G
groot 已提交
3061
		hamming, err := distance.CalcHammingDistance(vectorsLeft.Dim, vectorsLeft.GetBinaryVector(), vectorsRight.GetBinaryVector())
3062
		if err != nil {
3063 3064 3065
			log.Debug("Failed to CalcHammingDistance",
				zap.Error(err),
				zap.String("traceID", traceID),
3066
				zap.String("role", typeutil.ProxyRole))
3067

3068 3069 3070 3071 3072 3073 3074 3075 3076
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    err.Error(),
				},
			}, nil
		}

		if metric == distance.HAMMING {
3077 3078
			log.Debug("CalcHammingDistance done",
				zap.String("traceID", traceID),
3079
				zap.String("role", typeutil.ProxyRole))
3080

3081 3082 3083 3084
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
				Array: &milvuspb.CalcDistanceResults_IntDist{
					IntDist: &schemapb.IntArray{
G
groot 已提交
3085
						Data: hamming,
3086 3087 3088 3089 3090 3091
					},
				},
			}, nil
		}

		if metric == distance.TANIMOTO {
G
groot 已提交
3092
			tanimoto, err := distance.CalcTanimotoCoefficient(vectorsLeft.Dim, hamming)
3093
			if err != nil {
3094 3095 3096
				log.Debug("Failed to CalcTanimotoCoefficient",
					zap.Error(err),
					zap.String("traceID", traceID),
3097
					zap.String("role", typeutil.ProxyRole))
3098

3099 3100 3101 3102 3103 3104 3105 3106
				return &milvuspb.CalcDistanceResults{
					Status: &commonpb.Status{
						ErrorCode: commonpb.ErrorCode_UnexpectedError,
						Reason:    err.Error(),
					},
				}, nil
			}

3107 3108
			log.Debug("CalcTanimotoCoefficient done",
				zap.String("traceID", traceID),
3109
				zap.String("role", typeutil.ProxyRole))
3110

3111 3112 3113 3114 3115 3116 3117 3118 3119 3120 3121
			return &milvuspb.CalcDistanceResults{
				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success, Reason: ""},
				Array: &milvuspb.CalcDistanceResults_FloatDist{
					FloatDist: &schemapb.FloatArray{
						Data: tanimoto,
					},
				},
			}, nil
		}
	}

3122
	err = errors.New("unexpected error")
3123
	if (vectorsLeft.GetBinaryVector() != nil && vectorsRight.GetFloatVector() != nil) || (vectorsLeft.GetFloatVector() != nil && vectorsRight.GetBinaryVector() != nil) {
3124
		err = errors.New("cannot calculate distance between binary vectors and float vectors")
3125 3126
	}

3127 3128 3129
	log.Debug("Failed to CalcDistance",
		zap.Error(err),
		zap.String("traceID", traceID),
3130
		zap.String("role", typeutil.ProxyRole))
3131

3132 3133 3134 3135 3136 3137 3138 3139
	return &milvuspb.CalcDistanceResults{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
			Reason:    err.Error(),
		},
	}, nil
}

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

3145
// GetPersistentSegmentInfo get the information of sealed segment.
C
Cai Yudong 已提交
3146
func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.GetPersistentSegmentInfoRequest) (*milvuspb.GetPersistentSegmentInfoResponse, error) {
D
dragondriver 已提交
3147
	log.Debug("GetPersistentSegmentInfo",
3148
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3149 3150 3151
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

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

J
jingkl 已提交
3200
// GetQuerySegmentInfo gets segment information from QueryCoord.
C
Cai Yudong 已提交
3201
func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQuerySegmentInfoRequest) (*milvuspb.GetQuerySegmentInfoResponse, error) {
D
dragondriver 已提交
3202
	log.Debug("GetQuerySegmentInfo",
3203
		zap.String("role", typeutil.ProxyRole),
D
dragondriver 已提交
3204 3205 3206
		zap.String("db", req.DbName),
		zap.Any("collection", req.CollectionName))

G
godchen 已提交
3207
	resp := &milvuspb.GetQuerySegmentInfoResponse{
Z
zhenshan.cao 已提交
3208
		Status: &commonpb.Status{
3209
			ErrorCode: commonpb.ErrorCode_UnexpectedError,
Z
zhenshan.cao 已提交
3210 3211
		},
	}
3212 3213 3214 3215
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		return resp, nil
	}
G
godchen 已提交
3216
	segments, err := node.getSegmentsOfCollection(ctx, req.DbName, req.CollectionName)
Z
zhenshan.cao 已提交
3217 3218 3219 3220
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3221 3222 3223 3224 3225
	collID, err := globalMetaCache.GetCollectionID(ctx, req.CollectionName)
	if err != nil {
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3226
	infoResp, err := node.queryCoord.GetSegmentInfo(ctx, &querypb.GetSegmentInfoRequest{
Z
zhenshan.cao 已提交
3227
		Base: &commonpb.MsgBase{
3228
			MsgType:   commonpb.MsgType_SegmentInfo,
Z
zhenshan.cao 已提交
3229 3230
			MsgID:     0,
			Timestamp: 0,
3231
			SourceID:  Params.ProxyCfg.ProxyID,
Z
zhenshan.cao 已提交
3232
		},
3233 3234
		CollectionID: collID,
		SegmentIDs:   segments,
Z
zhenshan.cao 已提交
3235 3236
	})
	if err != nil {
3237
		log.Error("Failed to get segment info from QueryCoord",
3238
			zap.Int64s("segmentIDs", segments), zap.Error(err))
Z
zhenshan.cao 已提交
3239 3240 3241
		resp.Status.Reason = err.Error()
		return resp, nil
	}
3242
	log.Debug("GetQuerySegmentInfo ", zap.Any("infos", infoResp.Infos), zap.Any("status", infoResp.Status))
3243
	if infoResp.Status.ErrorCode != commonpb.ErrorCode_Success {
3244
		log.Error("Failed to get segment info from QueryCoord", zap.String("errMsg", infoResp.Status.Reason))
Z
zhenshan.cao 已提交
3245 3246 3247 3248 3249 3250 3251 3252 3253 3254 3255 3256 3257
		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,
3258
			NodeID:       info.NodeID,
X
xige-16 已提交
3259
			State:        info.SegmentState,
Z
zhenshan.cao 已提交
3260 3261
		}
	}
3262
	resp.Status.ErrorCode = commonpb.ErrorCode_Success
Z
zhenshan.cao 已提交
3263 3264 3265 3266
	resp.Infos = queryInfos
	return resp, nil
}

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

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

J
jingkl 已提交
3326
// Dummy handles dummy request
C
Cai Yudong 已提交
3327
func (node *Proxy) Dummy(ctx context.Context, req *milvuspb.DummyRequest) (*milvuspb.DummyResponse, error) {
3328 3329 3330 3331 3332 3333 3334 3335 3336 3337 3338
	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
	}

3339 3340
	if drt.RequestType == "query" {
		drr, err := parseDummyQueryRequest(req.RequestType)
3341
		if err != nil {
3342
			log.Debug("Failed to parse dummy query request")
3343 3344 3345
			return failedResponse, nil
		}

3346
		request := &milvuspb.QueryRequest{
3347 3348 3349
			DbName:         drr.DbName,
			CollectionName: drr.CollectionName,
			PartitionNames: drr.PartitionNames,
3350
			OutputFields:   drr.OutputFields,
X
Xiangyu Wang 已提交
3351 3352
		}

3353
		_, err = node.Query(ctx, request)
3354
		if err != nil {
3355
			log.Debug("Failed to execute dummy query")
3356 3357
			return failedResponse, err
		}
X
Xiangyu Wang 已提交
3358 3359 3360 3361 3362 3363

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

3364 3365
	log.Debug("cannot find specify dummy request type")
	return failedResponse, nil
X
Xiangyu Wang 已提交
3366 3367
}

J
jingkl 已提交
3368
// RegisterLink registers a link
C
Cai Yudong 已提交
3369
func (node *Proxy) RegisterLink(ctx context.Context, req *milvuspb.RegisterLinkRequest) (*milvuspb.RegisterLinkResponse, error) {
G
godchen 已提交
3370
	code := node.stateCode.Load().(internalpb.StateCode)
D
dragondriver 已提交
3371
	log.Debug("RegisterLink",
3372
		zap.String("role", typeutil.ProxyRole),
C
Cai Yudong 已提交
3373
		zap.Any("state code of proxy", code))
D
dragondriver 已提交
3374

G
godchen 已提交
3375
	if code != internalpb.StateCode_Healthy {
3376 3377 3378
		return &milvuspb.RegisterLinkResponse{
			Address: nil,
			Status: &commonpb.Status{
3379
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3380
				Reason:    "proxy not healthy",
3381 3382 3383 3384 3385 3386
			},
		}, nil
	}
	return &milvuspb.RegisterLinkResponse{
		Address: nil,
		Status: &commonpb.Status{
3387
			ErrorCode: commonpb.ErrorCode_Success,
3388
			Reason:    os.Getenv(metricsinfo.DeployModeEnvKey),
3389 3390 3391
		},
	}, nil
}
3392

3393 3394 3395
// 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",
3396
		zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3397 3398 3399 3400
		zap.String("req", req.Request))

	if !node.checkHealthy() {
		log.Warn("Proxy.GetMetrics failed",
3401
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3402
			zap.String("req", req.Request),
3403
			zap.Error(errProxyIsUnhealthy(Params.ProxyCfg.ProxyID)))
3404 3405 3406 3407

		return &milvuspb.GetMetricsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
3408
				Reason:    msgProxyIsUnhealthy(Params.ProxyCfg.ProxyID),
3409 3410 3411 3412 3413 3414 3415 3416
			},
			Response: "",
		}, nil
	}

	metricType, err := metricsinfo.ParseMetricType(req.Request)
	if err != nil {
		log.Warn("Proxy.GetMetrics failed to parse metric type",
3417
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3418 3419 3420 3421 3422 3423 3424 3425 3426 3427 3428 3429 3430 3431 3432
			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 已提交
3433 3434 3435 3436 3437 3438 3439 3440 3441 3442
	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,
3443
		SourceID:  Params.ProxyCfg.ProxyID,
D
dragondriver 已提交
3444 3445
	}

3446
	if metricType == metricsinfo.SystemInfoMetrics {
3447 3448 3449 3450 3451 3452 3453
		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))

3454
		metrics, err := getSystemInfoMetrics(ctx, req, node)
3455 3456

		log.Debug("Proxy.GetMetrics",
3457
			zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3458 3459 3460 3461 3462
			zap.String("req", req.Request),
			zap.String("metric_type", metricType),
			zap.Any("metrics", metrics), // TODO(dragondriver): necessary? may be very large
			zap.Error(err))

3463 3464
		node.metricsCacheManager.UpdateSystemInfoMetrics(metrics)

G
godchen 已提交
3465
		return metrics, nil
3466 3467 3468
	}

	log.Debug("Proxy.GetMetrics failed, request metric type is not implemented yet",
3469
		zap.Int64("node_id", Params.ProxyCfg.ProxyID),
3470 3471 3472 3473 3474 3475 3476 3477 3478 3479 3480 3481
		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 已提交
3482 3483 3484
// 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",
3485
		zap.Int64("proxy_id", Params.ProxyCfg.ProxyID),
B
bigsheeper 已提交
3486 3487 3488 3489 3490 3491 3492 3493 3494 3495 3496 3497 3498 3499
		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,
3500
			SourceID:  Params.ProxyCfg.ProxyID,
B
bigsheeper 已提交
3501 3502 3503
		},
		SourceNodeIDs:    []int64{req.SrcNodeID},
		DstNodeIDs:       req.DstNodeIDs,
X
xige-16 已提交
3504
		BalanceReason:    querypb.TriggerCondition_GrpcRequest,
B
bigsheeper 已提交
3505 3506 3507 3508 3509 3510 3511 3512 3513 3514 3515 3516 3517 3518 3519 3520 3521 3522
		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 已提交
3523
//GetCompactionState gets the compaction state of multiple segments
3524 3525 3526 3527 3528 3529 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 3556 3557 3558 3559 3560 3561 3562
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
}

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 已提交
3563 3564 3565
// 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))
3566
	var err error
B
Bingyi Sun 已提交
3567 3568 3569 3570 3571 3572 3573
	resp := &milvuspb.GetFlushStateResponse{}
	if !node.checkHealthy() {
		resp.Status = unhealthyStatus()
		log.Info("unable to get flush state because of closed server")
		return resp, nil
	}

3574
	resp, err = node.dataCoord.GetFlushState(ctx, req)
X
Xiaofan 已提交
3575 3576 3577 3578
	if err != nil {
		log.Info("failed to get flush state response", zap.Error(err))
		return nil, err
	}
B
Bingyi Sun 已提交
3579 3580 3581 3582
	log.Info("received get flush state response", zap.Any("response", resp))
	return resp, err
}

C
Cai Yudong 已提交
3583 3584
// checkHealthy checks proxy state is Healthy
func (node *Proxy) checkHealthy() bool {
3585 3586 3587 3588
	code := node.stateCode.Load().(internalpb.StateCode)
	return code == internalpb.StateCode_Healthy
}

J
jingkl 已提交
3589
//unhealthyStatus returns the proxy not healthy status
3590 3591 3592
func unhealthyStatus() *commonpb.Status {
	return &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_UnexpectedError,
C
Cai Yudong 已提交
3593
		Reason:    "proxy not healthy",
3594 3595
	}
}