From b98b3b16a3885c1616bbbbce79b0900452771dcb Mon Sep 17 00:00:00 2001 From: Xianhui Lin <35839735+JsDove@users.noreply.github.com> Date: Mon, 18 Aug 2025 01:23:45 +0800 Subject: [PATCH] feat:add BatchDescribeCollection interface (#43786) feat:add BatchDescribeCollection interface issue: https://github.com/milvus-io/milvus/issues/43781 Signed-off-by: Xianhui.Lin --- client/go.mod | 6 +- client/go.sum | 6 +- .../milvusclient/mock_milvus_server_test.go | 236 ++++++++++++++++++ go.mod | 2 +- go.sum | 4 +- internal/distributed/proxy/service.go | 5 + internal/distributed/proxy/service_test.go | 6 + internal/mocks/mock_proxy.go | 59 +++++ internal/proxy/database_interceptor.go | 5 + internal/proxy/database_interceptor_test.go | 1 + internal/proxy/impl.go | 43 +++- internal/proxy/impl_test.go | 12 + internal/proxy/task.go | 2 + internal/rootcoord/show_collection_task.go | 1 + .../rootcoord/show_collection_task_test.go | 10 +- pkg/go.mod | 2 +- pkg/go.sum | 4 +- tests/go_client/go.mod | 2 +- tests/go_client/go.sum | 6 +- 19 files changed, 390 insertions(+), 22 deletions(-) diff --git a/client/go.mod b/client/go.mod index 0e0c735a25..27a259d1ba 100644 --- a/client/go.mod +++ b/client/go.mod @@ -6,15 +6,13 @@ require ( github.com/blang/semver/v4 v4.0.0 github.com/cockroachdb/errors v1.9.1 github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 - github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0 + github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e github.com/quasilyte/go-ruleguard/dsl v0.3.22 github.com/samber/lo v1.27.0 github.com/stretchr/testify v1.10.0 github.com/tidwall/gjson v1.17.1 - go.opentelemetry.io/otel v1.28.0 go.uber.org/atomic v1.11.0 - golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 google.golang.org/grpc v1.65.0 google.golang.org/protobuf v1.34.2 ) @@ -89,6 +87,7 @@ require ( go.etcd.io/etcd/raft/v3 v3.5.5 // indirect go.etcd.io/etcd/server/v3 v3.5.5 // indirect go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect + go.opentelemetry.io/otel v1.28.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.20.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.20.0 // indirect go.opentelemetry.io/otel/metric v1.28.0 // indirect @@ -99,6 +98,7 @@ require ( go.uber.org/multierr v1.11.0 // indirect go.uber.org/zap v1.27.0 // indirect golang.org/x/crypto v0.36.0 // indirect + golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect golang.org/x/net v0.38.0 // indirect golang.org/x/sync v0.12.0 // indirect golang.org/x/sys v0.31.0 // indirect diff --git a/client/go.sum b/client/go.sum index a2af448136..4609a4ed92 100644 --- a/client/go.sum +++ b/client/go.sum @@ -318,8 +318,10 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8= github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc= github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0 h1:SJLdKkifvPDT31jw0hzC7/KSlZ5tv5AhPt77jNHMAQY= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807040333-531631e7fce6 h1:qTBOTsZ3OwEXkrHRqPn562ddkDqeToIY6CstLIaVQYs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807040333-531631e7fce6/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 h1:zyrKuc0rwT5xWIFkZr/bFWXXYbYvSBMT3iFITnaR8IE= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e h1:VCr43pG4efacDbM4au70fh8/5hNTftoWzm1iEumvDWM= github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e/go.mod h1:37AWzxVs2NS4QUJrkcbeLUwi+4Av0h5mEdjLI62EANU= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc= diff --git a/client/milvusclient/mock_milvus_server_test.go b/client/milvusclient/mock_milvus_server_test.go index ad36e1f090..e16fffb3da 100644 --- a/client/milvusclient/mock_milvus_server_test.go +++ b/client/milvusclient/mock_milvus_server_test.go @@ -86,6 +86,65 @@ func (_c *MilvusServiceServer_AddCollectionField_Call) RunAndReturn(run func(con return _c } +// AddFileResource provides a mock function with given fields: _a0, _a1 +func (_m *MilvusServiceServer) AddFileResource(_a0 context.Context, _a1 *milvuspb.AddFileResourceRequest) (*commonpb.Status, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for AddFileResource") + } + + var r0 *commonpb.Status + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.AddFileResourceRequest) (*commonpb.Status, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.AddFileResourceRequest) *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, *milvuspb.AddFileResourceRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MilvusServiceServer_AddFileResource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddFileResource' +type MilvusServiceServer_AddFileResource_Call struct { + *mock.Call +} + +// AddFileResource is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.AddFileResourceRequest +func (_e *MilvusServiceServer_Expecter) AddFileResource(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_AddFileResource_Call { + return &MilvusServiceServer_AddFileResource_Call{Call: _e.mock.On("AddFileResource", _a0, _a1)} +} + +func (_c *MilvusServiceServer_AddFileResource_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.AddFileResourceRequest)) *MilvusServiceServer_AddFileResource_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.AddFileResourceRequest)) + }) + return _c +} + +func (_c *MilvusServiceServer_AddFileResource_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_AddFileResource_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MilvusServiceServer_AddFileResource_Call) RunAndReturn(run func(context.Context, *milvuspb.AddFileResourceRequest) (*commonpb.Status, error)) *MilvusServiceServer_AddFileResource_Call { + _c.Call.Return(run) + return _c +} + // AddUserTags provides a mock function with given fields: _a0, _a1 func (_m *MilvusServiceServer) AddUserTags(_a0 context.Context, _a1 *milvuspb.AddUserTagsRequest) (*commonpb.Status, error) { ret := _m.Called(_a0, _a1) @@ -558,6 +617,65 @@ func (_c *MilvusServiceServer_BackupRBAC_Call) RunAndReturn(run func(context.Con return _c } +// BatchDescribeCollection provides a mock function with given fields: _a0, _a1 +func (_m *MilvusServiceServer) BatchDescribeCollection(_a0 context.Context, _a1 *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for BatchDescribeCollection") + } + + var r0 *milvuspb.BatchDescribeCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) *milvuspb.BatchDescribeCollectionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*milvuspb.BatchDescribeCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MilvusServiceServer_BatchDescribeCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BatchDescribeCollection' +type MilvusServiceServer_BatchDescribeCollection_Call struct { + *mock.Call +} + +// BatchDescribeCollection is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.BatchDescribeCollectionRequest +func (_e *MilvusServiceServer_Expecter) BatchDescribeCollection(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_BatchDescribeCollection_Call { + return &MilvusServiceServer_BatchDescribeCollection_Call{Call: _e.mock.On("BatchDescribeCollection", _a0, _a1)} +} + +func (_c *MilvusServiceServer_BatchDescribeCollection_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.BatchDescribeCollectionRequest)) *MilvusServiceServer_BatchDescribeCollection_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.BatchDescribeCollectionRequest)) + }) + return _c +} + +func (_c *MilvusServiceServer_BatchDescribeCollection_Call) Return(_a0 *milvuspb.BatchDescribeCollectionResponse, _a1 error) *MilvusServiceServer_BatchDescribeCollection_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MilvusServiceServer_BatchDescribeCollection_Call) RunAndReturn(run func(context.Context, *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error)) *MilvusServiceServer_BatchDescribeCollection_Call { + _c.Call.Return(run) + return _c +} + // CalcDistance provides a mock function with given fields: _a0, _a1 func (_m *MilvusServiceServer) CalcDistance(_a0 context.Context, _a1 *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) { ret := _m.Called(_a0, _a1) @@ -4157,6 +4275,65 @@ func (_c *MilvusServiceServer_ListDatabases_Call) RunAndReturn(run func(context. return _c } +// ListFileResources provides a mock function with given fields: _a0, _a1 +func (_m *MilvusServiceServer) ListFileResources(_a0 context.Context, _a1 *milvuspb.ListFileResourcesRequest) (*milvuspb.ListFileResourcesResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for ListFileResources") + } + + var r0 *milvuspb.ListFileResourcesResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.ListFileResourcesRequest) (*milvuspb.ListFileResourcesResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.ListFileResourcesRequest) *milvuspb.ListFileResourcesResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*milvuspb.ListFileResourcesResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.ListFileResourcesRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MilvusServiceServer_ListFileResources_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListFileResources' +type MilvusServiceServer_ListFileResources_Call struct { + *mock.Call +} + +// ListFileResources is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.ListFileResourcesRequest +func (_e *MilvusServiceServer_Expecter) ListFileResources(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_ListFileResources_Call { + return &MilvusServiceServer_ListFileResources_Call{Call: _e.mock.On("ListFileResources", _a0, _a1)} +} + +func (_c *MilvusServiceServer_ListFileResources_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.ListFileResourcesRequest)) *MilvusServiceServer_ListFileResources_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.ListFileResourcesRequest)) + }) + return _c +} + +func (_c *MilvusServiceServer_ListFileResources_Call) Return(_a0 *milvuspb.ListFileResourcesResponse, _a1 error) *MilvusServiceServer_ListFileResources_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MilvusServiceServer_ListFileResources_Call) RunAndReturn(run func(context.Context, *milvuspb.ListFileResourcesRequest) (*milvuspb.ListFileResourcesResponse, error)) *MilvusServiceServer_ListFileResources_Call { + _c.Call.Return(run) + return _c +} + // ListImportTasks provides a mock function with given fields: _a0, _a1 func (_m *MilvusServiceServer) ListImportTasks(_a0 context.Context, _a1 *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) { ret := _m.Called(_a0, _a1) @@ -5219,6 +5396,65 @@ func (_c *MilvusServiceServer_ReleasePartitions_Call) RunAndReturn(run func(cont return _c } +// RemoveFileResource provides a mock function with given fields: _a0, _a1 +func (_m *MilvusServiceServer) RemoveFileResource(_a0 context.Context, _a1 *milvuspb.RemoveFileResourceRequest) (*commonpb.Status, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for RemoveFileResource") + } + + var r0 *commonpb.Status + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RemoveFileResourceRequest) (*commonpb.Status, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RemoveFileResourceRequest) *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, *milvuspb.RemoveFileResourceRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MilvusServiceServer_RemoveFileResource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveFileResource' +type MilvusServiceServer_RemoveFileResource_Call struct { + *mock.Call +} + +// RemoveFileResource is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.RemoveFileResourceRequest +func (_e *MilvusServiceServer_Expecter) RemoveFileResource(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_RemoveFileResource_Call { + return &MilvusServiceServer_RemoveFileResource_Call{Call: _e.mock.On("RemoveFileResource", _a0, _a1)} +} + +func (_c *MilvusServiceServer_RemoveFileResource_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.RemoveFileResourceRequest)) *MilvusServiceServer_RemoveFileResource_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.RemoveFileResourceRequest)) + }) + return _c +} + +func (_c *MilvusServiceServer_RemoveFileResource_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_RemoveFileResource_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MilvusServiceServer_RemoveFileResource_Call) RunAndReturn(run func(context.Context, *milvuspb.RemoveFileResourceRequest) (*commonpb.Status, error)) *MilvusServiceServer_RemoveFileResource_Call { + _c.Call.Return(run) + return _c +} + // RenameCollection provides a mock function with given fields: _a0, _a1 func (_m *MilvusServiceServer) RenameCollection(_a0 context.Context, _a1 *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) { ret := _m.Called(_a0, _a1) diff --git a/go.mod b/go.mod index fe55b78c49..460c938c3b 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 github.com/klauspost/compress v1.17.9 github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d - github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3 + github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 github.com/minio/minio-go/v7 v7.0.73 github.com/panjf2000/ants/v2 v2.11.3 // indirect github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 diff --git a/go.sum b/go.sum index 039ef57cf5..50cf7cc405 100644 --- a/go.sum +++ b/go.sum @@ -740,8 +740,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6 h1:YHMFI6L github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg= github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8= github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3 h1:AK0kGIj/nUiGkSodqtaOJPLeIvEYV2HICr8eaChQ4RA= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 h1:zyrKuc0rwT5xWIFkZr/bFWXXYbYvSBMT3iFITnaR8IE= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/pulsar-client-go v0.12.1 h1:O2JZp1tsYiO7C0MQ4hrUY/aJXnn2Gry6hpm7UodghmE= github.com/milvus-io/pulsar-client-go v0.12.1/go.mod h1:dkutuH4oS2pXiGm+Ti7fQZ4MRjrMPZ8IJeEGAWMeckk= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= diff --git a/internal/distributed/proxy/service.go b/internal/distributed/proxy/service.go index e5b250c435..32e148cd63 100644 --- a/internal/distributed/proxy/service.go +++ b/internal/distributed/proxy/service.go @@ -641,6 +641,11 @@ func (s *Server) DescribeCollection(ctx context.Context, request *milvuspb.Descr return s.proxy.DescribeCollection(ctx, request) } +// BatchDescribeCollection notifies Proxy to describe a collection +func (s *Server) BatchDescribeCollection(ctx context.Context, request *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error) { + return s.proxy.BatchDescribeCollection(ctx, request) +} + // AddCollectionField add a field to collection func (s *Server) AddCollectionField(ctx context.Context, request *milvuspb.AddCollectionFieldRequest) (*commonpb.Status, error) { return s.proxy.AddCollectionField(ctx, request) diff --git a/internal/distributed/proxy/service_test.go b/internal/distributed/proxy/service_test.go index 8de7c7b709..6ad45d3af1 100644 --- a/internal/distributed/proxy/service_test.go +++ b/internal/distributed/proxy/service_test.go @@ -274,6 +274,12 @@ func Test_NewServer(t *testing.T) { assert.NoError(t, err) }) + t.Run("BatchDescribeCollection", func(t *testing.T) { + mockProxy.EXPECT().BatchDescribeCollection(mock.Anything, mock.Anything).Return(nil, nil) + _, err := server.BatchDescribeCollection(ctx, nil) + assert.NoError(t, err) + }) + t.Run("GetCollectionStatistics", func(t *testing.T) { mockProxy.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(nil, nil) _, err := server.GetCollectionStatistics(ctx, nil) diff --git a/internal/mocks/mock_proxy.go b/internal/mocks/mock_proxy.go index b2d7bdda5c..b0fd9a62fd 100644 --- a/internal/mocks/mock_proxy.go +++ b/internal/mocks/mock_proxy.go @@ -623,6 +623,65 @@ func (_c *MockProxy_BackupRBAC_Call) RunAndReturn(run func(context.Context, *mil return _c } +// BatchDescribeCollection provides a mock function with given fields: _a0, _a1 +func (_m *MockProxy) BatchDescribeCollection(_a0 context.Context, _a1 *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error) { + ret := _m.Called(_a0, _a1) + + if len(ret) == 0 { + panic("no return value specified for BatchDescribeCollection") + } + + var r0 *milvuspb.BatchDescribeCollectionResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) *milvuspb.BatchDescribeCollectionResponse); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*milvuspb.BatchDescribeCollectionResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.BatchDescribeCollectionRequest) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockProxy_BatchDescribeCollection_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BatchDescribeCollection' +type MockProxy_BatchDescribeCollection_Call struct { + *mock.Call +} + +// BatchDescribeCollection is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.BatchDescribeCollectionRequest +func (_e *MockProxy_Expecter) BatchDescribeCollection(_a0 interface{}, _a1 interface{}) *MockProxy_BatchDescribeCollection_Call { + return &MockProxy_BatchDescribeCollection_Call{Call: _e.mock.On("BatchDescribeCollection", _a0, _a1)} +} + +func (_c *MockProxy_BatchDescribeCollection_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.BatchDescribeCollectionRequest)) *MockProxy_BatchDescribeCollection_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.BatchDescribeCollectionRequest)) + }) + return _c +} + +func (_c *MockProxy_BatchDescribeCollection_Call) Return(_a0 *milvuspb.BatchDescribeCollectionResponse, _a1 error) *MockProxy_BatchDescribeCollection_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockProxy_BatchDescribeCollection_Call) RunAndReturn(run func(context.Context, *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error)) *MockProxy_BatchDescribeCollection_Call { + _c.Call.Return(run) + return _c +} + // CalcDistance provides a mock function with given fields: _a0, _a1 func (_m *MockProxy) CalcDistance(_a0 context.Context, _a1 *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) { ret := _m.Called(_a0, _a1) diff --git a/internal/proxy/database_interceptor.go b/internal/proxy/database_interceptor.go index 31917941fb..4831eede4d 100644 --- a/internal/proxy/database_interceptor.go +++ b/internal/proxy/database_interceptor.go @@ -48,6 +48,11 @@ func fillDatabase(ctx context.Context, req interface{}) (context.Context, interf r.DbName = GetCurDBNameFromContextOrDefault(ctx) } return ctx, r + case *milvuspb.BatchDescribeCollectionRequest: + if r.DbName == "" { + r.DbName = GetCurDBNameFromContextOrDefault(ctx) + } + return ctx, r case *milvuspb.GetStatisticsRequest: if r.DbName == "" { r.DbName = GetCurDBNameFromContextOrDefault(ctx) diff --git a/internal/proxy/database_interceptor_test.go b/internal/proxy/database_interceptor_test.go index 08a6c328e3..69e1a25c01 100644 --- a/internal/proxy/database_interceptor_test.go +++ b/internal/proxy/database_interceptor_test.go @@ -53,6 +53,7 @@ func TestDatabaseInterceptor(t *testing.T) { &milvuspb.LoadCollectionRequest{}, &milvuspb.ReleaseCollectionRequest{}, &milvuspb.DescribeCollectionRequest{}, + &milvuspb.BatchDescribeCollectionRequest{}, &milvuspb.GetStatisticsRequest{}, &milvuspb.GetCollectionStatisticsRequest{}, &milvuspb.ShowCollectionsRequest{}, diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index bb22538b40..e4be0fc0eb 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -960,6 +960,40 @@ func (node *Proxy) DescribeCollection(ctx context.Context, request *milvuspb.Des return interceptor.Call(ctx, request) } +func (node *Proxy) BatchDescribeCollection(ctx context.Context, request *milvuspb.BatchDescribeCollectionRequest) (*milvuspb.BatchDescribeCollectionResponse, error) { + collectionNames := request.GetCollectionName() + if len(collectionNames) == 0 { + return &milvuspb.BatchDescribeCollectionResponse{ + Status: merr.Status(merr.WrapErrParameterInvalidMsg("collection names cannot be empty")), + }, nil + } + + responses := make([]*milvuspb.DescribeCollectionResponse, 0, len(collectionNames)) + + for _, collectionName := range collectionNames { + describeCollectionRequest := &milvuspb.DescribeCollectionRequest{ + DbName: request.GetDbName(), + CollectionName: collectionName, + } + + describeCollectionResponse, err := node.DescribeCollection(ctx, describeCollectionRequest) + // If there's an error, create a response with error status + if err != nil { + describeCollectionResponse = &milvuspb.DescribeCollectionResponse{ + Status: merr.Status(err), + CollectionName: collectionName, + } + } + + responses = append(responses, describeCollectionResponse) + } + + return &milvuspb.BatchDescribeCollectionResponse{ + Status: merr.Success(), + Responses: responses, + }, nil +} + // AddCollectionField add a field to collection func (node *Proxy) AddCollectionField(ctx context.Context, request *milvuspb.AddCollectionFieldRequest) (*commonpb.Status, error) { if err := merr.CheckHealthy(node.GetStateCode()); err != nil { @@ -4261,8 +4295,7 @@ func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.G SegmentIDs: getSegmentsByStatesResponse.Segments, }) if err != nil { - metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, - metrics.FailLabel, req.GetDbName(), req.GetCollectionName()).Inc() + metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel, req.GetDbName(), req.GetCollectionName()).Inc() log.Warn("GetPersistentSegmentInfo fail", zap.Error(err)) resp.Status = merr.Status(err) @@ -4270,8 +4303,7 @@ func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.G } err = merr.Error(infoResp.GetStatus()) if err != nil { - metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, - metrics.FailLabel, req.GetDbName(), req.GetCollectionName()).Inc() + metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel, req.GetDbName(), req.GetCollectionName()).Inc() resp.Status = merr.Status(err) return resp, nil } @@ -4291,8 +4323,7 @@ func (node *Proxy) GetPersistentSegmentInfo(ctx context.Context, req *milvuspb.G StorageVersion: info.GetStorageVersion(), } } - metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, - metrics.SuccessLabel, req.GetDbName(), req.GetCollectionName()).Inc() + metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.SuccessLabel, req.GetDbName(), req.GetCollectionName()).Inc() metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds())) resp.Infos = persistentInfos return resp, nil diff --git a/internal/proxy/impl_test.go b/internal/proxy/impl_test.go index 1d8a1c0bb6..c8e366f503 100644 --- a/internal/proxy/impl_test.go +++ b/internal/proxy/impl_test.go @@ -1203,6 +1203,18 @@ func TestProxyDescribeCollection(t *testing.T) { assert.Empty(t, resp.GetStatus().GetReason()) assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) }) + + t.Run("batch describe collection ok", func(t *testing.T) { + resp, err := node.BatchDescribeCollection(ctx, &milvuspb.BatchDescribeCollectionRequest{ + DbName: "test_1", + CollectionName: []string{"test_collection"}, + }) + assert.NoError(t, err) + assert.Empty(t, resp.GetStatus().GetReason()) + assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) + assert.Equal(t, 1, len(resp.GetResponses())) + assert.Equal(t, commonpb.ErrorCode_Success, resp.GetResponses()[0].GetStatus().GetErrorCode()) + }) } func TestProxy_AllocTimestamp(t *testing.T) { diff --git a/internal/proxy/task.go b/internal/proxy/task.go index 0669b75a0a..8a1d543e67 100644 --- a/internal/proxy/task.go +++ b/internal/proxy/task.go @@ -1008,6 +1008,7 @@ func (t *showCollectionsTask) Execute(ctx context.Context) error { CreatedUtcTimestamps: make([]uint64, 0, len(resp.CollectionIDs)), InMemoryPercentages: make([]int64, 0, len(resp.CollectionIDs)), QueryServiceAvailable: make([]bool, 0, len(resp.CollectionIDs)), + ShardsNum: make([]int32, 0, len(resp.CollectionIDs)), } for offset, id := range resp.CollectionIDs { @@ -1030,6 +1031,7 @@ func (t *showCollectionsTask) Execute(ctx context.Context) error { t.result.CreatedUtcTimestamps = append(t.result.CreatedUtcTimestamps, collectionInfo.createdUtcTimestamp) t.result.InMemoryPercentages = append(t.result.InMemoryPercentages, resp.InMemoryPercentages[offset]) t.result.QueryServiceAvailable = append(t.result.QueryServiceAvailable, resp.QueryServiceAvailable[offset]) + t.result.ShardsNum = append(t.result.ShardsNum, collectionInfo.shardsNum) } } else { t.result = respFromRootCoord diff --git a/internal/rootcoord/show_collection_task.go b/internal/rootcoord/show_collection_task.go index 9a8a19b009..6ca296fa70 100644 --- a/internal/rootcoord/show_collection_task.go +++ b/internal/rootcoord/show_collection_task.go @@ -77,6 +77,7 @@ func (t *showCollectionTask) Execute(ctx context.Context) error { t.Rsp.CreatedTimestamps = append(t.Rsp.CreatedTimestamps, coll.CreateTime) physical, _ := tsoutil.ParseHybridTs(coll.CreateTime) t.Rsp.CreatedUtcTimestamps = append(t.Rsp.CreatedUtcTimestamps, uint64(physical)) + t.Rsp.ShardsNum = append(t.Rsp.ShardsNum, coll.ShardsNum) } return nil } diff --git a/internal/rootcoord/show_collection_task_test.go b/internal/rootcoord/show_collection_task_test.go index e133c94ecf..570ba4a476 100644 --- a/internal/rootcoord/show_collection_task_test.go +++ b/internal/rootcoord/show_collection_task_test.go @@ -82,10 +82,12 @@ func Test_showCollectionTask_Execute(t *testing.T) { meta.ListCollectionsFunc = func(ctx context.Context, ts Timestamp) ([]*model.Collection, error) { return []*model.Collection{ { - Name: "test coll", + Name: "test coll", + ShardsNum: 2, }, { - Name: "test coll2", + Name: "test coll2", + ShardsNum: 3, }, }, nil } @@ -103,6 +105,8 @@ func Test_showCollectionTask_Execute(t *testing.T) { assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, task.Rsp.GetStatus().GetErrorCode()) assert.Equal(t, 2, len(task.Rsp.GetCollectionNames())) + assert.Equal(t, int32(2), task.Rsp.GetShardsNum()[0]) + assert.Equal(t, int32(3), task.Rsp.GetShardsNum()[1]) }) } @@ -581,6 +585,7 @@ func TestShowCollectionsAuth(t *testing.T) { CollectionID: 100, Name: "test_collection", CreateTime: tsoutil.GetCurrentTime(), + ShardsNum: 2, }, }, nil).Once() @@ -594,5 +599,6 @@ func TestShowCollectionsAuth(t *testing.T) { assert.NoError(t, err) assert.Equal(t, 1, len(task.Rsp.GetCollectionNames())) assert.Equal(t, "test_collection", task.Rsp.GetCollectionNames()[0]) + assert.Equal(t, int32(2), task.Rsp.GetShardsNum()[0]) }) } diff --git a/pkg/go.mod b/pkg/go.mod index e4eca7c241..ee6e95e2ef 100644 --- a/pkg/go.mod +++ b/pkg/go.mod @@ -21,7 +21,7 @@ require ( github.com/jolestar/go-commons-pool/v2 v2.1.2 github.com/json-iterator/go v1.1.12 github.com/klauspost/compress v1.17.9 - github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3 + github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 github.com/minio/minio-go/v7 v7.0.73 github.com/panjf2000/ants/v2 v2.11.3 github.com/prometheus/client_golang v1.14.0 diff --git a/pkg/go.sum b/pkg/go.sum index 6094ae3f11..54848a02d3 100644 --- a/pkg/go.sum +++ b/pkg/go.sum @@ -559,8 +559,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6 h1:YHMFI6L github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg= github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8= github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3 h1:AK0kGIj/nUiGkSodqtaOJPLeIvEYV2HICr8eaChQ4RA= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250806032933-1b94535b80c3/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 h1:zyrKuc0rwT5xWIFkZr/bFWXXYbYvSBMT3iFITnaR8IE= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/pulsar-client-go v0.12.1 h1:O2JZp1tsYiO7C0MQ4hrUY/aJXnn2Gry6hpm7UodghmE= github.com/milvus-io/pulsar-client-go v0.12.1/go.mod h1:dkutuH4oS2pXiGm+Ti7fQZ4MRjrMPZ8IJeEGAWMeckk= github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34= diff --git a/tests/go_client/go.mod b/tests/go_client/go.mod index 3e84f1ac29..990ef40db6 100644 --- a/tests/go_client/go.mod +++ b/tests/go_client/go.mod @@ -51,7 +51,7 @@ require ( github.com/kr/text v0.2.0 // indirect github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect - github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0 // indirect + github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/opencontainers/runtime-spec v1.0.2 // indirect diff --git a/tests/go_client/go.sum b/tests/go_client/go.sum index a2af448136..4609a4ed92 100644 --- a/tests/go_client/go.sum +++ b/tests/go_client/go.sum @@ -318,8 +318,10 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8= github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc= github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0 h1:SJLdKkifvPDT31jw0hzC7/KSlZ5tv5AhPt77jNHMAQY= -github.com/milvus-io/milvus-proto/go-api/v2 v2.6.0/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807040333-531631e7fce6 h1:qTBOTsZ3OwEXkrHRqPn562ddkDqeToIY6CstLIaVQYs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807040333-531631e7fce6/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17 h1:zyrKuc0rwT5xWIFkZr/bFWXXYbYvSBMT3iFITnaR8IE= +github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1-0.20250807065533-ebdc11f5df17/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e h1:VCr43pG4efacDbM4au70fh8/5hNTftoWzm1iEumvDWM= github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e/go.mod h1:37AWzxVs2NS4QUJrkcbeLUwi+4Av0h5mEdjLI62EANU= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=