diff --git a/internal/querynodev2/services.go b/internal/querynodev2/services.go index b987d723b3..e1dadc0b84 100644 --- a/internal/querynodev2/services.go +++ b/internal/querynodev2/services.go @@ -393,6 +393,8 @@ func (node *QueryNode) LoadPartitions(ctx context.Context, req *querypb.LoadPart PartitionIDs: req.GetPartitionIDs(), LoadType: querypb.LoadType_LoadCollection, // TODO: dyh, remove loadType in querynode }) + + log.Info("load partitions done") return merr.Status(nil), nil } @@ -483,13 +485,33 @@ func (node *QueryNode) ReleaseCollection(ctx context.Context, in *querypb.Releas } // ReleasePartitions clears all data related to this partition on the querynode -func (node *QueryNode) ReleasePartitions(ctx context.Context, in *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) { - if !node.lifetime.Add(commonpbutil.IsHealthyOrStopping) { - msg := fmt.Sprintf("query node %d is not ready", paramtable.GetNodeID()) - return util.WrapStatus(commonpb.ErrorCode_UnexpectedError, msg), nil +func (node *QueryNode) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) { + log := log.Ctx(ctx).With( + zap.Int64("collection", req.GetCollectionID()), + zap.Int64s("partitions", req.GetPartitionIDs()), + ) + + log.Info("received release partitions request") + + // check node healthy + if !node.lifetime.Add(commonpbutil.IsHealthy) { + err := fmt.Errorf("query node %d is not ready", paramtable.GetNodeID()) + status := &commonpb.Status{ + ErrorCode: commonpb.ErrorCode_UnexpectedError, + Reason: err.Error(), + } + return status, nil } defer node.lifetime.Done() + collection := node.manager.Collection.Get(req.GetCollectionID()) + if collection != nil { + for _, partition := range req.GetPartitionIDs() { + collection.RemovePartition(partition) + } + } + + log.Info("release partitions done") return util.SuccessStatus(), nil } diff --git a/internal/querynodev2/services_test.go b/internal/querynodev2/services_test.go index b015285bb4..b604d3d4c2 100644 --- a/internal/querynodev2/services_test.go +++ b/internal/querynodev2/services_test.go @@ -689,10 +689,18 @@ func (suite *ServiceSuite) TestReleaseCollection_Failed() { func (suite *ServiceSuite) TestReleasePartitions_Normal() { ctx := context.Background() - req := &querypb.ReleasePartitionsRequest{} + suite.TestLoadPartition() + req := &querypb.ReleasePartitionsRequest{ + CollectionID: suite.collectionID, + PartitionIDs: suite.partitionIDs, + } status, err := suite.node.ReleasePartitions(ctx, req) suite.NoError(err) suite.Equal(commonpb.ErrorCode_Success, status.GetErrorCode()) + collection := suite.node.manager.Collection.Get(suite.collectionID) + for _, partition := range suite.partitionIDs { + suite.False(collection.ExistPartition(partition)) + } } func (suite *ServiceSuite) TestReleasePartitions_Failed() {