Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

1028 lines
28 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
import (
"context"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
const (
CreateSnapshotTaskName = "CreateSnapshotTask"
DropSnapshotTaskName = "DropSnapshotTask"
DescribeSnapshotTaskName = "DescribeSnapshotTask"
ListSnapshotsTaskName = "ListSnapshotsTask"
RestoreSnapshotTaskName = "RestoreSnapshotTask"
GetRestoreSnapshotStateTaskName = "GetRestoreSnapshotStateTask"
ListRestoreSnapshotJobsTaskName = "ListRestoreSnapshotJobsTask"
PinSnapshotDataTaskName = "PinSnapshotDataTask"
UnpinSnapshotDataTaskName = "UnpinSnapshotDataTask"
)
// resolveCollectionNames resolves collection ID to (dbName, collectionName) via MetaCache.
// Returns empty strings on failure (best-effort for display purposes).
func resolveCollectionNames(ctx context.Context, collectionID int64) (string, string) {
if collectionID == 0 {
return "", ""
}
collInfo, err := globalMetaCache.GetCollectionInfo(ctx, "", "", collectionID)
if err != nil {
mlog.Warn(ctx, "failed to resolve collection names from ID",
mlog.FieldCollectionID(collectionID), mlog.Err(err))
return "", ""
}
return collInfo.dbName, collInfo.schema.Name
}
type createSnapshotTask struct {
baseTask
Condition
req *milvuspb.CreateSnapshotRequest
ctx context.Context
mixCoord types.MixCoordClient
result *commonpb.Status
collectionID UniqueID
}
func (cst *createSnapshotTask) TraceCtx() context.Context {
return cst.ctx
}
func (cst *createSnapshotTask) ID() UniqueID {
return cst.req.GetBase().GetMsgID()
}
func (cst *createSnapshotTask) SetID(uid UniqueID) {
cst.req.GetBase().MsgID = uid
}
func (cst *createSnapshotTask) Name() string {
return CreateSnapshotTaskName
}
func (cst *createSnapshotTask) Type() commonpb.MsgType {
return cst.req.GetBase().GetMsgType()
}
func (cst *createSnapshotTask) BeginTs() Timestamp {
return cst.req.GetBase().GetTimestamp()
}
func (cst *createSnapshotTask) EndTs() Timestamp {
return cst.req.GetBase().GetTimestamp()
}
func (cst *createSnapshotTask) SetTs(ts Timestamp) {
cst.req.Base.Timestamp = ts
}
func (cst *createSnapshotTask) OnEnqueue() error {
if cst.req.Base == nil {
cst.req.Base = commonpbutil.NewMsgBase()
}
cst.req.Base.MsgType = commonpb.MsgType_CreateSnapshot
cst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (cst *createSnapshotTask) PreExecute(ctx context.Context) error {
// Validate snapshot_name using standard naming rules
if err := ValidateSnapshotName(cst.req.GetName()); err != nil {
return err
}
// Validate compaction protection duration
maxCompactionProtectionSeconds := paramtable.Get().DataCoordCfg.SnapshotMaxCompactionProtectionSeconds.GetAsInt64()
if cst.req.GetCompactionProtectionSeconds() < 0 {
return merr.WrapErrParameterInvalidMsg("compaction_protection_seconds must be non-negative")
}
if cst.req.GetCompactionProtectionSeconds() > maxCompactionProtectionSeconds {
return merr.WrapErrParameterInvalidMsg("compaction_protection_seconds must not exceed %d", maxCompactionProtectionSeconds)
}
collectionID, err := globalMetaCache.GetCollectionID(ctx, cst.req.GetDbName(), cst.req.GetCollectionName())
if err != nil {
return err
}
cst.collectionID = collectionID
return nil
}
func (cst *createSnapshotTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy create snapshot",
mlog.String("snapshotName", cst.req.GetName()),
mlog.FieldCollectionName(cst.req.GetCollectionName()),
mlog.FieldCollectionID(cst.collectionID),
)
var err error
cst.result, err = cst.mixCoord.CreateSnapshot(ctx, &datapb.CreateSnapshotRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_CreateSnapshot),
),
Name: cst.req.GetName(),
Description: cst.req.GetDescription(),
CollectionId: cst.collectionID,
CompactionProtectionSeconds: cst.req.GetCompactionProtectionSeconds(),
})
if err = merr.CheckRPCCall(cst.result, err); err != nil {
return err
}
return nil
}
func (cst *createSnapshotTask) PostExecute(ctx context.Context) error {
return nil
}
type dropSnapshotTask struct {
baseTask
Condition
req *milvuspb.DropSnapshotRequest
ctx context.Context
mixCoord types.MixCoordClient
result *commonpb.Status
collectionID UniqueID
}
func (dst *dropSnapshotTask) TraceCtx() context.Context {
return dst.ctx
}
func (dst *dropSnapshotTask) ID() UniqueID {
return dst.req.GetBase().GetMsgID()
}
func (dst *dropSnapshotTask) SetID(uid UniqueID) {
dst.req.GetBase().MsgID = uid
}
func (dst *dropSnapshotTask) Name() string {
return DropSnapshotTaskName
}
func (dst *dropSnapshotTask) Type() commonpb.MsgType {
return dst.req.GetBase().GetMsgType()
}
func (dst *dropSnapshotTask) BeginTs() Timestamp {
return dst.req.GetBase().GetTimestamp()
}
func (dst *dropSnapshotTask) EndTs() Timestamp {
return dst.req.GetBase().GetTimestamp()
}
func (dst *dropSnapshotTask) SetTs(ts Timestamp) {
dst.req.Base.Timestamp = ts
}
func (dst *dropSnapshotTask) OnEnqueue() error {
if dst.req.Base == nil {
dst.req.Base = commonpbutil.NewMsgBase()
}
dst.req.Base.MsgType = commonpb.MsgType_DropSnapshot
dst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (dst *dropSnapshotTask) PreExecute(ctx context.Context) error {
// Validate snapshot_name using standard naming rules
if err := ValidateSnapshotName(dst.req.GetName()); err != nil {
return err
}
if dst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for drop snapshot")
}
collectionID, err := globalMetaCache.GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
if err != nil {
return err
}
dst.collectionID = collectionID
return nil
}
func (dst *dropSnapshotTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy drop snapshot",
mlog.String("snapshotName", dst.req.GetName()),
mlog.FieldCollectionName(dst.req.GetCollectionName()),
mlog.FieldCollectionID(dst.collectionID),
)
var err error
dst.result, err = dst.mixCoord.DropSnapshot(ctx, &datapb.DropSnapshotRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_DropSnapshot),
),
Name: dst.req.GetName(),
CollectionId: dst.collectionID,
})
if err = merr.CheckRPCCall(dst.result, err); err != nil {
return err
}
return nil
}
func (dst *dropSnapshotTask) PostExecute(ctx context.Context) error {
return nil
}
type describeSnapshotTask struct {
baseTask
Condition
req *milvuspb.DescribeSnapshotRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.DescribeSnapshotResponse
collectionID UniqueID
}
func (dst *describeSnapshotTask) TraceCtx() context.Context {
return dst.ctx
}
func (dst *describeSnapshotTask) ID() UniqueID {
return dst.req.GetBase().GetMsgID()
}
func (dst *describeSnapshotTask) SetID(uid UniqueID) {
dst.req.GetBase().MsgID = uid
}
func (dst *describeSnapshotTask) Name() string {
return DescribeSnapshotTaskName
}
func (dst *describeSnapshotTask) Type() commonpb.MsgType {
return dst.req.GetBase().GetMsgType()
}
func (dst *describeSnapshotTask) BeginTs() Timestamp {
return dst.req.GetBase().GetTimestamp()
}
func (dst *describeSnapshotTask) EndTs() Timestamp {
return dst.req.GetBase().GetTimestamp()
}
func (dst *describeSnapshotTask) SetTs(ts Timestamp) {
dst.req.Base.Timestamp = ts
}
func (dst *describeSnapshotTask) OnEnqueue() error {
if dst.req.Base == nil {
dst.req.Base = commonpbutil.NewMsgBase()
}
dst.req.Base.MsgType = commonpb.MsgType_DescribeSnapshot
dst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (dst *describeSnapshotTask) PreExecute(ctx context.Context) error {
// Validate snapshot_name using standard naming rules
if err := ValidateSnapshotName(dst.req.GetName()); err != nil {
return err
}
if dst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for describe snapshot")
}
collectionID, err := globalMetaCache.GetCollectionID(ctx, dst.req.GetDbName(), dst.req.GetCollectionName())
if err != nil {
return err
}
dst.collectionID = collectionID
return nil
}
func (dst *describeSnapshotTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy describe snapshot",
mlog.String("snapshotName", dst.req.GetName()),
)
result, err := dst.mixCoord.DescribeSnapshot(ctx, &datapb.DescribeSnapshotRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_DescribeSnapshot),
),
Name: dst.req.GetName(),
CollectionId: dst.collectionID,
IncludeCollectionInfo: false,
})
if err = merr.CheckRPCCall(result, err); err != nil {
dst.result = &milvuspb.DescribeSnapshotResponse{
Status: merr.Status(err),
}
return err
}
snapshotInfo := result.GetSnapshotInfo()
collectionName, err := globalMetaCache.GetCollectionName(ctx, "", snapshotInfo.GetCollectionId())
if err != nil {
mlog.Warn(ctx, "DescribeSnapshot fail to get collection name",
mlog.Err(err))
return err
}
var partitionNames []string
for _, partitionID := range snapshotInfo.GetPartitionIds() {
partitionName, err := globalMetaCache.GetPartitionName(ctx, "", collectionName, partitionID)
if err != nil {
mlog.Warn(ctx, "DescribeSnapshot fail to get partition name",
mlog.FieldPartitionID(partitionID),
mlog.Err(err))
}
partitionNames = append(partitionNames, partitionName)
}
dst.result = &milvuspb.DescribeSnapshotResponse{
Status: result.GetStatus(),
Name: snapshotInfo.GetName(),
Description: snapshotInfo.GetDescription(),
CreateTs: snapshotInfo.GetCreateTs(),
CollectionName: collectionName,
PartitionNames: partitionNames,
S3Location: snapshotInfo.GetS3Location(),
}
return nil
}
func (dst *describeSnapshotTask) PostExecute(ctx context.Context) error {
return nil
}
type listSnapshotsTask struct {
baseTask
Condition
req *milvuspb.ListSnapshotsRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.ListSnapshotsResponse
collectionID UniqueID
dbID UniqueID
}
func (lst *listSnapshotsTask) TraceCtx() context.Context {
return lst.ctx
}
func (lst *listSnapshotsTask) ID() UniqueID {
return lst.req.GetBase().GetMsgID()
}
func (lst *listSnapshotsTask) SetID(uid UniqueID) {
lst.req.GetBase().MsgID = uid
}
func (lst *listSnapshotsTask) Name() string {
return ListSnapshotsTaskName
}
func (lst *listSnapshotsTask) Type() commonpb.MsgType {
return lst.req.GetBase().GetMsgType()
}
func (lst *listSnapshotsTask) BeginTs() Timestamp {
return lst.req.GetBase().GetTimestamp()
}
func (lst *listSnapshotsTask) EndTs() Timestamp {
return lst.req.GetBase().GetTimestamp()
}
func (lst *listSnapshotsTask) SetTs(ts Timestamp) {
lst.req.Base.Timestamp = ts
}
func (lst *listSnapshotsTask) OnEnqueue() error {
if lst.req.Base == nil {
lst.req.Base = commonpbutil.NewMsgBase()
}
lst.req.Base.MsgType = commonpb.MsgType_ListSnapshots
lst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (lst *listSnapshotsTask) PreExecute(ctx context.Context) error {
// Resolve database ID for db-level filtering
if lst.req.GetDbName() != "" {
dbInfo, err := globalMetaCache.GetDatabaseInfo(ctx, lst.req.GetDbName())
if err != nil {
return err
}
lst.dbID = dbInfo.dbID
}
if lst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for ListSnapshots")
}
collectionID, err := globalMetaCache.GetCollectionID(ctx, lst.req.GetDbName(), lst.req.GetCollectionName())
if err != nil {
return err
}
lst.collectionID = collectionID
return nil
}
func (lst *listSnapshotsTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy list snapshots",
mlog.FieldCollectionName(lst.req.GetCollectionName()),
mlog.FieldCollectionID(lst.collectionID),
)
result, err := lst.mixCoord.ListSnapshots(ctx, &datapb.ListSnapshotsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_ListSnapshots),
),
CollectionId: lst.collectionID,
DbId: lst.dbID,
})
if err = merr.CheckRPCCall(result, err); err != nil {
lst.result = &milvuspb.ListSnapshotsResponse{
Status: merr.Status(err),
}
return err
}
lst.result = &milvuspb.ListSnapshotsResponse{
Status: result.GetStatus(),
Snapshots: result.GetSnapshots(),
}
return nil
}
func (lst *listSnapshotsTask) PostExecute(ctx context.Context) error {
return nil
}
type restoreSnapshotTask struct {
baseTask
Condition
req *milvuspb.RestoreSnapshotRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.RestoreSnapshotResponse
collectionID UniqueID // source collection ID for per-collection snapshot lookup
}
func (rst *restoreSnapshotTask) TraceCtx() context.Context {
return rst.ctx
}
func (rst *restoreSnapshotTask) ID() UniqueID {
return rst.req.GetBase().GetMsgID()
}
func (rst *restoreSnapshotTask) SetID(uid UniqueID) {
rst.req.GetBase().MsgID = uid
}
func (rst *restoreSnapshotTask) Name() string {
return RestoreSnapshotTaskName
}
func (rst *restoreSnapshotTask) Type() commonpb.MsgType {
return rst.req.GetBase().GetMsgType()
}
func (rst *restoreSnapshotTask) BeginTs() Timestamp {
return rst.req.GetBase().GetTimestamp()
}
func (rst *restoreSnapshotTask) EndTs() Timestamp {
return rst.req.GetBase().GetTimestamp()
}
func (rst *restoreSnapshotTask) SetTs(ts Timestamp) {
rst.req.Base.Timestamp = ts
}
func (rst *restoreSnapshotTask) OnEnqueue() error {
if rst.req.Base == nil {
rst.req.Base = commonpbutil.NewMsgBase()
}
rst.req.Base.MsgType = commonpb.MsgType_RestoreSnapshot
rst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (rst *restoreSnapshotTask) PreExecute(ctx context.Context) error {
// Validate snapshot_name using standard naming rules
if err := ValidateSnapshotName(rst.req.GetName()); err != nil {
return err
}
// Validate source collection name
if rst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for restore snapshot")
}
// Validate target collection name (required, cheap checks before RPC)
if rst.req.GetTargetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("target_collection_name is required for restore snapshot")
}
if err := ValidateCollectionName(rst.req.GetTargetCollectionName()); err != nil {
return err
}
// Resolve source collection ID for per-collection snapshot lookup (RPC call)
collectionID, err := globalMetaCache.GetCollectionID(ctx, rst.req.GetDbName(), rst.req.GetCollectionName())
if err != nil {
return err
}
rst.collectionID = collectionID
return nil
}
func (rst *restoreSnapshotTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy restore snapshot",
mlog.String("snapshotName", rst.req.GetName()),
mlog.String("sourceCollection", rst.req.GetCollectionName()),
mlog.String("sourceDb", rst.req.GetDbName()),
mlog.String("targetCollection", rst.req.GetTargetCollectionName()),
mlog.String("targetDb", rst.req.GetTargetDbName()),
)
// Delegate directly to DataCoord which handles the entire restore process
resp, err := rst.mixCoord.RestoreSnapshot(ctx, &datapb.RestoreSnapshotRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_RestoreSnapshot),
),
Name: rst.req.GetName(),
TargetDbName: rst.req.GetTargetDbName(),
TargetCollectionName: rst.req.GetTargetCollectionName(),
SourceCollectionId: rst.collectionID,
})
if err = merr.CheckRPCCall(resp, err); err != nil {
mlog.Warn(ctx, "RestoreSnapshot failed",
mlog.Err(err))
rst.result = &milvuspb.RestoreSnapshotResponse{Status: merr.Status(err)}
return err
}
rst.result = &milvuspb.RestoreSnapshotResponse{
Status: merr.Success(),
JobId: resp.GetJobId(),
}
return nil
}
func (rst *restoreSnapshotTask) PostExecute(ctx context.Context) error {
return nil
}
type getRestoreSnapshotStateTask struct {
baseTask
Condition
req *milvuspb.GetRestoreSnapshotStateRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.GetRestoreSnapshotStateResponse
}
func (grst *getRestoreSnapshotStateTask) TraceCtx() context.Context {
return grst.ctx
}
func (grst *getRestoreSnapshotStateTask) ID() UniqueID {
return grst.req.GetBase().GetMsgID()
}
func (grst *getRestoreSnapshotStateTask) SetID(uid UniqueID) {
grst.req.GetBase().MsgID = uid
}
func (grst *getRestoreSnapshotStateTask) Name() string {
return GetRestoreSnapshotStateTaskName
}
func (grst *getRestoreSnapshotStateTask) Type() commonpb.MsgType {
return grst.req.GetBase().GetMsgType()
}
func (grst *getRestoreSnapshotStateTask) BeginTs() Timestamp {
return grst.req.GetBase().GetTimestamp()
}
func (grst *getRestoreSnapshotStateTask) EndTs() Timestamp {
return grst.req.GetBase().GetTimestamp()
}
func (grst *getRestoreSnapshotStateTask) SetTs(ts Timestamp) {
grst.req.Base.Timestamp = ts
}
func (grst *getRestoreSnapshotStateTask) OnEnqueue() error {
if grst.req.Base == nil {
grst.req.Base = commonpbutil.NewMsgBase()
}
grst.req.Base.MsgType = commonpb.MsgType_GetRestoreSnapshotState
grst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (grst *getRestoreSnapshotStateTask) PreExecute(ctx context.Context) error {
// No additional validation needed for get restore snapshot state
return nil
}
func (grst *getRestoreSnapshotStateTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy get restore snapshot state",
mlog.FieldJobID(grst.req.GetJobId()),
)
result, err := grst.mixCoord.GetRestoreSnapshotState(ctx, &datapb.GetRestoreSnapshotStateRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_Undefined),
),
JobId: grst.req.GetJobId(),
})
if err = merr.CheckRPCCall(result, err); err != nil {
grst.result = &milvuspb.GetRestoreSnapshotStateResponse{
Status: merr.Status(err),
}
return err
}
// Convert datapb.RestoreSnapshotInfo to milvuspb.RestoreSnapshotInfo
info := result.GetInfo()
var milvusInfo *milvuspb.RestoreSnapshotInfo
if info != nil {
dbName, collectionName := resolveCollectionNames(ctx, info.GetCollectionId())
milvusInfo = &milvuspb.RestoreSnapshotInfo{
JobId: info.GetJobId(),
SnapshotName: info.GetSnapshotName(),
DbName: dbName,
CollectionName: collectionName,
State: milvuspb.RestoreSnapshotState(info.GetState()),
Progress: info.GetProgress(),
Reason: info.GetReason(),
StartTime: info.GetStartTime(),
TimeCost: info.GetTimeCost(),
}
}
grst.result = &milvuspb.GetRestoreSnapshotStateResponse{
Status: result.GetStatus(),
Info: milvusInfo,
}
return nil
}
func (grst *getRestoreSnapshotStateTask) PostExecute(ctx context.Context) error {
return nil
}
type listRestoreSnapshotJobsTask struct {
baseTask
Condition
req *milvuspb.ListRestoreSnapshotJobsRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.ListRestoreSnapshotJobsResponse
collectionID UniqueID
dbID UniqueID
}
func (lrst *listRestoreSnapshotJobsTask) TraceCtx() context.Context {
return lrst.ctx
}
func (lrst *listRestoreSnapshotJobsTask) ID() UniqueID {
return lrst.req.GetBase().GetMsgID()
}
func (lrst *listRestoreSnapshotJobsTask) SetID(uid UniqueID) {
lrst.req.GetBase().MsgID = uid
}
func (lrst *listRestoreSnapshotJobsTask) Name() string {
return ListRestoreSnapshotJobsTaskName
}
func (lrst *listRestoreSnapshotJobsTask) Type() commonpb.MsgType {
return lrst.req.GetBase().GetMsgType()
}
func (lrst *listRestoreSnapshotJobsTask) BeginTs() Timestamp {
return lrst.req.GetBase().GetTimestamp()
}
func (lrst *listRestoreSnapshotJobsTask) EndTs() Timestamp {
return lrst.req.GetBase().GetTimestamp()
}
func (lrst *listRestoreSnapshotJobsTask) SetTs(ts Timestamp) {
lrst.req.Base.Timestamp = ts
}
func (lrst *listRestoreSnapshotJobsTask) OnEnqueue() error {
if lrst.req.Base == nil {
lrst.req.Base = commonpbutil.NewMsgBase()
}
lrst.req.Base.MsgType = commonpb.MsgType_ListRestoreSnapshotJobs
lrst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (lrst *listRestoreSnapshotJobsTask) PreExecute(ctx context.Context) error {
// Resolve database ID for db-level filtering
if lrst.req.GetDbName() != "" {
dbInfo, err := globalMetaCache.GetDatabaseInfo(ctx, lrst.req.GetDbName())
if err != nil {
return err
}
lrst.dbID = dbInfo.dbID
}
if lrst.req.GetCollectionName() != "" {
collectionID, err := globalMetaCache.GetCollectionID(ctx, lrst.req.GetDbName(), lrst.req.GetCollectionName())
if err != nil {
return err
}
lrst.collectionID = collectionID
}
return nil
}
func (lrst *listRestoreSnapshotJobsTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy list restore snapshot jobs",
mlog.FieldCollectionName(lrst.req.GetCollectionName()),
mlog.FieldCollectionID(lrst.collectionID),
)
result, err := lrst.mixCoord.ListRestoreSnapshotJobs(ctx, &datapb.ListRestoreSnapshotJobsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_Undefined),
),
CollectionId: lrst.collectionID,
DbId: lrst.dbID,
})
if err = merr.CheckRPCCall(result, err); err != nil {
lrst.result = &milvuspb.ListRestoreSnapshotJobsResponse{
Status: merr.Status(err),
}
return err
}
// Convert datapb.RestoreSnapshotInfo to milvuspb.RestoreSnapshotInfo
jobs := result.GetJobs()
milvusJobs := make([]*milvuspb.RestoreSnapshotInfo, 0, len(jobs))
for _, job := range jobs {
dbName, collectionName := resolveCollectionNames(ctx, job.GetCollectionId())
milvusJobs = append(milvusJobs, &milvuspb.RestoreSnapshotInfo{
JobId: job.GetJobId(),
SnapshotName: job.GetSnapshotName(),
DbName: dbName,
CollectionName: collectionName,
State: milvuspb.RestoreSnapshotState(job.GetState()),
Progress: job.GetProgress(),
Reason: job.GetReason(),
StartTime: job.GetStartTime(),
TimeCost: job.GetTimeCost(),
})
}
lrst.result = &milvuspb.ListRestoreSnapshotJobsResponse{
Status: result.GetStatus(),
Jobs: milvusJobs,
}
return nil
}
func (lrst *listRestoreSnapshotJobsTask) PostExecute(ctx context.Context) error {
return nil
}
// pinSnapshotDataTask pins snapshot data to prevent GC from cleaning up segments
// referenced by a snapshot. Accepts milvuspb request from user, resolves
// collection_name to collection_id, and forwards to DataCoord.
type pinSnapshotDataTask struct {
baseTask
Condition
req *milvuspb.PinSnapshotDataRequest
ctx context.Context
mixCoord types.MixCoordClient
result *milvuspb.PinSnapshotDataResponse
collectionID UniqueID
}
func (pst *pinSnapshotDataTask) TraceCtx() context.Context {
return pst.ctx
}
func (pst *pinSnapshotDataTask) ID() UniqueID {
return pst.req.GetBase().GetMsgID()
}
func (pst *pinSnapshotDataTask) SetID(uid UniqueID) {
pst.req.GetBase().MsgID = uid
}
func (pst *pinSnapshotDataTask) Name() string {
return PinSnapshotDataTaskName
}
func (pst *pinSnapshotDataTask) Type() commonpb.MsgType {
return pst.req.GetBase().GetMsgType()
}
func (pst *pinSnapshotDataTask) BeginTs() Timestamp {
return pst.req.GetBase().GetTimestamp()
}
func (pst *pinSnapshotDataTask) EndTs() Timestamp {
return pst.req.GetBase().GetTimestamp()
}
func (pst *pinSnapshotDataTask) SetTs(ts Timestamp) {
pst.req.Base.Timestamp = ts
}
func (pst *pinSnapshotDataTask) OnEnqueue() error {
if pst.req.Base == nil {
pst.req.Base = commonpbutil.NewMsgBase()
}
pst.req.Base.MsgType = commonpb.MsgType_PinSnapshotData
pst.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
const maxPinTTLSeconds = 30 * 24 * 3600 // 30 days
func (pst *pinSnapshotDataTask) PreExecute(ctx context.Context) error {
if err := ValidateSnapshotName(pst.req.GetName()); err != nil {
return err
}
if pst.req.GetCollectionName() == "" {
return merr.WrapErrParameterMissingMsg("collection_name is required for pin snapshot data")
}
if pst.req.GetTtlSeconds() < 0 {
return merr.WrapErrParameterInvalidMsg("ttl_seconds must be non-negative")
}
if pst.req.GetTtlSeconds() > maxPinTTLSeconds {
return merr.WrapErrParameterInvalidMsg("ttl_seconds exceeds maximum of %d (30 days)", maxPinTTLSeconds)
}
collectionID, err := globalMetaCache.GetCollectionID(ctx, pst.req.GetDbName(), pst.req.GetCollectionName())
if err != nil {
return err
}
pst.collectionID = collectionID
return nil
}
func (pst *pinSnapshotDataTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy pin snapshot data",
mlog.String("snapshotName", pst.req.GetName()),
mlog.FieldCollectionName(pst.req.GetCollectionName()),
mlog.FieldCollectionID(pst.collectionID),
)
resp, err := pst.mixCoord.PinSnapshotData(ctx, &datapb.PinSnapshotDataRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_PinSnapshotData),
),
Name: pst.req.GetName(),
CollectionId: pst.collectionID,
TtlSeconds: pst.req.GetTtlSeconds(),
})
if err = merr.CheckRPCCall(resp, err); err != nil {
pst.result = &milvuspb.PinSnapshotDataResponse{
Status: merr.Status(err),
}
return err
}
pst.result = &milvuspb.PinSnapshotDataResponse{
Status: merr.Success(),
PinId: resp.GetPinId(),
}
return nil
}
func (pst *pinSnapshotDataTask) PostExecute(ctx context.Context) error {
return nil
}
// unpinSnapshotDataTask unpins previously pinned snapshot data, allowing GC
// to clean up segments if no other pins reference them.
type unpinSnapshotDataTask struct {
baseTask
Condition
req *milvuspb.UnpinSnapshotDataRequest
ctx context.Context
mixCoord types.MixCoordClient
result *commonpb.Status
}
func (ust *unpinSnapshotDataTask) TraceCtx() context.Context {
return ust.ctx
}
func (ust *unpinSnapshotDataTask) ID() UniqueID {
return ust.req.GetBase().GetMsgID()
}
func (ust *unpinSnapshotDataTask) SetID(uid UniqueID) {
ust.req.GetBase().MsgID = uid
}
func (ust *unpinSnapshotDataTask) Name() string {
return UnpinSnapshotDataTaskName
}
func (ust *unpinSnapshotDataTask) Type() commonpb.MsgType {
return ust.req.GetBase().GetMsgType()
}
func (ust *unpinSnapshotDataTask) BeginTs() Timestamp {
return ust.req.GetBase().GetTimestamp()
}
func (ust *unpinSnapshotDataTask) EndTs() Timestamp {
return ust.req.GetBase().GetTimestamp()
}
func (ust *unpinSnapshotDataTask) SetTs(ts Timestamp) {
ust.req.Base.Timestamp = ts
}
func (ust *unpinSnapshotDataTask) OnEnqueue() error {
if ust.req.Base == nil {
ust.req.Base = commonpbutil.NewMsgBase()
}
ust.req.Base.MsgType = commonpb.MsgType_UnpinSnapshotData
ust.req.Base.SourceID = paramtable.GetNodeID()
return nil
}
func (ust *unpinSnapshotDataTask) PreExecute(ctx context.Context) error {
if ust.req.GetPinId() == 0 {
return merr.WrapErrParameterMissingMsg("pin_id is required for unpin snapshot data")
}
return nil
}
func (ust *unpinSnapshotDataTask) Execute(ctx context.Context) error {
mlog.Info(ctx, "proxy unpin snapshot data",
mlog.Int64("pinID", ust.req.GetPinId()),
)
var err error
ust.result, err = ust.mixCoord.UnpinSnapshotData(ctx, &datapb.UnpinSnapshotDataRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_UnpinSnapshotData),
),
PinId: ust.req.GetPinId(),
})
if err = merr.CheckRPCCall(ust.result, err); err != nil {
return err
}
return nil
}
func (ust *unpinSnapshotDataTask) PostExecute(ctx context.Context) error {
return nil
}