feat: Define import API and metadata (#28731)

Define the new rpc and metadata for ImportV2. see also: https://github.com/milvus-io/milvus/issues/28521 --------- Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
2023-12-04 19:56:35 +08:00 · 2023-12-04 19:56:35 +08:00 · d26b563a8b
parent 393b1f943c
commit d26b563a8b
10 changed files with 2589 additions and 406 deletions
--- a/internal/datacoord/mock_test.go
+++ b/internal/datacoord/mock_test.go
@ -295,6 +295,26 @@ func (c *mockDataNodeClient) CheckChannelOperationProgress(ctx context.Context,
 	return &datapb.ChannelOperationProgressResponse{Status: merr.Success()}, nil
 }

+func (c *mockDataNodeClient) PreImport(ctx context.Context, req *datapb.PreImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
+}
+
+func (c *mockDataNodeClient) ImportV2(ctx context.Context, req *datapb.ImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
+}
+
+func (c *mockDataNodeClient) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest, opts ...grpc.CallOption) (*datapb.QueryPreImportResponse, error) {
+	return &datapb.QueryPreImportResponse{Status: merr.Success()}, nil
+}
+
+func (c *mockDataNodeClient) QueryImport(ctx context.Context, req *datapb.QueryImportRequest, opts ...grpc.CallOption) (*datapb.QueryImportResponse, error) {
+	return &datapb.QueryImportResponse{Status: merr.Success()}, nil
+}
+
+func (c *mockDataNodeClient) DropImport(ctx context.Context, req *datapb.DropImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
+}
+
 func (c *mockDataNodeClient) Stop() error {
 	c.state = commonpb.StateCode_Abnormal
 	return nil
--- a/internal/datanode/services.go
+++ b/internal/datanode/services.go
@ -943,3 +943,23 @@ func logDupFlush(cID, segID int64) {
 		zap.Int64("collectionID", cID),
 		zap.Int64("segmentID", segID))
 }
+
+func (node *DataNode) PreImport(ctx context.Context, req *datapb.PreImportRequest) (*commonpb.Status, error) {
+	return nil, merr.ErrServiceUnimplemented
+}
+
+func (node *DataNode) ImportV2(ctx context.Context, req *datapb.ImportRequest) (*commonpb.Status, error) {
+	return nil, merr.ErrServiceUnimplemented
+}
+
+func (node *DataNode) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error) {
+	return nil, merr.ErrServiceUnimplemented
+}
+
+func (node *DataNode) QueryImport(ctx context.Context, req *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error) {
+	return nil, merr.ErrServiceUnimplemented
+}
+
+func (node *DataNode) DropImport(ctx context.Context, req *datapb.DropImportRequest) (*commonpb.Status, error) {
+	return nil, merr.ErrServiceUnimplemented
+}
--- a/internal/distributed/datanode/client/client.go
+++ b/internal/distributed/datanode/client/client.go
@ -244,3 +244,33 @@ func (c *Client) CheckChannelOperationProgress(ctx context.Context, req *datapb.
 		return client.CheckChannelOperationProgress(ctx, req)
 	})
 }
+
+func (c *Client) PreImport(ctx context.Context, req *datapb.PreImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return wrapGrpcCall(ctx, c, func(client datapb.DataNodeClient) (*commonpb.Status, error) {
+		return client.PreImport(ctx, req)
+	})
+}
+
+func (c *Client) ImportV2(ctx context.Context, req *datapb.ImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return wrapGrpcCall(ctx, c, func(client datapb.DataNodeClient) (*commonpb.Status, error) {
+		return client.ImportV2(ctx, req)
+	})
+}
+
+func (c *Client) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest, opts ...grpc.CallOption) (*datapb.QueryPreImportResponse, error) {
+	return wrapGrpcCall(ctx, c, func(client datapb.DataNodeClient) (*datapb.QueryPreImportResponse, error) {
+		return client.QueryPreImport(ctx, req)
+	})
+}
+
+func (c *Client) QueryImport(ctx context.Context, req *datapb.QueryImportRequest, opts ...grpc.CallOption) (*datapb.QueryImportResponse, error) {
+	return wrapGrpcCall(ctx, c, func(client datapb.DataNodeClient) (*datapb.QueryImportResponse, error) {
+		return client.QueryImport(ctx, req)
+	})
+}
+
+func (c *Client) DropImport(ctx context.Context, req *datapb.DropImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return wrapGrpcCall(ctx, c, func(client datapb.DataNodeClient) (*commonpb.Status, error) {
+		return client.DropImport(ctx, req)
+	})
+}
--- a/internal/distributed/datanode/service.go
+++ b/internal/distributed/datanode/service.go
@ -379,3 +379,23 @@ func (s *Server) NotifyChannelOperation(ctx context.Context, req *datapb.Channel
 func (s *Server) CheckChannelOperationProgress(ctx context.Context, req *datapb.ChannelWatchInfo) (*datapb.ChannelOperationProgressResponse, error) {
 	return s.datanode.CheckChannelOperationProgress(ctx, req)
 }
+
+func (s *Server) PreImport(ctx context.Context, req *datapb.PreImportRequest) (*commonpb.Status, error) {
+	return s.datanode.PreImport(ctx, req)
+}
+
+func (s *Server) ImportV2(ctx context.Context, req *datapb.ImportRequest) (*commonpb.Status, error) {
+	return s.datanode.ImportV2(ctx, req)
+}
+
+func (s *Server) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error) {
+	return s.datanode.QueryPreImport(ctx, req)
+}
+
+func (s *Server) QueryImport(ctx context.Context, req *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error) {
+	return s.datanode.QueryImport(ctx, req)
+}
+
+func (s *Server) DropImport(ctx context.Context, req *datapb.DropImportRequest) (*commonpb.Status, error) {
+	return s.datanode.DropImport(ctx, req)
+}
--- a/internal/distributed/datanode/service_test.go
+++ b/internal/distributed/datanode/service_test.go
@ -162,6 +162,26 @@ func (m *MockDataNode) CheckChannelOperationProgress(ctx context.Context, req *d
 	return &datapb.ChannelOperationProgressResponse{}, m.err
 }

