Add load collection for AssignSegmentID (#12077)

Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>

internal/datacoord/server_test.go

@@ -122,6 +122,10 @@ func TestAssignSegmentID(t *testing.T) {
 	t.Run("assign segment with invalid collection", func(t *testing.T) {
 		svr := newTestServer(t, nil)
 		defer closeTestServer(t, svr)
+		svr.rootCoordClient = &mockDescribeCollRoot{
+			RootCoord: svr.rootCoordClient,
+			collID:    collID,
+		}
 		schema := newTestSchema()
 		svr.meta.AddCollection(&datapb.CollectionInfo{
 			ID:         collID,
@@ -145,6 +149,23 @@ func TestAssignSegmentID(t *testing.T) {
 	})
 }
 
+type mockDescribeCollRoot struct {
+	types.RootCoord
+	collID UniqueID
+}
+
+func (r *mockDescribeCollRoot) DescribeCollection(ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
+	if req.CollectionID != r.collID {
+		return &milvuspb.DescribeCollectionResponse{
+			Status: &commonpb.Status{
+				ErrorCode: commonpb.ErrorCode_UnexpectedError,
+				Reason:    "Collection not found",
+			},
+		}, nil
+	}
+	return r.RootCoord.DescribeCollection(ctx, req)
+}
+
 func TestFlush(t *testing.T) {
 	req := &datapb.FlushRequest{
 		Base: &commonpb.MsgBase{

internal/datacoord/services.go

@@ -114,6 +114,14 @@ func (s *Server) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentI
 			zap.String("channelName", r.GetChannelName()),
 			zap.Uint32("count", r.GetCount()))
 
+		if s.meta.GetCollection(r.GetCollectionID()) == nil {
+			err := s.loadCollectionFromRootCoord(ctx, r.GetCollectionID())
+			if err != nil {
+				log.Warn("failed to load collection in alloc segment", zap.Any("request", r), zap.Error(err))
+				continue
+			}
+		}
+
 		s.cluster.Watch(r.ChannelName, r.CollectionID)
 
 		allocations, err := s.segmentManager.AllocSegment(ctx,