fix: [2.4] add the db information in the dml message

Signed-off-by: SimFG <[email protected]>
milvus-io · Dec 5, 2024 · c2569c0 · c2569c0
1 parent 33eee32
commit c2569c0
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/internal/proxy/meta_cache.go b/internal/proxy/meta_cache.go
@@ -42,6 +42,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/commonpbutil"
 	"github.com/milvus-io/milvus/pkg/util/conc"
+	"github.com/milvus-io/milvus/pkg/util/expr"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
@@ -347,6 +348,7 @@ func InitMetaCache(ctx context.Context, rootCoord types.RootCoordClient, queryCo
 	if err != nil {
 		return err
 	}
+	expr.Register("cache", globalMetaCache)
 
 	// The privilege info is a little more. And to get this info, the query operation of involving multiple table queries is required.
 	resp, err := rootCoord.ListPolicy(ctx, &internalpb.ListPolicyRequest{})

diff --git a/internal/proxy/msg_pack.go b/internal/proxy/msg_pack.go
@@ -58,6 +58,7 @@ func genInsertMsgsByPartition(ctx context.Context,
 			),
 			CollectionID:   insertMsg.CollectionID,
 			PartitionID:    partitionID,
+			DbName:         insertMsg.DbName,
 			CollectionName: insertMsg.CollectionName,
 			PartitionName:  partitionName,
 			SegmentID:      segmentID,

diff --git a/internal/proxy/task_delete.go b/internal/proxy/task_delete.go
@@ -55,6 +55,7 @@ type deleteTask struct {
 	primaryKeys      *schemapb.IDs
 	collectionID     UniqueID
 	partitionID      UniqueID
+	dbID             UniqueID
 	partitionKeyMode bool
 
 	// set by scheduler
@@ -145,14 +146,11 @@ func (dt *deleteTask) Execute(ctx context.Context) (err error) {
 
 	result, numRows, err := repackDeleteMsgByHash(
 		ctx,
-		dt.primaryKeys,
-		dt.vChannels,
-		dt.idAllocator,
-		dt.ts,
-		dt.collectionID,
-		dt.req.GetCollectionName(),
-		dt.partitionID,
-		dt.req.GetPartitionName(),
+		dt.primaryKeys, dt.vChannels,
+		dt.idAllocator, dt.ts,
+		dt.collectionID, dt.req.GetCollectionName(),
+		dt.partitionID, dt.req.GetPartitionName(),
+		dt.req.GetDbName(),
 	)
 	if err != nil {
 		return err
@@ -200,6 +198,7 @@ func repackDeleteMsgByHash(
 	collectionName string,
 	partitionID int64,
 	partitionName string,
+	dbName string,
 ) (map[uint32][]*msgstream.DeleteMsg, int64, error) {
 	maxSize := Params.PulsarCfg.MaxMessageSize.GetAsInt()
 	hashValues := typeutil.HashPK2Channels(primaryKeys, vChannels)
@@ -229,6 +228,7 @@ func repackDeleteMsgByHash(
 				PartitionID:    partitionID,
 				CollectionName: collectionName,
 				PartitionName:  partitionName,
+				DbName:         dbName,
 				PrimaryKeys:    &schemapb.IDs{},
 				ShardName:      vchannel,
 			},
@@ -408,6 +408,7 @@ func (dr *deleteRunner) produce(ctx context.Context, primaryKeys *schemapb.IDs)
 		partitionKeyMode: dr.partitionKeyMode,
 		vChannels:        dr.vChannels,
 		primaryKeys:      primaryKeys,
+		dbID:             dr.dbID,
 	}
 
 	if err := dr.queue.Enqueue(task); err != nil {