+func (m *MockDataNode) PreImport(ctx context.Context, req *datapb.PreImportRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockDataNode) ImportV2(ctx context.Context, req *datapb.ImportRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
+func (m *MockDataNode) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error) {
+	return &datapb.QueryPreImportResponse{}, m.err
+}
+
+func (m *MockDataNode) QueryImport(ctx context.Context, req *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error) {
+	return &datapb.QueryImportResponse{}, m.err
+}
+
+func (m *MockDataNode) DropImport(ctx context.Context, req *datapb.DropImportRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
 // /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 type mockDataCoord struct {
 	types.DataCoordClient
--- a/internal/mocks/mock_datanode.go
+++ b/internal/mocks/mock_datanode.go
@ -64,8 +64,8 @@ type MockDataNode_AddImportSegment_Call struct {
 }

 // AddImportSegment is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.AddImportSegmentRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.AddImportSegmentRequest
 func (_e *MockDataNode_Expecter) AddImportSegment(_a0 interface{}, _a1 interface{}) *MockDataNode_AddImportSegment_Call {
 	return &MockDataNode_AddImportSegment_Call{Call: _e.mock.On("AddImportSegment", _a0, _a1)}
 }
@ -119,8 +119,8 @@ type MockDataNode_CheckChannelOperationProgress_Call struct {
 }

 // CheckChannelOperationProgress is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.ChannelWatchInfo
+//  - _a0 context.Context
+//  - _a1 *datapb.ChannelWatchInfo
 func (_e *MockDataNode_Expecter) CheckChannelOperationProgress(_a0 interface{}, _a1 interface{}) *MockDataNode_CheckChannelOperationProgress_Call {
 	return &MockDataNode_CheckChannelOperationProgress_Call{Call: _e.mock.On("CheckChannelOperationProgress", _a0, _a1)}
 }
@ -174,8 +174,8 @@ type MockDataNode_Compaction_Call struct {
 }

 // Compaction is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.CompactionPlan
+//  - _a0 context.Context
+//  - _a1 *datapb.CompactionPlan
 func (_e *MockDataNode_Expecter) Compaction(_a0 interface{}, _a1 interface{}) *MockDataNode_Compaction_Call {
 	return &MockDataNode_Compaction_Call{Call: _e.mock.On("Compaction", _a0, _a1)}
 }
@ -197,6 +197,61 @@ func (_c *MockDataNode_Compaction_Call) RunAndReturn(run func(context.Context, *
 	return _c
 }

+// DropImport provides a mock function with given fields: _a0, _a1
+func (_m *MockDataNode) DropImport(_a0 context.Context, _a1 *datapb.DropImportRequest) (*commonpb.Status, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.DropImportRequest) (*commonpb.Status, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.DropImportRequest) *commonpb.Status); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.DropImportRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNode_DropImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropImport'
+type MockDataNode_DropImport_Call struct {
+	*mock.Call
+}
+
+// DropImport is a helper method to define mock.On call
+//  - _a0 context.Context
+//  - _a1 *datapb.DropImportRequest
+func (_e *MockDataNode_Expecter) DropImport(_a0 interface{}, _a1 interface{}) *MockDataNode_DropImport_Call {
+	return &MockDataNode_DropImport_Call{Call: _e.mock.On("DropImport", _a0, _a1)}
+}
+
+func (_c *MockDataNode_DropImport_Call) Run(run func(_a0 context.Context, _a1 *datapb.DropImportRequest)) *MockDataNode_DropImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*datapb.DropImportRequest))
+	})
+	return _c
+}
+
+func (_c *MockDataNode_DropImport_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNode_DropImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNode_DropImport_Call) RunAndReturn(run func(context.Context, *datapb.DropImportRequest) (*commonpb.Status, error)) *MockDataNode_DropImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // FlushChannels provides a mock function with given fields: _a0, _a1
 func (_m *MockDataNode) FlushChannels(_a0 context.Context, _a1 *datapb.FlushChannelsRequest) (*commonpb.Status, error) {
 	ret := _m.Called(_a0, _a1)
@ -229,8 +284,8 @@ type MockDataNode_FlushChannels_Call struct {
 }

 // FlushChannels is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.FlushChannelsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.FlushChannelsRequest
 func (_e *MockDataNode_Expecter) FlushChannels(_a0 interface{}, _a1 interface{}) *MockDataNode_FlushChannels_Call {
 	return &MockDataNode_FlushChannels_Call{Call: _e.mock.On("FlushChannels", _a0, _a1)}
 }
@ -284,8 +339,8 @@ type MockDataNode_FlushSegments_Call struct {
 }

 // FlushSegments is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.FlushSegmentsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.FlushSegmentsRequest
 func (_e *MockDataNode_Expecter) FlushSegments(_a0 interface{}, _a1 interface{}) *MockDataNode_FlushSegments_Call {
 	return &MockDataNode_FlushSegments_Call{Call: _e.mock.On("FlushSegments", _a0, _a1)}
 }
@ -380,8 +435,8 @@ type MockDataNode_GetCompactionState_Call struct {
 }

 // GetCompactionState is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.CompactionStateRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.CompactionStateRequest
 func (_e *MockDataNode_Expecter) GetCompactionState(_a0 interface{}, _a1 interface{}) *MockDataNode_GetCompactionState_Call {
 	return &MockDataNode_GetCompactionState_Call{Call: _e.mock.On("GetCompactionState", _a0, _a1)}
 }
@ -435,8 +490,8 @@ type MockDataNode_GetComponentStates_Call struct {
 }

 // GetComponentStates is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *milvuspb.GetComponentStatesRequest
+//  - _a0 context.Context
+//  - _a1 *milvuspb.GetComponentStatesRequest
 func (_e *MockDataNode_Expecter) GetComponentStates(_a0 interface{}, _a1 interface{}) *MockDataNode_GetComponentStates_Call {
 	return &MockDataNode_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates", _a0, _a1)}
 }
@ -490,8 +545,8 @@ type MockDataNode_GetMetrics_Call struct {
 }

 // GetMetrics is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *milvuspb.GetMetricsRequest
+//  - _a0 context.Context
+//  - _a1 *milvuspb.GetMetricsRequest
 func (_e *MockDataNode_Expecter) GetMetrics(_a0 interface{}, _a1 interface{}) *MockDataNode_GetMetrics_Call {
 	return &MockDataNode_GetMetrics_Call{Call: _e.mock.On("GetMetrics", _a0, _a1)}
 }
@ -586,8 +641,8 @@ type MockDataNode_GetStatisticsChannel_Call struct {
 }

 // GetStatisticsChannel is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *internalpb.GetStatisticsChannelRequest
+//  - _a0 context.Context
+//  - _a1 *internalpb.GetStatisticsChannelRequest
 func (_e *MockDataNode_Expecter) GetStatisticsChannel(_a0 interface{}, _a1 interface{}) *MockDataNode_GetStatisticsChannel_Call {
 	return &MockDataNode_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel", _a0, _a1)}
 }
@ -641,8 +696,8 @@ type MockDataNode_Import_Call struct {
 }

 // Import is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.ImportTaskRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.ImportTaskRequest
 func (_e *MockDataNode_Expecter) Import(_a0 interface{}, _a1 interface{}) *MockDataNode_Import_Call {
 	return &MockDataNode_Import_Call{Call: _e.mock.On("Import", _a0, _a1)}
 }
@ -664,6 +719,61 @@ func (_c *MockDataNode_Import_Call) RunAndReturn(run func(context.Context, *data
 	return _c
 }

+// ImportV2 provides a mock function with given fields: _a0, _a1
+func (_m *MockDataNode) ImportV2(_a0 context.Context, _a1 *datapb.ImportRequest) (*commonpb.Status, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.ImportRequest) (*commonpb.Status, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.ImportRequest) *commonpb.Status); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.ImportRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNode_ImportV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImportV2'
+type MockDataNode_ImportV2_Call struct {
+	*mock.Call
+}
+
+// ImportV2 is a helper method to define mock.On call
+//  - _a0 context.Context
+//  - _a1 *datapb.ImportRequest
+func (_e *MockDataNode_Expecter) ImportV2(_a0 interface{}, _a1 interface{}) *MockDataNode_ImportV2_Call {
+	return &MockDataNode_ImportV2_Call{Call: _e.mock.On("ImportV2", _a0, _a1)}
+}
+
+func (_c *MockDataNode_ImportV2_Call) Run(run func(_a0 context.Context, _a1 *datapb.ImportRequest)) *MockDataNode_ImportV2_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*datapb.ImportRequest))
+	})
+	return _c
+}
+
+func (_c *MockDataNode_ImportV2_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNode_ImportV2_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNode_ImportV2_Call) RunAndReturn(run func(context.Context, *datapb.ImportRequest) (*commonpb.Status, error)) *MockDataNode_ImportV2_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Init provides a mock function with given fields:
 func (_m *MockDataNode) Init() error {
 	ret := _m.Called()
@ -737,8 +847,8 @@ type MockDataNode_NotifyChannelOperation_Call struct {
 }

 // NotifyChannelOperation is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.ChannelOperationsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.ChannelOperationsRequest
 func (_e *MockDataNode_Expecter) NotifyChannelOperation(_a0 interface{}, _a1 interface{}) *MockDataNode_NotifyChannelOperation_Call {
 	return &MockDataNode_NotifyChannelOperation_Call{Call: _e.mock.On("NotifyChannelOperation", _a0, _a1)}
 }
@ -760,6 +870,171 @@ func (_c *MockDataNode_NotifyChannelOperation_Call) RunAndReturn(run func(contex
 	return _c
 }

+// PreImport provides a mock function with given fields: _a0, _a1
+func (_m *MockDataNode) PreImport(_a0 context.Context, _a1 *datapb.PreImportRequest) (*commonpb.Status, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.PreImportRequest) (*commonpb.Status, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.PreImportRequest) *commonpb.Status); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.PreImportRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNode_PreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PreImport'
+type MockDataNode_PreImport_Call struct {
+	*mock.Call
+}
+
+// PreImport is a helper method to define mock.On call
+//  - _a0 context.Context
+//  - _a1 *datapb.PreImportRequest
+func (_e *MockDataNode_Expecter) PreImport(_a0 interface{}, _a1 interface{}) *MockDataNode_PreImport_Call {
+	return &MockDataNode_PreImport_Call{Call: _e.mock.On("PreImport", _a0, _a1)}
+}
+
+func (_c *MockDataNode_PreImport_Call) Run(run func(_a0 context.Context, _a1 *datapb.PreImportRequest)) *MockDataNode_PreImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*datapb.PreImportRequest))
+	})
+	return _c
+}
+
+func (_c *MockDataNode_PreImport_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNode_PreImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNode_PreImport_Call) RunAndReturn(run func(context.Context, *datapb.PreImportRequest) (*commonpb.Status, error)) *MockDataNode_PreImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// QueryImport provides a mock function with given fields: _a0, _a1
+func (_m *MockDataNode) QueryImport(_a0 context.Context, _a1 *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *datapb.QueryImportResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryImportRequest) *datapb.QueryImportResponse); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*datapb.QueryImportResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.QueryImportRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNode_QueryImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryImport'
+type MockDataNode_QueryImport_Call struct {
+	*mock.Call
+}
+
+// QueryImport is a helper method to define mock.On call
+//  - _a0 context.Context
+//  - _a1 *datapb.QueryImportRequest
+func (_e *MockDataNode_Expecter) QueryImport(_a0 interface{}, _a1 interface{}) *MockDataNode_QueryImport_Call {
+	return &MockDataNode_QueryImport_Call{Call: _e.mock.On("QueryImport", _a0, _a1)}
+}
+
+func (_c *MockDataNode_QueryImport_Call) Run(run func(_a0 context.Context, _a1 *datapb.QueryImportRequest)) *MockDataNode_QueryImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*datapb.QueryImportRequest))
+	})
+	return _c
+}
+
+func (_c *MockDataNode_QueryImport_Call) Return(_a0 *datapb.QueryImportResponse, _a1 error) *MockDataNode_QueryImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNode_QueryImport_Call) RunAndReturn(run func(context.Context, *datapb.QueryImportRequest) (*datapb.QueryImportResponse, error)) *MockDataNode_QueryImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// QueryPreImport provides a mock function with given fields: _a0, _a1
+func (_m *MockDataNode) QueryPreImport(_a0 context.Context, _a1 *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *datapb.QueryPreImportResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryPreImportRequest) *datapb.QueryPreImportResponse); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*datapb.QueryPreImportResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.QueryPreImportRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNode_QueryPreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryPreImport'
+type MockDataNode_QueryPreImport_Call struct {
+	*mock.Call
+}
+
+// QueryPreImport is a helper method to define mock.On call
+//  - _a0 context.Context
+//  - _a1 *datapb.QueryPreImportRequest
+func (_e *MockDataNode_Expecter) QueryPreImport(_a0 interface{}, _a1 interface{}) *MockDataNode_QueryPreImport_Call {
+	return &MockDataNode_QueryPreImport_Call{Call: _e.mock.On("QueryPreImport", _a0, _a1)}
+}
+
+func (_c *MockDataNode_QueryPreImport_Call) Run(run func(_a0 context.Context, _a1 *datapb.QueryPreImportRequest)) *MockDataNode_QueryPreImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*datapb.QueryPreImportRequest))
+	})
+	return _c
+}
+
+func (_c *MockDataNode_QueryPreImport_Call) Return(_a0 *datapb.QueryPreImportResponse, _a1 error) *MockDataNode_QueryPreImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNode_QueryPreImport_Call) RunAndReturn(run func(context.Context, *datapb.QueryPreImportRequest) (*datapb.QueryPreImportResponse, error)) *MockDataNode_QueryPreImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Register provides a mock function with given fields:
 func (_m *MockDataNode) Register() error {
 	ret := _m.Called()
@ -833,8 +1108,8 @@ type MockDataNode_ResendSegmentStats_Call struct {
 }

 // ResendSegmentStats is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.ResendSegmentStatsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.ResendSegmentStatsRequest
 func (_e *MockDataNode_Expecter) ResendSegmentStats(_a0 interface{}, _a1 interface{}) *MockDataNode_ResendSegmentStats_Call {
 	return &MockDataNode_ResendSegmentStats_Call{Call: _e.mock.On("ResendSegmentStats", _a0, _a1)}
 }
@ -867,7 +1142,7 @@ type MockDataNode_SetAddress_Call struct {
 }

 // SetAddress is a helper method to define mock.On call
-//   - address string
+//  - address string
 func (_e *MockDataNode_Expecter) SetAddress(address interface{}) *MockDataNode_SetAddress_Call {
 	return &MockDataNode_SetAddress_Call{Call: _e.mock.On("SetAddress", address)}
 }
@ -909,7 +1184,7 @@ type MockDataNode_SetDataCoordClient_Call struct {
 }

 // SetDataCoordClient is a helper method to define mock.On call
-//   - dataCoord types.DataCoordClient
+//  - dataCoord types.DataCoordClient
 func (_e *MockDataNode_Expecter) SetDataCoordClient(dataCoord interface{}) *MockDataNode_SetDataCoordClient_Call {
 	return &MockDataNode_SetDataCoordClient_Call{Call: _e.mock.On("SetDataCoordClient", dataCoord)}
 }
@ -942,7 +1217,7 @@ type MockDataNode_SetEtcdClient_Call struct {
 }

 // SetEtcdClient is a helper method to define mock.On call
-//   - etcdClient *clientv3.Client
+//  - etcdClient *clientv3.Client
 func (_e *MockDataNode_Expecter) SetEtcdClient(etcdClient interface{}) *MockDataNode_SetEtcdClient_Call {
 	return &MockDataNode_SetEtcdClient_Call{Call: _e.mock.On("SetEtcdClient", etcdClient)}
 }
@ -984,7 +1259,7 @@ type MockDataNode_SetRootCoordClient_Call struct {
 }

 // SetRootCoordClient is a helper method to define mock.On call
-//   - rootCoord types.RootCoordClient
+//  - rootCoord types.RootCoordClient
 func (_e *MockDataNode_Expecter) SetRootCoordClient(rootCoord interface{}) *MockDataNode_SetRootCoordClient_Call {
 	return &MockDataNode_SetRootCoordClient_Call{Call: _e.mock.On("SetRootCoordClient", rootCoord)}
 }
@ -1038,8 +1313,8 @@ type MockDataNode_ShowConfigurations_Call struct {
 }

 // ShowConfigurations is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *internalpb.ShowConfigurationsRequest
+//  - _a0 context.Context
+//  - _a1 *internalpb.ShowConfigurationsRequest
 func (_e *MockDataNode_Expecter) ShowConfigurations(_a0 interface{}, _a1 interface{}) *MockDataNode_ShowConfigurations_Call {
 	return &MockDataNode_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations", _a0, _a1)}
 }
@ -1175,8 +1450,8 @@ type MockDataNode_SyncSegments_Call struct {
 }

 // SyncSegments is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.SyncSegmentsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.SyncSegmentsRequest
 func (_e *MockDataNode_Expecter) SyncSegments(_a0 interface{}, _a1 interface{}) *MockDataNode_SyncSegments_Call {
 	return &MockDataNode_SyncSegments_Call{Call: _e.mock.On("SyncSegments", _a0, _a1)}
 }
@ -1209,7 +1484,7 @@ type MockDataNode_UpdateStateCode_Call struct {
 }

 // UpdateStateCode is a helper method to define mock.On call
-//   - stateCode commonpb.StateCode
+//  - stateCode commonpb.StateCode
 func (_e *MockDataNode_Expecter) UpdateStateCode(stateCode interface{}) *MockDataNode_UpdateStateCode_Call {
 	return &MockDataNode_UpdateStateCode_Call{Call: _e.mock.On("UpdateStateCode", stateCode)}
 }
@ -1263,8 +1538,8 @@ type MockDataNode_WatchDmChannels_Call struct {
 }

 // WatchDmChannels is a helper method to define mock.On call
-//   - _a0 context.Context
-//   - _a1 *datapb.WatchDmChannelsRequest
+//  - _a0 context.Context
+//  - _a1 *datapb.WatchDmChannelsRequest
 func (_e *MockDataNode_Expecter) WatchDmChannels(_a0 interface{}, _a1 interface{}) *MockDataNode_WatchDmChannels_Call {
 	return &MockDataNode_WatchDmChannels_Call{Call: _e.mock.On("WatchDmChannels", _a0, _a1)}
 }
--- a/internal/mocks/mock_datanode_client.go
+++ b/internal/mocks/mock_datanode_client.go
@ -70,9 +70,9 @@ type MockDataNodeClient_AddImportSegment_Call struct {
 }

 // AddImportSegment is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.AddImportSegmentRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.AddImportSegmentRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) AddImportSegment(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_AddImportSegment_Call {
 	return &MockDataNodeClient_AddImportSegment_Call{Call: _e.mock.On("AddImportSegment",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -140,9 +140,9 @@ type MockDataNodeClient_CheckChannelOperationProgress_Call struct {
 }

 // CheckChannelOperationProgress is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.ChannelWatchInfo
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.ChannelWatchInfo
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) CheckChannelOperationProgress(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_CheckChannelOperationProgress_Call {
 	return &MockDataNodeClient_CheckChannelOperationProgress_Call{Call: _e.mock.On("CheckChannelOperationProgress",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -251,9 +251,9 @@ type MockDataNodeClient_Compaction_Call struct {
 }

 // Compaction is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.CompactionPlan
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.CompactionPlan
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) Compaction(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_Compaction_Call {
 	return &MockDataNodeClient_Compaction_Call{Call: _e.mock.On("Compaction",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -282,6 +282,76 @@ func (_c *MockDataNodeClient_Compaction_Call) RunAndReturn(run func(context.Cont
 	return _c
 }

+// DropImport provides a mock function with given fields: ctx, in, opts
+func (_m *MockDataNodeClient) DropImport(ctx context.Context, in *datapb.DropImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.DropImportRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.DropImportRequest, ...grpc.CallOption) *commonpb.Status); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.DropImportRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNodeClient_DropImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropImport'
+type MockDataNodeClient_DropImport_Call struct {
+	*mock.Call
+}
+
+// DropImport is a helper method to define mock.On call
+//  - ctx context.Context
+//  - in *datapb.DropImportRequest
+//  - opts ...grpc.CallOption
+func (_e *MockDataNodeClient_Expecter) DropImport(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_DropImport_Call {
+	return &MockDataNodeClient_DropImport_Call{Call: _e.mock.On("DropImport",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockDataNodeClient_DropImport_Call) Run(run func(ctx context.Context, in *datapb.DropImportRequest, opts ...grpc.CallOption)) *MockDataNodeClient_DropImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*datapb.DropImportRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockDataNodeClient_DropImport_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNodeClient_DropImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNodeClient_DropImport_Call) RunAndReturn(run func(context.Context, *datapb.DropImportRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockDataNodeClient_DropImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // FlushChannels provides a mock function with given fields: ctx, in, opts
 func (_m *MockDataNodeClient) FlushChannels(ctx context.Context, in *datapb.FlushChannelsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
 	_va := make([]interface{}, len(opts))
@ -321,9 +391,9 @@ type MockDataNodeClient_FlushChannels_Call struct {
 }

 // FlushChannels is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.FlushChannelsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.FlushChannelsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) FlushChannels(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_FlushChannels_Call {
 	return &MockDataNodeClient_FlushChannels_Call{Call: _e.mock.On("FlushChannels",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -391,9 +461,9 @@ type MockDataNodeClient_FlushSegments_Call struct {
 }

 // FlushSegments is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.FlushSegmentsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.FlushSegmentsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) FlushSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_FlushSegments_Call {
 	return &MockDataNodeClient_FlushSegments_Call{Call: _e.mock.On("FlushSegments",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -461,9 +531,9 @@ type MockDataNodeClient_GetCompactionState_Call struct {
 }

 // GetCompactionState is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.CompactionStateRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.CompactionStateRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) GetCompactionState(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_GetCompactionState_Call {
 	return &MockDataNodeClient_GetCompactionState_Call{Call: _e.mock.On("GetCompactionState",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -531,9 +601,9 @@ type MockDataNodeClient_GetComponentStates_Call struct {
 }

 // GetComponentStates is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *milvuspb.GetComponentStatesRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *milvuspb.GetComponentStatesRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) GetComponentStates(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_GetComponentStates_Call {
 	return &MockDataNodeClient_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -601,9 +671,9 @@ type MockDataNodeClient_GetMetrics_Call struct {
 }

 // GetMetrics is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *milvuspb.GetMetricsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *milvuspb.GetMetricsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) GetMetrics(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_GetMetrics_Call {
 	return &MockDataNodeClient_GetMetrics_Call{Call: _e.mock.On("GetMetrics",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -671,9 +741,9 @@ type MockDataNodeClient_GetStatisticsChannel_Call struct {
 }

 // GetStatisticsChannel is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *internalpb.GetStatisticsChannelRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *internalpb.GetStatisticsChannelRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) GetStatisticsChannel(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_GetStatisticsChannel_Call {
 	return &MockDataNodeClient_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -741,9 +811,9 @@ type MockDataNodeClient_Import_Call struct {
 }

 // Import is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.ImportTaskRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.ImportTaskRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) Import(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_Import_Call {
 	return &MockDataNodeClient_Import_Call{Call: _e.mock.On("Import",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -772,6 +842,76 @@ func (_c *MockDataNodeClient_Import_Call) RunAndReturn(run func(context.Context,
 	return _c
 }

+// ImportV2 provides a mock function with given fields: ctx, in, opts
+func (_m *MockDataNodeClient) ImportV2(ctx context.Context, in *datapb.ImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.ImportRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.ImportRequest, ...grpc.CallOption) *commonpb.Status); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.ImportRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNodeClient_ImportV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImportV2'
+type MockDataNodeClient_ImportV2_Call struct {
+	*mock.Call
+}
+
+// ImportV2 is a helper method to define mock.On call
+//  - ctx context.Context
+//  - in *datapb.ImportRequest
+//  - opts ...grpc.CallOption
+func (_e *MockDataNodeClient_Expecter) ImportV2(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_ImportV2_Call {
+	return &MockDataNodeClient_ImportV2_Call{Call: _e.mock.On("ImportV2",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockDataNodeClient_ImportV2_Call) Run(run func(ctx context.Context, in *datapb.ImportRequest, opts ...grpc.CallOption)) *MockDataNodeClient_ImportV2_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*datapb.ImportRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockDataNodeClient_ImportV2_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNodeClient_ImportV2_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNodeClient_ImportV2_Call) RunAndReturn(run func(context.Context, *datapb.ImportRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockDataNodeClient_ImportV2_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // NotifyChannelOperation provides a mock function with given fields: ctx, in, opts
 func (_m *MockDataNodeClient) NotifyChannelOperation(ctx context.Context, in *datapb.ChannelOperationsRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
 	_va := make([]interface{}, len(opts))
@ -811,9 +951,9 @@ type MockDataNodeClient_NotifyChannelOperation_Call struct {
 }

 // NotifyChannelOperation is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.ChannelOperationsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.ChannelOperationsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) NotifyChannelOperation(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_NotifyChannelOperation_Call {
 	return &MockDataNodeClient_NotifyChannelOperation_Call{Call: _e.mock.On("NotifyChannelOperation",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -842,6 +982,216 @@ func (_c *MockDataNodeClient_NotifyChannelOperation_Call) RunAndReturn(run func(
 	return _c
 }

+// PreImport provides a mock function with given fields: ctx, in, opts
+func (_m *MockDataNodeClient) PreImport(ctx context.Context, in *datapb.PreImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.PreImportRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.PreImportRequest, ...grpc.CallOption) *commonpb.Status); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.PreImportRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNodeClient_PreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PreImport'
+type MockDataNodeClient_PreImport_Call struct {
+	*mock.Call
+}
+
+// PreImport is a helper method to define mock.On call
+//  - ctx context.Context
+//  - in *datapb.PreImportRequest
+//  - opts ...grpc.CallOption
+func (_e *MockDataNodeClient_Expecter) PreImport(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_PreImport_Call {
+	return &MockDataNodeClient_PreImport_Call{Call: _e.mock.On("PreImport",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockDataNodeClient_PreImport_Call) Run(run func(ctx context.Context, in *datapb.PreImportRequest, opts ...grpc.CallOption)) *MockDataNodeClient_PreImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*datapb.PreImportRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockDataNodeClient_PreImport_Call) Return(_a0 *commonpb.Status, _a1 error) *MockDataNodeClient_PreImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNodeClient_PreImport_Call) RunAndReturn(run func(context.Context, *datapb.PreImportRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockDataNodeClient_PreImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// QueryImport provides a mock function with given fields: ctx, in, opts
+func (_m *MockDataNodeClient) QueryImport(ctx context.Context, in *datapb.QueryImportRequest, opts ...grpc.CallOption) (*datapb.QueryImportResponse, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *datapb.QueryImportResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryImportRequest, ...grpc.CallOption) (*datapb.QueryImportResponse, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryImportRequest, ...grpc.CallOption) *datapb.QueryImportResponse); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*datapb.QueryImportResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.QueryImportRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNodeClient_QueryImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryImport'
+type MockDataNodeClient_QueryImport_Call struct {
+	*mock.Call
+}
+
+// QueryImport is a helper method to define mock.On call
+//  - ctx context.Context
+//  - in *datapb.QueryImportRequest
+//  - opts ...grpc.CallOption
+func (_e *MockDataNodeClient_Expecter) QueryImport(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_QueryImport_Call {
+	return &MockDataNodeClient_QueryImport_Call{Call: _e.mock.On("QueryImport",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockDataNodeClient_QueryImport_Call) Run(run func(ctx context.Context, in *datapb.QueryImportRequest, opts ...grpc.CallOption)) *MockDataNodeClient_QueryImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*datapb.QueryImportRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockDataNodeClient_QueryImport_Call) Return(_a0 *datapb.QueryImportResponse, _a1 error) *MockDataNodeClient_QueryImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNodeClient_QueryImport_Call) RunAndReturn(run func(context.Context, *datapb.QueryImportRequest, ...grpc.CallOption) (*datapb.QueryImportResponse, error)) *MockDataNodeClient_QueryImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// QueryPreImport provides a mock function with given fields: ctx, in, opts
+func (_m *MockDataNodeClient) QueryPreImport(ctx context.Context, in *datapb.QueryPreImportRequest, opts ...grpc.CallOption) (*datapb.QueryPreImportResponse, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *datapb.QueryPreImportResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryPreImportRequest, ...grpc.CallOption) (*datapb.QueryPreImportResponse, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *datapb.QueryPreImportRequest, ...grpc.CallOption) *datapb.QueryPreImportResponse); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*datapb.QueryPreImportResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *datapb.QueryPreImportRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockDataNodeClient_QueryPreImport_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'QueryPreImport'
+type MockDataNodeClient_QueryPreImport_Call struct {
+	*mock.Call
+}
+
+// QueryPreImport is a helper method to define mock.On call
+//  - ctx context.Context
+//  - in *datapb.QueryPreImportRequest
+//  - opts ...grpc.CallOption
+func (_e *MockDataNodeClient_Expecter) QueryPreImport(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_QueryPreImport_Call {
+	return &MockDataNodeClient_QueryPreImport_Call{Call: _e.mock.On("QueryPreImport",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockDataNodeClient_QueryPreImport_Call) Run(run func(ctx context.Context, in *datapb.QueryPreImportRequest, opts ...grpc.CallOption)) *MockDataNodeClient_QueryPreImport_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*datapb.QueryPreImportRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockDataNodeClient_QueryPreImport_Call) Return(_a0 *datapb.QueryPreImportResponse, _a1 error) *MockDataNodeClient_QueryPreImport_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockDataNodeClient_QueryPreImport_Call) RunAndReturn(run func(context.Context, *datapb.QueryPreImportRequest, ...grpc.CallOption) (*datapb.QueryPreImportResponse, error)) *MockDataNodeClient_QueryPreImport_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // ResendSegmentStats provides a mock function with given fields: ctx, in, opts
 func (_m *MockDataNodeClient) ResendSegmentStats(ctx context.Context, in *datapb.ResendSegmentStatsRequest, opts ...grpc.CallOption) (*datapb.ResendSegmentStatsResponse, error) {
 	_va := make([]interface{}, len(opts))
@ -881,9 +1231,9 @@ type MockDataNodeClient_ResendSegmentStats_Call struct {
 }

 // ResendSegmentStats is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.ResendSegmentStatsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.ResendSegmentStatsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) ResendSegmentStats(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_ResendSegmentStats_Call {
 	return &MockDataNodeClient_ResendSegmentStats_Call{Call: _e.mock.On("ResendSegmentStats",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -951,9 +1301,9 @@ type MockDataNodeClient_ShowConfigurations_Call struct {
 }

 // ShowConfigurations is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *internalpb.ShowConfigurationsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *internalpb.ShowConfigurationsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) ShowConfigurations(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_ShowConfigurations_Call {
 	return &MockDataNodeClient_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -1021,9 +1371,9 @@ type MockDataNodeClient_SyncSegments_Call struct {
 }

 // SyncSegments is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.SyncSegmentsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.SyncSegmentsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) SyncSegments(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_SyncSegments_Call {
 	return &MockDataNodeClient_SyncSegments_Call{Call: _e.mock.On("SyncSegments",
 		append([]interface{}{ctx, in}, opts...)...)}
@ -1091,9 +1441,9 @@ type MockDataNodeClient_WatchDmChannels_Call struct {
 }

 // WatchDmChannels is a helper method to define mock.On call
-//   - ctx context.Context
-//   - in *datapb.WatchDmChannelsRequest
-//   - opts ...grpc.CallOption
+//  - ctx context.Context
+//  - in *datapb.WatchDmChannelsRequest
+//  - opts ...grpc.CallOption
 func (_e *MockDataNodeClient_Expecter) WatchDmChannels(ctx interface{}, in interface{}, opts ...interface{}) *MockDataNodeClient_WatchDmChannels_Call {
 	return &MockDataNodeClient_WatchDmChannels_Call{Call: _e.mock.On("WatchDmChannels",
 		append([]interface{}{ctx, in}, opts...)...)}
--- a/internal/proto/data_coord.proto
+++ b/internal/proto/data_coord.proto
@ -119,6 +119,13 @@ service DataNode {
  rpc FlushChannels(FlushChannelsRequest) returns(common.Status) {}
  rpc NotifyChannelOperation(ChannelOperationsRequest) returns(common.Status) {}
  rpc CheckChannelOperationProgress(ChannelWatchInfo) returns(ChannelOperationProgressResponse) {}
+
+  // import v2
+  rpc PreImport(PreImportRequest) returns(common.Status) {}
+  rpc ImportV2(ImportRequest) returns(common.Status) {}
+  rpc QueryPreImport(QueryPreImportRequest) returns(QueryPreImportResponse) {}
+  rpc QueryImport(QueryImportRequest) returns(QueryImportResponse) {}
+  rpc DropImport(DropImportRequest) returns(common.Status) {}
 }

 message FlushRequest {
@ -737,3 +744,124 @@ message ChannelOperationProgressResponse {
  ChannelWatchState state = 3;
  int32 progress = 4;
 }
+
+enum ImportState {
+  None = 0;
+  Pending = 1;
+  InProgress = 2;
+  Failed = 3;
+  Completed = 4;
+}
+
+message ColumnBasedFile {
+  repeated string files = 1;
+}
+
+message ImportFile {
+  oneof file {
+    string row_based_file = 1;
+    ColumnBasedFile column_based_file = 2;
+  }
+}
+
+message PreImportRequest {
+  string clusterID = 1;
+  int64 requestID = 2;
+  int64 taskID = 3;
+  int64 collectionID = 4;
+  int64 partitionID = 5;
+  schema.CollectionSchema schema = 6;
+  repeated ImportFile import_files = 7;
+}
+
+message autoIDRange {
+  int64 begin = 1;
+  int64 end = 2;
+}
+
+message ImportRequest {
+  string clusterID = 1;
+  int64 requestID = 2;
+  int64 taskID = 3;
+  int64 collectionID = 4;
+  int64 partitionID = 5;
+  schema.CollectionSchema schema = 6;
+  repeated ImportFile import_files = 7;
+  map<int64, autoIDRange> autoID_ranges = 8;
+  map<int64, string> segment_channels = 9;
+}
+
+message QueryPreImportRequest {
+  string clusterID = 1;
+  int64 requestID = 2;
+  int64 taskID = 3;
+}
+
+message ImportFileStats {
+  ImportFile import_file = 1;
+  int64 file_size = 2;
+  int64 total_rows = 3;
+  map<string, int64> channel_rows = 4;
+}
+
+message QueryPreImportResponse {
+  common.Status status = 1;
+  int64 taskID = 2;
+  ImportState state = 3;
+  string reason = 4;
+  int64 slots = 5;
+  repeated ImportFileStats file_stats = 6;
+}
+
+message QueryImportRequest {
+  string clusterID = 1;
+  int64 requestID = 2;
+  int64 taskID = 3;
+}
+
+message ImportSegmentInfo {
+  int64 segmentID = 1;
+  int64 total_rows = 2;
+  int64 imported_rows = 3;
+  repeated FieldBinlog binlogs = 4;
+  repeated FieldBinlog statslogs = 5;
+  repeated index.IndexFilePathInfo index_infos = 6;
+}
+
+message QueryImportResponse {
+  common.Status status = 1;
+  int64 taskID = 2;
+  ImportState state = 3;
+  string reason = 4;
+  int64 slots = 5;
+  repeated ImportSegmentInfo import_segments_info = 6;
+}
+
+message DropImportRequest {
+  string clusterID = 1;
+  int64 requestID = 2;
+  int64 taskID = 3;
+}
+
+message PreImportTask {
+  int64 requestID = 1;
+  int64 taskID = 2;
+  int64 collectionID = 3;
+  int64 partitionID = 4;
+  int64 nodeID = 5;
+  ImportState state = 6;
+  string reason = 7;
+  repeated ImportFileStats file_stats = 8;
+}
+
+message ImportTaskV2 {
+  int64 requestID = 1;
+  int64 taskID = 2;
+  int64 collectionID = 3;
+  int64 partitionID = 4;
+  repeated int64 segmentIDs = 5;
+  int64 nodeID = 6;
+  ImportState state = 7;
+  string reason = 8;
+  repeated ImportFile files = 9;
+}
--- a/internal/proto/datapb/data_coord.pb.go
+++ b/internal/proto/datapb/data_coord.pb.go
--- a/internal/util/mock/grpc_datanode_client.go
+++ b/internal/util/mock/grpc_datanode_client.go
@ -92,3 +92,23 @@ func (m *GrpcDataNodeClient) NotifyChannelOperation(ctx context.Context, in *dat
 func (m *GrpcDataNodeClient) CheckChannelOperationProgress(ctx context.Context, req *datapb.ChannelWatchInfo, opts ...grpc.CallOption) (*datapb.ChannelOperationProgressResponse, error) {
 	return &datapb.ChannelOperationProgressResponse{}, m.Err
 }
+
+func (m *GrpcDataNodeClient) PreImport(ctx context.Context, req *datapb.PreImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{}, m.Err
+}
+
+func (m *GrpcDataNodeClient) ImportV2(ctx context.Context, req *datapb.ImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{}, m.Err
+}
+
+func (m *GrpcDataNodeClient) QueryPreImport(ctx context.Context, req *datapb.QueryPreImportRequest, opts ...grpc.CallOption) (*datapb.QueryPreImportResponse, error) {
+	return &datapb.QueryPreImportResponse{}, m.Err
+}
+
+func (m *GrpcDataNodeClient) QueryImport(ctx context.Context, req *datapb.QueryImportRequest, opts ...grpc.CallOption) (*datapb.QueryImportResponse, error) {
+	return &datapb.QueryImportResponse{}, m.Err
+}
+
+func (m *GrpcDataNodeClient) DropImport(ctx context.Context, req *datapb.DropImportRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
+	return &commonpb.Status{}, m.Err
+}