enhance: tidy code in import plugin context (#46374)

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
This commit is contained in:
XuanYang-cn 2026-01-04 16:51:23 +08:00 committed by GitHub
parent 27525d57cc
commit 146f173eea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 9 additions and 30 deletions

View File

@ -51,34 +51,10 @@ func PreAllocateBinlogIDs(allocator allocator.Allocator, segmentInfos []*Segment
return &datapb.IDRange{Begin: begin, End: end}, err
}
func WrapPluginContextWithImport(collectionID int64, dbProperties []*commonpb.KeyValuePair, options importutilv2.Options, msg proto.Message) {
pluginContext := make([]*commonpb.KeyValuePair, 0)
// Return None nil slice
func GetReadPluginContext(options importutilv2.Options) []*commonpb.KeyValuePair {
importEzk, _ := importutilv2.GetEZK(options)
readPluginContext := hookutil.GetReadStoragePluginContext(importEzk)
if readPluginContext != nil {
pluginContext = append(pluginContext, readPluginContext...)
}
writePluginContext := hookutil.GetStoragePluginContext(dbProperties, collectionID)
if writePluginContext != nil {
pluginContext = append(pluginContext, writePluginContext...)
}
if len(pluginContext) == 0 {
return
}
switch msg.(type) {
case *datapb.ImportRequest:
job := msg.(*datapb.ImportRequest)
job.PluginContext = append(job.PluginContext, pluginContext...)
case *datapb.PreImportRequest:
job := msg.(*datapb.PreImportRequest)
job.PluginContext = append(job.PluginContext, pluginContext...)
default:
return
}
return hookutil.GetReadStoragePluginContext(importEzk)
}
func WrapPluginContext(collectionID int64, properties []*commonpb.KeyValuePair, msg proto.Message) {

View File

@ -300,8 +300,9 @@ func AssemblePreImportRequest(task ImportTask, job ImportJob) *datapb.PreImportR
Options: job.GetOptions(),
TaskSlot: task.GetTaskSlot(),
StorageConfig: createStorageConfig(),
PluginContext: GetReadPluginContext(job.GetOptions()),
}
WrapPluginContextWithImport(task.GetCollectionID(), job.GetSchema().GetProperties(), job.GetOptions(), req)
WrapPluginContext(task.GetCollectionID(), job.GetSchema().GetProperties(), req)
return req
}
@ -379,9 +380,10 @@ func AssembleImportRequest(task ImportTask, job ImportJob, meta *meta, alloc all
StorageConfig: createStorageConfig(),
TaskSlot: task.GetTaskSlot(),
StorageVersion: storageVersion,
PluginContext: GetReadPluginContext(job.GetOptions()),
UseLoonFfi: Params.CommonCfg.UseLoonFFI.GetAsBool(),
}
WrapPluginContextWithImport(task.GetCollectionID(), job.GetSchema().GetProperties(), job.GetOptions(), req)
WrapPluginContext(task.GetCollectionID(), job.GetSchema().GetProperties(), req)
return req, nil
}

View File

@ -249,11 +249,12 @@ func GetStoragePluginContext(collProps []*commonpb.KeyValuePair, collectionID in
return nil
}
// Non nill return
func GetReadStoragePluginContext(importEzk string) []*commonpb.KeyValuePair {
readContext, err := ImportEZ(importEzk)
if err != nil {
log.Error("failed to import ezk", zap.Error(err))
return nil
return []*commonpb.KeyValuePair{}
}
return readContext
}