fix: resolve collection alias in RBAC and clean up grants on collection lifecycle (#47851)

related: https://github.com/milvus-io/milvus/issues/47850
1. Privilege interceptor: resolve collection alias to real collection
name
   before RBAC permission check, so that operating via alias checks
   permission against the real collection, not the alias name.

2. MetaCache alias cache: add aliasInfo cache with positive/negative
   entries to avoid repeated DescribeAlias RPC calls. Cache is
invalidated on alias removal, collection removal, and database removal.

3. Catalog grant cleanup: add DeleteGrantByCollectionName and
   MigrateGrantCollectionName to RootCoordCatalog interface and kv
   implementation. On collection drop, delete all associated grants;
   on collection rename, migrate grants to the new name.

4. Feature flag: add proxy.resolveAliasForPrivilege config to
   enable/disable alias resolution in the privilege interceptor.

---------

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
sthuang
2026-03-05 18:39:22 +08:00
committed by GitHub
parent 96df7c9043
commit f5b6db2929
24 changed files with 1797 additions and 18 deletions
+1
View File
@@ -337,6 +337,7 @@ proxy:
ddlConcurrency: 16 # The concurrent execution number of DDL at proxy.
dclConcurrency: 16 # The concurrent execution number of DCL at proxy.
mustUsePartitionKey: false # switch for whether proxy must use partition key for the collection
resolveAliasForPrivilege: false # switch for whether proxy shall resolve alias to actual collection name during RBAC privilege checks
# maximum number of result entries, typically Nq * TopK * GroupSize.
# It costs additional memory and time to process a large number of result entries.
# If the number of result entries exceeds this limit, the search will be rejected.
+5
View File
@@ -84,6 +84,11 @@ type RootCoordCatalog interface {
// For example []string{"user1/role1"}
ListUserRole(ctx context.Context, tenant string) ([]string, error)
// DeleteGrantByCollectionName deletes all grants for a specific collection.
DeleteGrantByCollectionName(ctx context.Context, tenant string, dbName string, collectionName string) error
// MigrateGrantCollectionName migrates all grants from oldName to newName when a collection is renamed.
MigrateGrantCollectionName(ctx context.Context, tenant string, oldDBName string, oldName string, newDBName string, newName string) error
ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error)
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
@@ -1421,6 +1421,107 @@ func (kc *Catalog) ListGrant(ctx context.Context, tenant string, entity *milvusp
return entities, nil
}
func (kc *Catalog) DeleteGrantByCollectionName(ctx context.Context, tenant string, dbName string, collectionName string) error {
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey)
if err != nil {
log.Ctx(ctx).Warn("fail to load grant privilege entities for collection cleanup",
zap.String("key", granteeKey), zap.Error(err))
return err
}
var exactRemoveKeys []string
var prefixRemoveKeys []string
for i, key := range keys {
grantInfos := typeutil.AfterN(key, granteeKey+"/", "/")
if len(grantInfos) != 3 {
continue
}
// grantInfos: [role, objectType, dbName.objectName]
if grantInfos[1] != "Collection" {
continue
}
grantDB, grantObj := funcutil.SplitObjectName(grantInfos[2])
if grantObj == collectionName && grantDB == dbName {
// Use exact deletion for the grantee key
exactRemoveKeys = append(exactRemoveKeys, key)
// Use prefix deletion for the granteeID key (has sub-keys)
granteeIDKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, values[i]+"/")
prefixRemoveKeys = append(prefixRemoveKeys, granteeIDKey)
}
}
if len(exactRemoveKeys) == 0 && len(prefixRemoveKeys) == 0 {
return nil
}
// Use prefix deletion for granteeID keys (which have sub-keys underneath).
if len(prefixRemoveKeys) > 0 {
if err = kc.Txn.MultiSaveAndRemoveWithPrefix(ctx, nil, prefixRemoveKeys); err != nil {
log.Ctx(ctx).Warn("fail to remove granteeID entries for collection",
zap.String("dbName", dbName), zap.String("collectionName", collectionName), zap.Error(err))
return err
}
}
// Use exact deletion for grantee keys (leaf keys with no sub-keys).
// Avoid MultiSaveAndRemoveWithPrefix here to prevent accidentally matching
// keys like col1_backup when removing col1.
if len(exactRemoveKeys) > 0 {
if err = kc.Txn.MultiSaveAndRemove(ctx, nil, exactRemoveKeys); err != nil {
log.Ctx(ctx).Warn("fail to remove grantee entries for collection",
zap.String("dbName", dbName), zap.String("collectionName", collectionName), zap.Error(err))
return err
}
}
return nil
}
func (kc *Catalog) MigrateGrantCollectionName(ctx context.Context, tenant string, oldDBName string, oldName string, newDBName string, newName string) error {
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
keys, values, err := kc.Txn.LoadWithPrefix(ctx, granteeKey)
if err != nil {
log.Ctx(ctx).Warn("fail to load grant privilege entities for collection migration",
zap.String("key", granteeKey), zap.Error(err))
return err
}
saves := make(map[string]string)
var removeKeys []string
for i, key := range keys {
grantInfos := typeutil.AfterN(key, granteeKey+"/", "/")
if len(grantInfos) != 3 {
continue
}
if grantInfos[1] != "Collection" {
continue
}
grantDB, grantObj := funcutil.SplitObjectName(grantInfos[2])
if grantObj == oldName && grantDB == oldDBName {
// Build new key with new collection name
newObjName := funcutil.CombineObjectName(newDBName, newName)
newKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
fmt.Sprintf("%s/%s/%s", grantInfos[0], grantInfos[1], newObjName))
saves[newKey] = values[i]
removeKeys = append(removeKeys, key)
}
}
if len(removeKeys) == 0 {
return nil
}
// Use MultiSaveAndRemove (exact deletion) instead of prefix-based deletion
// to avoid accidentally matching keys like col1_backup when removing col1
if err = kc.Txn.MultiSaveAndRemove(ctx, saves, removeKeys); err != nil {
log.Ctx(ctx).Warn("fail to migrate grants for renamed collection",
zap.String("oldDBName", oldDBName), zap.String("oldName", oldName),
zap.String("newDBName", newDBName), zap.String("newName", newName), zap.Error(err))
}
return err
}
func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
var (
k = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, role.Name+"/")
@@ -3488,3 +3488,208 @@ func TestCatalog_FileResource(t *testing.T) {
assert.Equal(t, 2, len(resources))
})
}
func TestDeleteGrantByCollectionName(t *testing.T) {
ctx := context.Background()
tenant := "test-tenant"
t.Run("load error", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(nil, nil, errors.New("load error"))
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.Error(t, err)
})
t.Run("no matching grants", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
// Only a grant for a different collection
key1 := granteeKey + "/role1/Collection/default.other_col"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"granteeID1"}, nil)
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.NoError(t, err)
})
t.Run("deletes matching grants only", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.col1"
key2 := granteeKey + "/role2/Collection/default.col1"
key3 := granteeKey + "/role1/Collection/default.other_col"
key4 := granteeKey + "/role1/Global/default.col1" // different objectType
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1, key2, key3, key4},
[]string{"gid1", "gid2", "gid3", "gid4"},
nil)
// Prefix deletion for granteeID keys (have sub-keys)
gid1Key := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "gid1/")
gid2Key := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "gid2/")
kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, (map[string]string)(nil),
[]string{gid1Key, gid2Key}).Return(nil)
// Exact deletion for grantee keys (leaf keys, no sub-keys)
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, (map[string]string)(nil),
[]string{key1, key2}).Return(nil)
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.NoError(t, err)
})
t.Run("handles old format without db prefix", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
// Old format: objectName without "." is treated as default db
key1 := granteeKey + "/role1/Collection/col1"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
gid1Key := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "gid1/")
kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, (map[string]string)(nil),
[]string{gid1Key}).Return(nil)
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, (map[string]string)(nil),
[]string{key1}).Return(nil)
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.NoError(t, err)
})
t.Run("MultiSaveAndRemoveWithPrefix error", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.col1"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
gid1Key := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "gid1/")
kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, (map[string]string)(nil),
[]string{gid1Key}).Return(errors.New("remove error"))
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.Error(t, err)
})
t.Run("MultiSaveAndRemove error", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.col1"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
gid1Key := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "gid1/")
kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, (map[string]string)(nil),
[]string{gid1Key}).Return(nil)
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, (map[string]string)(nil),
[]string{key1}).Return(errors.New("remove error"))
err := c.DeleteGrantByCollectionName(ctx, tenant, "default", "col1")
assert.Error(t, err)
})
}
func TestMigrateGrantCollectionName(t *testing.T) {
ctx := context.Background()
tenant := "test-tenant"
t.Run("load error", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(nil, nil, errors.New("load error"))
err := c.MigrateGrantCollectionName(ctx, tenant, "default", "old_col", "default", "new_col")
assert.Error(t, err)
})
t.Run("no matching grants", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.other_col"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
err := c.MigrateGrantCollectionName(ctx, tenant, "default", "old_col", "default", "new_col")
assert.NoError(t, err)
})
t.Run("migrates matching grants", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.old_col"
key2 := granteeKey + "/role2/Collection/default.old_col"
key3 := granteeKey + "/role1/Collection/default.other_col"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1, key2, key3},
[]string{"gid1", "gid2", "gid3"},
nil)
newKey1 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
fmt.Sprintf("role1/Collection/%s", funcutil.CombineObjectName("default", "new_col")))
newKey2 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
fmt.Sprintf("role2/Collection/%s", funcutil.CombineObjectName("default", "new_col")))
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything,
map[string]string{newKey1: "gid1", newKey2: "gid2"},
[]string{key1, key2}).Return(nil)
err := c.MigrateGrantCollectionName(ctx, tenant, "default", "old_col", "default", "new_col")
assert.NoError(t, err)
})
t.Run("cross-db rename", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/db1.col1"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
newKey1 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
fmt.Sprintf("role1/Collection/%s", funcutil.CombineObjectName("db2", "col2")))
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything,
map[string]string{newKey1: "gid1"},
[]string{key1}).Return(nil)
err := c.MigrateGrantCollectionName(ctx, tenant, "db1", "col1", "db2", "col2")
assert.NoError(t, err)
})
t.Run("save error", func(t *testing.T) {
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock, nil)
granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
key1 := granteeKey + "/role1/Collection/default.old_col"
kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeKey).Return(
[]string{key1}, []string{"gid1"}, nil)
kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).
Return(errors.New("save error"))
err := c.MigrateGrantCollectionName(ctx, tenant, "default", "old_col", "default", "new_col")
assert.Error(t, err)
})
}
@@ -3527,8 +3527,7 @@ func (_c *DataCoordCatalog_ShouldDropChannel_Call) RunAndReturn(run func(context
func NewDataCoordCatalog(t interface {
mock.TestingT
Cleanup(func())
},
) *DataCoordCatalog {
}) *DataCoordCatalog {
mock := &DataCoordCatalog{}
mock.Mock.Test(t)
@@ -850,6 +850,55 @@ func (_c *RootCoordCatalog_DeleteGrant_Call) RunAndReturn(run func(context.Conte
return _c
}
// DeleteGrantByCollectionName provides a mock function with given fields: ctx, tenant, dbName, collectionName
func (_m *RootCoordCatalog) DeleteGrantByCollectionName(ctx context.Context, tenant string, dbName string, collectionName string) error {
ret := _m.Called(ctx, tenant, dbName, collectionName)
if len(ret) == 0 {
panic("no return value specified for DeleteGrantByCollectionName")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, string, string) error); ok {
r0 = rf(ctx, tenant, dbName, collectionName)
} else {
r0 = ret.Error(0)
}
return r0
}
// RootCoordCatalog_DeleteGrantByCollectionName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteGrantByCollectionName'
type RootCoordCatalog_DeleteGrantByCollectionName_Call struct {
*mock.Call
}
// DeleteGrantByCollectionName is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
// - dbName string
// - collectionName string
func (_e *RootCoordCatalog_Expecter) DeleteGrantByCollectionName(ctx interface{}, tenant interface{}, dbName interface{}, collectionName interface{}) *RootCoordCatalog_DeleteGrantByCollectionName_Call {
return &RootCoordCatalog_DeleteGrantByCollectionName_Call{Call: _e.mock.On("DeleteGrantByCollectionName", ctx, tenant, dbName, collectionName)}
}
func (_c *RootCoordCatalog_DeleteGrantByCollectionName_Call) Run(run func(ctx context.Context, tenant string, dbName string, collectionName string)) *RootCoordCatalog_DeleteGrantByCollectionName_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(string))
})
return _c
}
func (_c *RootCoordCatalog_DeleteGrantByCollectionName_Call) Return(_a0 error) *RootCoordCatalog_DeleteGrantByCollectionName_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *RootCoordCatalog_DeleteGrantByCollectionName_Call) RunAndReturn(run func(context.Context, string, string, string) error) *RootCoordCatalog_DeleteGrantByCollectionName_Call {
_c.Call.Return(run)
return _c
}
// DropAlias provides a mock function with given fields: ctx, dbID, alias, ts
func (_m *RootCoordCatalog) DropAlias(ctx context.Context, dbID int64, alias string, ts uint64) error {
ret := _m.Called(ctx, dbID, alias, ts)
@@ -2146,6 +2195,57 @@ func (_c *RootCoordCatalog_ListUserRole_Call) RunAndReturn(run func(context.Cont
return _c
}
// MigrateGrantCollectionName provides a mock function with given fields: ctx, tenant, oldDBName, oldName, newDBName, newName
func (_m *RootCoordCatalog) MigrateGrantCollectionName(ctx context.Context, tenant string, oldDBName string, oldName string, newDBName string, newName string) error {
ret := _m.Called(ctx, tenant, oldDBName, oldName, newDBName, newName)
if len(ret) == 0 {
panic("no return value specified for MigrateGrantCollectionName")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, string, string, string, string) error); ok {
r0 = rf(ctx, tenant, oldDBName, oldName, newDBName, newName)
} else {
r0 = ret.Error(0)
}
return r0
}
// RootCoordCatalog_MigrateGrantCollectionName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'MigrateGrantCollectionName'
type RootCoordCatalog_MigrateGrantCollectionName_Call struct {
*mock.Call
}
// MigrateGrantCollectionName is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
// - oldDBName string
// - oldName string
// - newDBName string
// - newName string
func (_e *RootCoordCatalog_Expecter) MigrateGrantCollectionName(ctx interface{}, tenant interface{}, oldDBName interface{}, oldName interface{}, newDBName interface{}, newName interface{}) *RootCoordCatalog_MigrateGrantCollectionName_Call {
return &RootCoordCatalog_MigrateGrantCollectionName_Call{Call: _e.mock.On("MigrateGrantCollectionName", ctx, tenant, oldDBName, oldName, newDBName, newName)}
}
func (_c *RootCoordCatalog_MigrateGrantCollectionName_Call) Run(run func(ctx context.Context, tenant string, oldDBName string, oldName string, newDBName string, newName string)) *RootCoordCatalog_MigrateGrantCollectionName_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(string), args[4].(string), args[5].(string))
})
return _c
}
func (_c *RootCoordCatalog_MigrateGrantCollectionName_Call) Return(_a0 error) *RootCoordCatalog_MigrateGrantCollectionName_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *RootCoordCatalog_MigrateGrantCollectionName_Call) RunAndReturn(run func(context.Context, string, string, string, string, string) error) *RootCoordCatalog_MigrateGrantCollectionName_Call {
_c.Call.Return(run)
return _c
}
// RemoveFileResource provides a mock function with given fields: ctx, resourceID, version
func (_m *RootCoordCatalog) RemoveFileResource(ctx context.Context, resourceID int64, version uint64) error {
ret := _m.Called(ctx, resourceID, version)
+658
View File
@@ -0,0 +1,658 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/pkg/v2/util/conc"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
// Tests for ResolveCollectionAlias in MetaCache and resolveCollectionAlias helper.
func TestResolveCollectionAlias_WildcardAndEmptySkippedByInterceptor(t *testing.T) {
// The interceptor guard (objectName != util.AnyWord && objectName != "") ensures
// that "*" and "" never reach resolveCollectionAlias. This test verifies those
// values are correctly guarded at the interceptor level by checking that the
// cache's DescribeAlias RPC is never called for these sentinel values.
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
// Wildcard "*" should not trigger resolution
for _, sentinel := range []string{"*", ""} {
// These sentinel values should be caught by the interceptor guard
// (objectName != util.AnyWord && objectName != "") before calling
// resolveCollectionAlias. Verify the guard logic directly.
shouldSkip := sentinel == "*" || sentinel == ""
assert.True(t, shouldSkip, "sentinel %q should be skipped by interceptor guard", sentinel)
}
// Normal name should resolve via RPC
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
Collection: "real_col",
}, nil)
result, err := resolveCollectionAlias(ctx, "default", "some_alias")
assert.NoError(t, err)
assert.Equal(t, "real_col", result)
mockCoord.AssertCalled(t, "DescribeAlias", mock.Anything, mock.Anything)
}
func TestResolveCollectionAlias_CachedCollection(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {
"test_collection": {
collID: 1,
schema: &schemaInfo{},
},
},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
result, err := cache.ResolveCollectionAlias(ctx, "default", "test_collection")
assert.NoError(t, err)
assert.Equal(t, "test_collection", result)
mockCoord.AssertNotCalled(t, "DescribeAlias")
}
func TestResolveCollectionAlias_ValidAlias(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
mockCoord.On("DescribeAlias", mock.Anything, mock.MatchedBy(func(req *milvuspb.DescribeAliasRequest) bool {
return req.DbName == "default" && req.Alias == "my_alias"
})).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Collection: "actual_collection",
}, nil)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
result, err := cache.ResolveCollectionAlias(ctx, "default", "my_alias")
assert.NoError(t, err)
assert.Equal(t, "actual_collection", result)
}
func TestResolveCollectionAlias_NotAnAlias(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
mockCoord.On("DescribeAlias", mock.Anything, mock.MatchedBy(func(req *milvuspb.DescribeAliasRequest) bool {
return req.DbName == "default" && req.Alias == "not_an_alias"
})).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: "alias not found",
},
}, nil)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
// Should return original name when alias not found (not error)
result, err := cache.ResolveCollectionAlias(ctx, "default", "not_an_alias")
assert.NoError(t, err)
assert.Equal(t, "not_an_alias", result)
}
func TestResolveCollectionAlias_RPCError(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).
Return(nil, assert.AnError)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
result, err := cache.ResolveCollectionAlias(ctx, "default", "some_name")
assert.Error(t, err)
assert.Equal(t, "", result)
}
func TestResolveCollectionAlias_InternalServerError(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Internal server error should be propagated, not swallowed
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "internal error",
},
}, nil)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
result, err := cache.ResolveCollectionAlias(ctx, "default", "some_name")
assert.Error(t, err)
assert.Equal(t, "", result)
}
func TestResolveCollectionAlias_EmptyCollectionInResponse(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Status is success but collection field is empty
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Collection: "",
}, nil)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
result, err := cache.ResolveCollectionAlias(ctx, "default", "some_alias")
assert.NoError(t, err)
assert.Equal(t, "some_alias", result)
}
func TestResolveCollectionAlias_NilGlobalMetaCache(t *testing.T) {
ctx := context.Background()
oldCache := globalMetaCache
globalMetaCache = nil
defer func() { globalMetaCache = oldCache }()
result, err := resolveCollectionAlias(ctx, "default", "test")
assert.Error(t, err)
assert.Equal(t, "test", result)
}
func TestResolveCollectionAlias_AliasCacheHit(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"my_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
// Should return cached alias without RPC call
result, err := cache.ResolveCollectionAlias(ctx, "default", "my_alias")
assert.NoError(t, err)
assert.Equal(t, "real_collection", result)
mockCoord.AssertNotCalled(t, "DescribeAlias")
}
func TestResolveCollectionAlias_NegativeCacheHit(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"regular_name": {collectionName: "", cachedAt: time.Now()},
},
},
}
// Negative cache: not an alias, return as-is without RPC
result, err := cache.ResolveCollectionAlias(ctx, "default", "regular_name")
assert.NoError(t, err)
assert.Equal(t, "regular_name", result)
mockCoord.AssertNotCalled(t, "DescribeAlias")
}
func TestResolveCollectionAlias_PopulatesPositiveCache(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Collection: "real_collection",
}, nil).Once()
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
// First call triggers RPC
result, err := cache.ResolveCollectionAlias(ctx, "default", "my_alias")
assert.NoError(t, err)
assert.Equal(t, "real_collection", result)
// Second call should hit cache, no additional RPC
result, err = cache.ResolveCollectionAlias(ctx, "default", "my_alias")
assert.NoError(t, err)
assert.Equal(t, "real_collection", result)
// DescribeAlias should have been called exactly once
mockCoord.AssertNumberOfCalls(t, "DescribeAlias", 1)
}
func TestResolveCollectionAlias_PopulatesNegativeCache(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
mockCoord.On("DescribeAlias", mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: "alias not found",
},
}, nil).Once()
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{},
}
// First call triggers RPC
result, err := cache.ResolveCollectionAlias(ctx, "default", "not_alias")
assert.NoError(t, err)
assert.Equal(t, "not_alias", result)
// Second call should hit negative cache
result, err = cache.ResolveCollectionAlias(ctx, "default", "not_alias")
assert.NoError(t, err)
assert.Equal(t, "not_alias", result)
mockCoord.AssertNumberOfCalls(t, "DescribeAlias", 1)
}
func TestRemoveAlias_InvalidatesCache(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"my_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
// Remove the alias
cache.RemoveAlias(ctx, "default", "my_alias")
// Verify the alias was removed from cache
_, ok := cache.getAlias("default", "my_alias")
assert.False(t, ok)
}
func TestRemoveCollectionByID_CleansUpAliases(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {
"my_collection": {
collID: 100,
schema: &schemaInfo{},
},
},
},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"alias1": {collectionName: "my_collection", cachedAt: time.Now()},
"alias2": {collectionName: "my_collection", cachedAt: time.Now()},
"alias3": {collectionName: "other_collection", cachedAt: time.Now()},
},
},
collectionCacheVersion: make(map[UniqueID]uint64),
sfGlobal: conc.Singleflight[*collectionInfo]{},
}
// Remove collection by ID
cache.RemoveCollectionsByID(ctx, 100, 0, true)
// Aliases pointing to my_collection should be removed
_, ok := cache.getAlias("default", "alias1")
assert.False(t, ok)
_, ok = cache.getAlias("default", "alias2")
assert.False(t, ok)
// Alias pointing to other_collection should remain
entry, ok := cache.getAlias("default", "alias3")
assert.True(t, ok)
assert.Equal(t, "other_collection", entry.collectionName)
}
func TestRemoveCollection_CleansUpAliasesWhenNotCached(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Collection is NOT in collInfo cache, but aliases pointing to it exist
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"default": {},
},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"alias1": {collectionName: "uncached_collection", cachedAt: time.Now()},
"alias2": {collectionName: "uncached_collection", cachedAt: time.Now()},
"alias3": {collectionName: "other_collection", cachedAt: time.Now()},
},
},
collectionCacheVersion: make(map[UniqueID]uint64),
sfGlobal: conc.Singleflight[*collectionInfo]{},
}
// RemoveCollection for an uncached collection
cache.RemoveCollection(ctx, "default", "uncached_collection", 0)
// Aliases pointing to uncached_collection should be removed
_, ok := cache.getAlias("default", "alias1")
assert.False(t, ok)
_, ok = cache.getAlias("default", "alias2")
assert.False(t, ok)
// Alias pointing to other_collection should remain
entry, ok := cache.getAlias("default", "alias3")
assert.True(t, ok)
assert.Equal(t, "other_collection", entry.collectionName)
}
func TestRemoveDatabase_CleansUpAliases(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{
"mydb": {},
},
aliasInfo: map[string]map[string]*aliasEntry{
"mydb": {
"alias1": {collectionName: "coll1", cachedAt: time.Now()},
},
},
dbInfo: map[string]*databaseInfo{
"mydb": {},
},
}
cache.RemoveDatabase(ctx, "mydb")
// Alias cache for the database should be gone
_, ok := cache.getAlias("mydb", "alias1")
assert.False(t, ok)
}
func TestCreateAliasTask_ResolvesCollectionAlias(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Set up cache: "existing_alias" -> "real_collection"
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"existing_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
paramtable.Init()
task := &CreateAliasTask{
Condition: NewTaskCondition(ctx),
CreateAliasRequest: &milvuspb.CreateAliasRequest{
Base: &commonpb.MsgBase{},
DbName: "default",
CollectionName: "existing_alias",
Alias: "new_alias",
},
ctx: ctx,
mixCoord: mockCoord,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
// CollectionName should always be resolved from alias to real collection
assert.Equal(t, "real_collection", task.CollectionName)
}
func TestAlterAliasTask_ResolvesCollectionAlias(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Set up cache: "existing_alias" -> "real_collection"
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"existing_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
paramtable.Init()
task := &AlterAliasTask{
Condition: NewTaskCondition(ctx),
AlterAliasRequest: &milvuspb.AlterAliasRequest{
Base: &commonpb.MsgBase{},
DbName: "default",
CollectionName: "existing_alias",
Alias: "some_alias",
},
ctx: ctx,
mixCoord: mockCoord,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
// CollectionName should always be resolved from alias to real collection
assert.Equal(t, "real_collection", task.CollectionName)
}
func TestCreateAliasTask_ResolvesEvenWhenRBACFlagDisabled(t *testing.T) {
// Alias resolution in CreateAlias/AlterAlias is unconditional (for correctness),
// independent of the RBAC feature flag.
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"existing_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
paramtable.Init()
paramtable.Get().Save(Params.ProxyCfg.ResolveAliasForPrivilege.Key, "false")
defer paramtable.Get().Reset(Params.ProxyCfg.ResolveAliasForPrivilege.Key)
task := &CreateAliasTask{
Condition: NewTaskCondition(ctx),
CreateAliasRequest: &milvuspb.CreateAliasRequest{
Base: &commonpb.MsgBase{},
DbName: "default",
CollectionName: "existing_alias",
Alias: "new_alias",
},
ctx: ctx,
mixCoord: mockCoord,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
// CollectionName should be resolved even when RBAC flag is disabled
assert.Equal(t, "real_collection", task.CollectionName)
}
func TestListAliasesTask_ResolvesCollectionAlias(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
// Set up cache: "existing_alias" -> "real_collection"
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{
"default": {
"existing_alias": {collectionName: "real_collection", cachedAt: time.Now()},
},
},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
paramtable.Init()
task := &ListAliasesTask{
Condition: NewTaskCondition(ctx),
ListAliasesRequest: &milvuspb.ListAliasesRequest{
Base: &commonpb.MsgBase{},
DbName: "default",
CollectionName: "existing_alias",
},
ctx: ctx,
mixCoord: mockCoord,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
// CollectionName should be resolved from alias to real collection
assert.Equal(t, "real_collection", task.CollectionName)
}
func TestListAliasesTask_NoResolveWhenCollectionNameEmpty(t *testing.T) {
ctx := context.Background()
mockCoord := mocks.NewMockMixCoordClient(t)
cache := &MetaCache{
mixCoord: mockCoord,
collInfo: map[string]map[string]*collectionInfo{"default": {}},
aliasInfo: map[string]map[string]*aliasEntry{},
}
oldCache := globalMetaCache
globalMetaCache = cache
defer func() { globalMetaCache = oldCache }()
paramtable.Init()
task := &ListAliasesTask{
Condition: NewTaskCondition(ctx),
ListAliasesRequest: &milvuspb.ListAliasesRequest{
Base: &commonpb.MsgBase{},
DbName: "default",
},
ctx: ctx,
mixCoord: mockCoord,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
// CollectionName should remain empty
assert.Equal(t, "", task.CollectionName)
mockCoord.AssertNotCalled(t, "DescribeAlias")
}
+6
View File
@@ -158,6 +158,12 @@ func (node *Proxy) InvalidateCollectionMetaCache(ctx context.Context, request *p
node.shardMgr.DeprecateShardCache(request.GetDbName(), name)
}
}
// Invalidate alias cache for alias operations
if msgType == commonpb.MsgType_CreateAlias || msgType == commonpb.MsgType_AlterAlias || msgType == commonpb.MsgType_DropAlias {
if collectionName != "" {
globalMetaCache.RemoveAlias(ctx, request.GetDbName(), collectionName)
}
}
log.Info("complete to invalidate collection meta cache with collection name", zap.String("type", request.GetBase().GetMsgType().String()))
case commonpb.MsgType_LoadCollection, commonpb.MsgType_ReleaseCollection:
// All the request from query use collectionID
+129 -2
View File
@@ -23,7 +23,9 @@ import (
"strconv"
"strings"
"sync"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.uber.org/zap"
@@ -69,6 +71,9 @@ type Cache interface {
GetPartitionsIndex(ctx context.Context, database, collectionName string) ([]string, error)
// GetCollectionSchema get collection's schema.
GetCollectionSchema(ctx context.Context, database, collectionName string) (*schemaInfo, error)
// ResolveCollectionAlias returns the actual collection name if input is an alias,
// or returns the input name if it's already a collection name.
ResolveCollectionAlias(ctx context.Context, database, nameOrAlias string) (string, error)
// GetShard(ctx context.Context, withCache bool, database, collectionName string, collectionID int64, channel string) ([]nodeInfo, error)
// GetShardLeaderList(ctx context.Context, database, collectionName string, collectionID int64, withCache bool) ([]string, error)
// DeprecateShardCache(database, collectionName string)
@@ -87,6 +92,8 @@ type Cache interface {
// RefreshPolicyInfo(op typeutil.CacheOp) error
// InitPolicyInfo(info []string, userRoles []string)
// RemoveAlias removes a cached alias entry.
RemoveAlias(ctx context.Context, database, alias string)
RemoveDatabase(ctx context.Context, database string)
HasDatabase(ctx context.Context, database string) bool
GetDatabaseInfo(ctx context.Context, database string) (*databaseInfo, error)
@@ -112,6 +119,13 @@ type collectionInfo struct {
properties []*commonpb.KeyValuePair
}
const aliasCacheNegativeTTL = 30 * time.Second
type aliasEntry struct {
collectionName string // real collection name; "" means negative cache (not an alias)
cachedAt time.Time // when this entry was cached; used for TTL on negative entries
}
type databaseInfo struct {
dbID typeutil.UniqueID
properties []*commonpb.KeyValuePair
@@ -339,8 +353,9 @@ var _ Cache = (*MetaCache)(nil)
type MetaCache struct {
mixCoord types.MixCoordClient
dbInfo map[string]*databaseInfo // database -> db_info
collInfo map[string]map[string]*collectionInfo // database -> collectionName -> collection_info
dbInfo map[string]*databaseInfo // database -> db_info
collInfo map[string]map[string]*collectionInfo // database -> collectionName -> collection_info
aliasInfo map[string]map[string]*aliasEntry // database -> alias -> entry
credMap map[string]*internalpb.CredentialInfo // cache for credential, lazy load
privilegeInfos map[string]struct{} // privileges cache
@@ -389,6 +404,7 @@ func NewMetaCache(mixCoord types.MixCoordClient) (*MetaCache, error) {
mixCoord: mixCoord,
dbInfo: map[string]*databaseInfo{},
collInfo: map[string]map[string]*collectionInfo{},
aliasInfo: map[string]map[string]*aliasEntry{},
credMap: map[string]*internalpb.CredentialInfo{},
privilegeInfos: map[string]struct{}{},
userToRoles: map[string]map[string]struct{}{},
@@ -643,6 +659,104 @@ func (m *MetaCache) GetCollectionSchema(ctx context.Context, database, collectio
return collInfo.schema, nil
}
func (m *MetaCache) getAlias(database, alias string) (*aliasEntry, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
if db, ok := m.aliasInfo[database]; ok {
if entry, ok := db[alias]; ok {
// Expire negative cache entries after TTL
if entry.collectionName == "" && time.Since(entry.cachedAt) > aliasCacheNegativeTTL {
return nil, false
}
return entry, true
}
}
return nil, false
}
// setAliasLocked sets an alias cache entry. Caller must hold m.mu write lock.
func (m *MetaCache) setAliasLocked(database, alias string, entry *aliasEntry) {
if _, ok := m.aliasInfo[database]; !ok {
m.aliasInfo[database] = make(map[string]*aliasEntry)
}
m.aliasInfo[database][alias] = entry
}
// removeAliasLocked removes an alias cache entry. Caller must hold m.mu write lock.
func (m *MetaCache) removeAliasLocked(database, alias string) {
if db, ok := m.aliasInfo[database]; ok {
delete(db, alias)
}
}
// removeAliasesForCollectionLocked removes all positive alias entries pointing to collectionName.
// Caller must hold m.mu write lock.
func (m *MetaCache) removeAliasesForCollectionLocked(database, collectionName string) {
if db, ok := m.aliasInfo[database]; ok {
for alias, entry := range db {
if entry.collectionName == collectionName {
delete(db, alias)
}
}
}
}
func (m *MetaCache) RemoveAlias(ctx context.Context, database, alias string) {
m.mu.Lock()
defer m.mu.Unlock()
m.removeAliasLocked(database, alias)
log.Ctx(ctx).Debug("remove alias from cache", zap.String("db", database), zap.String("alias", alias))
}
func (m *MetaCache) ResolveCollectionAlias(ctx context.Context, database, nameOrAlias string) (string, error) {
// Level 1: Found in collection cache, return as-is
if _, ok := m.getCollection(database, nameOrAlias, 0); ok {
return nameOrAlias, nil
}
// Level 2: Found in alias cache
if entry, ok := m.getAlias(database, nameOrAlias); ok {
if entry.collectionName == "" {
// Negative cache: not an alias, return as-is
return nameOrAlias, nil
}
return entry.collectionName, nil
}
// Level 3: Cache miss, call DescribeAlias RPC
resp, err := m.mixCoord.DescribeAlias(ctx, &milvuspb.DescribeAliasRequest{
DbName: database,
Alias: nameOrAlias,
})
if err != nil {
return "", err
}
if err = merr.CheckRPCCall(resp, nil); err != nil {
if errors.Is(err, merr.ErrAliasNotFound) || errors.Is(err, merr.ErrCollectionNotFound) {
// Negative cache: this name is not an alias
m.mu.Lock()
m.setAliasLocked(database, nameOrAlias, &aliasEntry{collectionName: "", cachedAt: time.Now()})
m.mu.Unlock()
return nameOrAlias, nil
}
return "", err
}
if resp.GetCollection() == "" {
// Negative cache
m.mu.Lock()
m.setAliasLocked(database, nameOrAlias, &aliasEntry{collectionName: "", cachedAt: time.Now()})
m.mu.Unlock()
return nameOrAlias, nil
}
// Positive cache: alias -> real collection name
m.mu.Lock()
m.setAliasLocked(database, nameOrAlias, &aliasEntry{collectionName: resp.GetCollection(), cachedAt: time.Now()})
m.mu.Unlock()
return resp.GetCollection(), nil
}
func (m *MetaCache) GetPartitionID(ctx context.Context, database, collectionName string, partitionName string) (typeutil.UniqueID, error) {
partInfo, err := m.GetPartitionInfo(ctx, database, collectionName, partitionName)
if err != nil {
@@ -842,18 +956,29 @@ func (m *MetaCache) RemoveCollection(ctx context.Context, database, collectionNa
m.mu.Lock()
defer m.mu.Unlock()
found := false
if db, dbOk := m.collInfo[database]; dbOk {
if coll, ok := db[collectionName]; ok {
m.removeCollectionByID(ctx, coll.collID, version, false)
found = true
}
}
if database == "" {
if db, dbOk := m.collInfo[defaultDB]; dbOk {
if coll, ok := db[collectionName]; ok {
m.removeCollectionByID(ctx, coll.collID, version, false)
found = true
}
}
}
// If the collection was not in cache, alias entries pointing to it won't
// have been cleaned up by removeCollectionByID. Clean them up here.
if !found {
m.removeAliasesForCollectionLocked(database, collectionName)
if database == "" {
m.removeAliasesForCollectionLocked(defaultDB, collectionName)
}
}
log.Ctx(ctx).Debug("remove collection", zap.String("db", database), zap.String("collection", collectionName))
}
@@ -875,6 +1000,7 @@ func (m *MetaCache) removeCollectionByID(ctx context.Context, collectionID Uniqu
collNames = append(collNames, k)
m.sfGlobal.Forget(buildSfKeyByName(database, k))
m.sfGlobal.Forget(buildSfKeyById(database, v.collID))
m.removeAliasesForCollectionLocked(database, k)
}
}
}
@@ -896,6 +1022,7 @@ func (m *MetaCache) RemoveDatabase(ctx context.Context, database string) {
defer m.mu.Unlock()
delete(m.collInfo, database)
delete(m.dbInfo, database)
delete(m.aliasInfo, database)
}
func (m *MetaCache) HasDatabase(ctx context.Context, database string) bool {
+6 -1
View File
@@ -361,7 +361,12 @@ func (c *MockMixCoordClientInterface) AlterAlias(ctx context.Context, req *milvu
// DescribeAlias describe alias
func (c *MockMixCoordClientInterface) DescribeAlias(ctx context.Context, req *milvuspb.DescribeAliasRequest, opts ...grpc.CallOption) (*milvuspb.DescribeAliasResponse, error) {
panic("implement me")
return &milvuspb.DescribeAliasResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: "alias not found",
},
}, nil
}
// ListAliases list all aliases of db or collection
+1
View File
@@ -55,6 +55,7 @@ func InitEmptyGlobalCache() {
emptyMock := common.NewEmptyMockT()
mixcoord := mocks.NewMockMixCoordClient(emptyMock)
mixcoord.EXPECT().DescribeCollection(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("collection not found"))
mixcoord.EXPECT().DescribeAlias(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("alias not found"))
globalMetaCache, err = NewMetaCache(mixcoord)
if err != nil {
panic(err)
+93
View File
@@ -719,6 +719,41 @@ func (_c *MockCache_HasDatabase_Call) RunAndReturn(run func(context.Context, str
return _c
}
// RemoveAlias provides a mock function with given fields: ctx, database, alias
func (_m *MockCache) RemoveAlias(ctx context.Context, database string, alias string) {
_m.Called(ctx, database, alias)
}
// MockCache_RemoveAlias_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveAlias'
type MockCache_RemoveAlias_Call struct {
*mock.Call
}
// RemoveAlias is a helper method to define mock.On call
// - ctx context.Context
// - database string
// - alias string
func (_e *MockCache_Expecter) RemoveAlias(ctx interface{}, database interface{}, alias interface{}) *MockCache_RemoveAlias_Call {
return &MockCache_RemoveAlias_Call{Call: _e.mock.On("RemoveAlias", ctx, database, alias)}
}
func (_c *MockCache_RemoveAlias_Call) Run(run func(ctx context.Context, database string, alias string)) *MockCache_RemoveAlias_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string))
})
return _c
}
func (_c *MockCache_RemoveAlias_Call) Return() *MockCache_RemoveAlias_Call {
_c.Call.Return()
return _c
}
func (_c *MockCache_RemoveAlias_Call) RunAndReturn(run func(context.Context, string, string)) *MockCache_RemoveAlias_Call {
_c.Run(run)
return _c
}
// RemoveCollection provides a mock function with given fields: ctx, database, collectionName, version
func (_m *MockCache) RemoveCollection(ctx context.Context, database string, collectionName string, version uint64) {
_m.Called(ctx, database, collectionName, version)
@@ -840,6 +875,64 @@ func (_c *MockCache_RemoveDatabase_Call) RunAndReturn(run func(context.Context,
return _c
}
// ResolveCollectionAlias provides a mock function with given fields: ctx, database, nameOrAlias
func (_m *MockCache) ResolveCollectionAlias(ctx context.Context, database string, nameOrAlias string) (string, error) {
ret := _m.Called(ctx, database, nameOrAlias)
if len(ret) == 0 {
panic("no return value specified for ResolveCollectionAlias")
}
var r0 string
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string, string) (string, error)); ok {
return rf(ctx, database, nameOrAlias)
}
if rf, ok := ret.Get(0).(func(context.Context, string, string) string); ok {
r0 = rf(ctx, database, nameOrAlias)
} else {
r0 = ret.Get(0).(string)
}
if rf, ok := ret.Get(1).(func(context.Context, string, string) error); ok {
r1 = rf(ctx, database, nameOrAlias)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockCache_ResolveCollectionAlias_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ResolveCollectionAlias'
type MockCache_ResolveCollectionAlias_Call struct {
*mock.Call
}
// ResolveCollectionAlias is a helper method to define mock.On call
// - ctx context.Context
// - database string
// - nameOrAlias string
func (_e *MockCache_Expecter) ResolveCollectionAlias(ctx interface{}, database interface{}, nameOrAlias interface{}) *MockCache_ResolveCollectionAlias_Call {
return &MockCache_ResolveCollectionAlias_Call{Call: _e.mock.On("ResolveCollectionAlias", ctx, database, nameOrAlias)}
}
func (_c *MockCache_ResolveCollectionAlias_Call) Run(run func(ctx context.Context, database string, nameOrAlias string)) *MockCache_ResolveCollectionAlias_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string))
})
return _c
}
func (_c *MockCache_ResolveCollectionAlias_Call) Return(_a0 string, _a1 error) *MockCache_ResolveCollectionAlias_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockCache_ResolveCollectionAlias_Call) RunAndReturn(run func(context.Context, string, string) (string, error)) *MockCache_ResolveCollectionAlias_Call {
_c.Call.Return(run)
return _c
}
// NewMockCache creates a new instance of MockCache. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockCache(t interface {
+44 -3
View File
@@ -6,7 +6,6 @@ import (
"reflect"
"sync"
"github.com/casbin/casbin/v2"
"go.uber.org/zap"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
@@ -19,12 +18,12 @@ import (
"github.com/milvus-io/milvus/pkg/v2/util"
"github.com/milvus-io/milvus/pkg/v2/util/contextutil"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
)
type PrivilegeFunc func(ctx context.Context, req interface{}) (context.Context, error)
var (
enforcer *casbin.SyncedEnforcer
initOnce sync.Once
initPrivilegeGroupsOnce sync.Once
)
@@ -71,6 +70,20 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
objectType := privilegeExt.ObjectType.String()
objectNameIndex := privilegeExt.ObjectNameIndex
objectName := funcutil.GetObjectName(req, objectNameIndex)
dbName := GetCurDBNameFromContextOrDefault(ctx)
// Resolve alias to actual collection name for RBAC checks
if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && objectType == commonpb.ObjectType_Collection.String() && objectNameIndex != 0 {
if objectName != util.AnyWord && objectName != "" {
if actualCollectionName, resolveErr := resolveCollectionAlias(ctx, dbName, objectName); resolveErr != nil {
log.RatedWarn(60, "failed to resolve collection alias for RBAC, using original name",
zap.String("objectName", objectName), zap.String("dbName", dbName), zap.Error(resolveErr))
} else {
objectName = actualCollectionName
}
}
}
if isCurUserObject(objectType, username, objectName) {
return ctx, nil
}
@@ -81,8 +94,27 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
objectNameIndexs := privilegeExt.ObjectNameIndexs
objectNames := funcutil.GetObjectNames(req, objectNameIndexs)
// Resolve aliases for operations that refer to multiple resources
if Params.ProxyCfg.ResolveAliasForPrivilege.GetAsBool() && objectType == commonpb.ObjectType_Collection.String() && objectNameIndexs != 0 && len(objectNames) > 0 {
resolvedNames := make([]string, 0, len(objectNames))
for _, name := range objectNames {
if name == util.AnyWord || name == "" {
resolvedNames = append(resolvedNames, name)
continue
}
if actualName, resolveErr := resolveCollectionAlias(ctx, dbName, name); resolveErr != nil {
log.RatedWarn(60, "failed to resolve collection alias for RBAC, using original name",
zap.String("objectName", name), zap.String("dbName", dbName), zap.Error(resolveErr))
resolvedNames = append(resolvedNames, name)
} else {
resolvedNames = append(resolvedNames, actualName)
}
}
objectNames = resolvedNames
}
objectPrivilege := privilegeExt.ObjectPrivilege.String()
dbName := GetCurDBNameFromContextOrDefault(ctx)
log = log.With(zap.String("username", username), zap.Strings("role_names", roleNames),
zap.String("object_type", objectType), zap.String("object_privilege", objectPrivilege),
@@ -167,3 +199,12 @@ func isSelectMyRoleGrants(req interface{}, roleNames []string) bool {
roleName := filterGrantEntity.GetRole().GetName()
return funcutil.SliceContain(roleNames, roleName)
}
// resolveCollectionAlias resolves an alias to its actual collection name
func resolveCollectionAlias(ctx context.Context, dbName, nameOrAlias string) (string, error) {
cache := globalMetaCache
if cache == nil {
return nameOrAlias, merr.WrapErrServiceInternal("meta cache not initialized")
}
return cache.ResolveCollectionAlias(ctx, dbName, nameOrAlias)
}
+21
View File
@@ -99,6 +99,14 @@ func (t *CreateAliasTask) PreExecute(ctx context.Context) error {
if err := validateCollectionName(collName); err != nil {
return err
}
// Always resolve CollectionName in case the user passed an alias instead of the real name.
// This is needed for correctness: rootcoord's CheckIfAliasCreatable only checks mt.names,
// so passing an alias as CollectionName would fail.
dbName := t.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, dbName, collName); err == nil {
t.CollectionName = resolved
}
return nil
}
@@ -255,6 +263,13 @@ func (t *AlterAliasTask) PreExecute(ctx context.Context) error {
return err
}
// Always resolve CollectionName in case the user passed an alias instead of the real name.
// This is needed for correctness: rootcoord's CheckIfAliasAlterable only checks mt.names,
// so passing an alias as CollectionName would fail.
dbName := t.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, dbName, collName); err == nil {
t.CollectionName = resolved
}
return nil
}
@@ -394,6 +409,12 @@ func (a *ListAliasesTask) PreExecute(ctx context.Context) error {
if err := validateCollectionName(a.GetCollectionName()); err != nil {
return err
}
// Resolve CollectionName in case the user passed an alias instead of the real name.
// rootcoord filters by collection name in mt.names, so an alias would return no results.
dbName := a.GetDbName()
if resolved, err := resolveCollectionAlias(ctx, dbName, a.GetCollectionName()); err == nil {
a.CollectionName = resolved
}
}
return nil
}
+14
View File
@@ -35,6 +35,13 @@ func TestCreateAlias_all(t *testing.T) {
defer rc.Close()
ctx := context.Background()
// Save and clear globalMetaCache to avoid resolveCollectionAlias calling
// DescribeAlias on a stale testify mock from a previous test.
oldCache := globalMetaCache
globalMetaCache = nil
defer func() { globalMetaCache = oldCache }()
prefix := "TestCreateAlias_all"
collectionName := prefix + funcutil.GenRandomStr()
task := &CreateAliasTask{
@@ -115,6 +122,13 @@ func TestAlterAlias_all(t *testing.T) {
rc := NewMixCoordMock()
defer rc.Close()
ctx := context.Background()
// Save and clear globalMetaCache to avoid resolveCollectionAlias calling
// DescribeAlias on a stale testify mock from a previous test.
oldCache := globalMetaCache
globalMetaCache = nil
defer func() { globalMetaCache = oldCache }()
prefix := "TestAlterAlias_all"
collectionName := prefix + funcutil.GenRandomStr()
task := &AlterAliasTask{
@@ -5,6 +5,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.uber.org/zap"
"google.golang.org/protobuf/types/known/fieldmaskpb"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@@ -16,6 +17,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v2/proto/proxypb"
"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message/ce"
@@ -339,5 +341,20 @@ func (c *DDLCallback) alterCollectionV2AckCallback(ctx context.Context, result m
if err := c.broker.BroadcastAlteredCollection(ctx, header.CollectionId); err != nil {
return errors.Wrap(err, "failed to broadcast altered collection")
}
// If the collection was renamed or moved to a different DB, grants were migrated
// in MetaTable.AlterCollection. Refresh the RBAC policy cache on all proxies so
// they pick up the new grant keys.
for _, path := range header.UpdateMask.GetPaths() {
if path == message.FieldMaskCollectionName || path == message.FieldMaskDB {
if err := c.proxyClientManager.RefreshPolicyInfoCache(ctx, &proxypb.RefreshPolicyInfoCacheRequest{
OpType: int32(typeutil.CacheRefresh),
}); err != nil {
log.Ctx(ctx).Warn("failed to refresh RBAC policy cache after collection rename, skipping", zap.Error(err))
}
break
}
}
return c.ExpireCaches(ctx, header)
}
@@ -21,19 +21,24 @@ import (
"fmt"
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/registry"
"github.com/milvus-io/milvus/internal/util/proxyutil"
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
"github.com/milvus-io/milvus/pkg/v2/proto/proxypb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message/ce"
"github.com/milvus-io/milvus/pkg/v2/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
func (c *Core) broadcastDropCollectionV1(ctx context.Context, req *milvuspb.DropCollectionRequest) error {
@@ -125,7 +130,7 @@ func (c *DDLCallback) dropCollectionV1AckCallback(ctx context.Context, result me
}
}
// add the collection tombstone to the sweeper.
c.tombstoneSweeper.AddTombstone(newCollectionTombstone(c.meta, c.broker, header.CollectionId))
c.tombstoneSweeper.AddTombstone(newCollectionTombstone(c.meta, c.broker, header.CollectionId, c.proxyClientManager))
// expire the collection meta cache on proxy.
return c.ExpireCaches(ctx, ce.NewBuilder().WithLegacyProxyCollectionMetaCache(
ce.OptLPCMDBName(body.DbName),
@@ -135,18 +140,20 @@ func (c *DDLCallback) dropCollectionV1AckCallback(ctx context.Context, result me
}
// newCollectionTombstone creates a new collection tombstone.
func newCollectionTombstone(meta IMetaTable, broker Broker, collectionID int64) *collectionTombstone {
func newCollectionTombstone(meta IMetaTable, broker Broker, collectionID int64, proxyClientManager proxyutil.ProxyClientManagerInterface) *collectionTombstone {
return &collectionTombstone{
meta: meta,
broker: broker,
collectionID: collectionID,
meta: meta,
broker: broker,
collectionID: collectionID,
proxyClientManager: proxyClientManager,
}
}
type collectionTombstone struct {
meta IMetaTable
broker Broker
collectionID int64
meta IMetaTable
broker Broker
collectionID int64
proxyClientManager proxyutil.ProxyClientManagerInterface
}
func (t *collectionTombstone) ID() string {
@@ -158,5 +165,16 @@ func (t *collectionTombstone) ConfirmCanBeRemoved(ctx context.Context) (bool, er
}
func (t *collectionTombstone) Remove(ctx context.Context) error {
return t.meta.RemoveCollection(ctx, t.collectionID, 0)
if err := t.meta.RemoveCollection(ctx, t.collectionID, 0); err != nil {
return err
}
// RemoveCollection deletes grants associated with the dropped collection.
// Refresh the RBAC policy cache on all proxies so they stop using stale grant entries.
if err := t.proxyClientManager.RefreshPolicyInfoCache(ctx, &proxypb.RefreshPolicyInfoCacheRequest{
OpType: int32(typeutil.CacheRefresh),
}); err != nil {
log.Ctx(ctx).Warn("failed to refresh RBAC policy cache after collection removal, skipping",
zap.Int64("collectionID", t.collectionID), zap.Error(err))
}
return nil
}
+15
View File
@@ -698,6 +698,12 @@ func (mt *MetaTable) RemoveCollection(ctx context.Context, collectionID UniqueID
return err
}
tenant := Params.CommonCfg.ClusterName.GetValue()
if err := mt.catalog.DeleteGrantByCollectionName(ctx1, tenant, coll.DBName, coll.Name); err != nil {
log.Ctx(ctx).Warn("failed to delete grants for dropped collection, skipping",
zap.String("dbName", coll.DBName), zap.String("collectionName", coll.Name), zap.Error(err))
}
allNames := common.CloneStringList(aliases)
allNames = append(allNames, coll.Name)
@@ -1033,6 +1039,15 @@ func (mt *MetaTable) AlterCollection(ctx context.Context, result message.Broadca
}
}
if oldColl.Name != newColl.Name || oldColl.DBName != newColl.DBName {
tenant := Params.CommonCfg.ClusterName.GetValue()
if err := mt.catalog.MigrateGrantCollectionName(ctx1, tenant, oldColl.DBName, oldColl.Name, newColl.DBName, newColl.Name); err != nil {
log.Ctx(ctx).Warn("failed to migrate grants for renamed collection, skipping",
zap.String("oldDBName", oldColl.DBName), zap.String("oldName", oldColl.Name),
zap.String("newDBName", newColl.DBName), zap.String("newName", newColl.Name), zap.Error(err))
}
}
mt.names.remove(oldColl.DBName, oldColl.Name)
mt.names.insert(newColl.DBName, newColl.Name, newColl.CollectionID)
mt.collID2Meta[header.CollectionId] = newColl
+31
View File
@@ -1194,6 +1194,9 @@ func TestMetaTable_RemoveCollection(t *testing.T) {
mock.Anything, // model.Collection
mock.AnythingOfType("uint64"),
).Return(nil)
catalog.On("DeleteGrantByCollectionName",
mock.Anything, mock.Anything, mock.Anything, mock.Anything,
).Return(nil)
meta := &MetaTable{
catalog: catalog,
names: newNameDb(),
@@ -1213,6 +1216,34 @@ func TestMetaTable_RemoveCollection(t *testing.T) {
})
}
func TestMetaTable_RemoveCollection_GrantDeleteBestEffort(t *testing.T) {
// When DeleteGrantByCollectionName fails, RemoveCollection should still succeed (best-effort)
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropCollection",
mock.Anything,
mock.Anything,
mock.AnythingOfType("uint64"),
).Return(nil)
catalog.On("DeleteGrantByCollectionName",
mock.Anything, mock.Anything, mock.Anything, mock.Anything,
).Return(errors.New("grant delete failed"))
meta := &MetaTable{
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {Name: "collection", State: pb.CollectionState_CollectionDropping},
},
}
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
meta.names.insert("", "collection", 100)
ctx := context.Background()
err := meta.RemoveCollection(ctx, 100, 9999)
assert.NoError(t, err)
}
func TestMetaTable_RemovePartition(t *testing.T) {
t.Run("catalog error", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
+1 -1
View File
@@ -661,7 +661,7 @@ func (c *Core) restore(ctx context.Context) error {
// CollectionCreating is a deprecated status,
// we cannot promise the coordinator handle it correctly, so just treat it as a tombstone.
if coll.State == pb.CollectionState_CollectionDropping || coll.State == pb.CollectionState_CollectionCreating {
c.tombstoneSweeper.AddTombstone(newCollectionTombstone(c.meta, c.broker, coll.CollectionID))
c.tombstoneSweeper.AddTombstone(newCollectionTombstone(c.meta, c.broker, coll.CollectionID, c.proxyClientManager))
continue
}
for _, part := range coll.Partitions {
+10
View File
@@ -1928,6 +1928,7 @@ type proxyConfig struct {
MustUsePartitionKey ParamItem `refreshable:"true"`
SkipAutoIDCheck ParamItem `refreshable:"true"`
SkipPartitionKeyCheck ParamItem `refreshable:"true"`
ResolveAliasForPrivilege ParamItem `refreshable:"true"`
MaxVarCharLength ParamItem `refreshable:"false"`
MaxTextLength ParamItem `refreshable:"false"`
MaxResultEntries ParamItem `refreshable:"true"`
@@ -2370,6 +2371,15 @@ please adjust in embedded Milvus: false`,
}
p.SkipPartitionKeyCheck.Init(base.mgr)
p.ResolveAliasForPrivilege = ParamItem{
Key: "proxy.resolveAliasForPrivilege",
Version: "2.6.9",
DefaultValue: "false",
Doc: "switch for whether proxy shall resolve alias to actual collection name during RBAC privilege checks",
Export: true,
}
p.ResolveAliasForPrivilege.Init(base.mgr)
p.MaxVarCharLength = ParamItem{
Key: "proxy.maxVarCharLength",
Version: "2.4.19", // hotfix
+125
View File
@@ -0,0 +1,125 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package rbac
import (
"context"
"fmt"
"time"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/pkg/v2/util"
"github.com/milvus-io/milvus/pkg/v2/util/crypto"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/tests/integration"
)
// TestAliasRBAC verifies that the RBAC interceptor resolves aliases to real
// collection names before checking permissions. When a user operates on a
// collection through an alias, the permission check must be against the real
// collection name, not the alias.
func (s *RBACBasicTestSuite) TestAliasRBAC() {
rootCtx := GetContext(context.Background(), defaultAuth)
realColName := "real_collection_rbac"
aliasName := "alias_collection_rbac"
roleName := "alias_role"
userName := "alias_user"
userPassword := "alias_passwd"
// Setup: create collection, alias, role, user
schema := integration.ConstructSchema(realColName, dim, true)
marshaledSchema, err := proto.Marshal(schema)
s.NoError(err)
createCollResp, err := s.Cluster.MilvusClient.CreateCollection(rootCtx, &milvuspb.CreateCollectionRequest{
CollectionName: realColName,
Schema: marshaledSchema,
})
s.NoError(err)
s.True(merr.Ok(createCollResp))
createAliasResp, err := s.Cluster.MilvusClient.CreateAlias(rootCtx, &milvuspb.CreateAliasRequest{
CollectionName: realColName,
Alias: aliasName,
})
s.NoError(err)
s.True(merr.Ok(createAliasResp))
resp, err := s.Cluster.MilvusClient.CreateRole(rootCtx, &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{Name: roleName},
})
s.NoError(err)
s.True(merr.Ok(resp))
resp, err = s.Cluster.MilvusClient.CreateCredential(rootCtx, &milvuspb.CreateCredentialRequest{
Username: userName,
Password: crypto.Base64Encode(userPassword),
})
s.NoError(err)
s.True(merr.Ok(resp))
resp, err = s.Cluster.MilvusClient.OperateUserRole(rootCtx, &milvuspb.OperateUserRoleRequest{
Username: userName,
RoleName: roleName,
Type: milvuspb.OperateUserRoleType_AddUserToRole,
})
s.NoError(err)
s.True(merr.Ok(resp))
defer func() {
s.Cluster.MilvusClient.DropAlias(rootCtx, &milvuspb.DropAliasRequest{Alias: aliasName}) //nolint
s.Cluster.MilvusClient.DropCollection(rootCtx, &milvuspb.DropCollectionRequest{CollectionName: realColName}) //nolint
s.Cluster.MilvusClient.DropRole(rootCtx, &milvuspb.DropRoleRequest{RoleName: roleName, ForceDrop: true}) //nolint
s.Cluster.MilvusClient.DeleteCredential(rootCtx, &milvuspb.DeleteCredentialRequest{Username: userName}) //nolint
}()
userCtx := GetContext(context.Background(), fmt.Sprintf("%s:%s", userName, userPassword))
// Without GetStatistics privilege, GetCollectionStatistics via alias should fail.
// The privilege interceptor resolves the alias to the real collection name,
// so using an alias does NOT bypass RBAC.
_, err = s.Cluster.MilvusClient.GetCollectionStatistics(userCtx, &milvuspb.GetCollectionStatisticsRequest{
CollectionName: aliasName,
})
s.Error(err, "GetCollectionStatistics via alias should fail without permission")
// Grant GetStatistics on the REAL collection name (not alias)
grantResp, err := s.Cluster.MilvusClient.OperatePrivilegeV2(rootCtx, &milvuspb.OperatePrivilegeV2Request{
Role: &milvuspb.RoleEntity{Name: roleName},
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: "GetStatistics"},
},
Type: milvuspb.OperatePrivilegeType_Grant,
DbName: util.AnyWord,
CollectionName: realColName,
})
s.NoError(err)
s.True(merr.Ok(grantResp))
// Wait for privilege cache to refresh
time.Sleep(3 * time.Second)
// Now GetCollectionStatistics via alias should succeed because the interceptor
// resolves the alias to the real collection name before checking permission.
statsResp, err := s.Cluster.MilvusClient.GetCollectionStatistics(userCtx, &milvuspb.GetCollectionStatisticsRequest{
CollectionName: aliasName,
})
s.NoError(err)
s.True(merr.Ok(statsResp.GetStatus()), "GetCollectionStatistics via alias should succeed when user has permission on real collection")
}
@@ -41,6 +41,7 @@ func (s *RBACBasicTestSuite) SetupSuite() {
s.WithMilvusConfig(paramtable.Get().QueryCoordCfg.BalanceCheckInterval.Key, "1000")
s.WithMilvusConfig(paramtable.Get().QueryNodeCfg.GracefulStopTimeout.Key, "1")
s.WithMilvusConfig(paramtable.Get().CommonCfg.AuthorizationEnabled.Key, "true")
s.WithMilvusConfig(paramtable.Get().ProxyCfg.ResolveAliasForPrivilege.Key, "true")
s.MiniClusterSuite.SetupSuite()
}
@@ -0,0 +1,185 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package rbac
import (
"context"
"time"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/pkg/v2/util"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/tests/integration"
)
// TestDropCollectionCleansGrants verifies that dropping a collection removes
// its associated RBAC grant entries, so re-creating a collection with the same
// name does NOT inherit stale permissions.
func (s *RBACBasicTestSuite) TestDropCollectionCleansGrants() {
ctx := GetContext(context.Background(), defaultAuth)
colName := "grant_cleanup_drop_col"
roleName := "grant_cleanup_drop_role"
// Create collection
schema := integration.ConstructSchema(colName, dim, true)
marshaledSchema, err := proto.Marshal(schema)
s.NoError(err)
createResp, err := s.Cluster.MilvusClient.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
CollectionName: colName,
Schema: marshaledSchema,
})
s.NoError(err)
s.True(merr.Ok(createResp))
// Create role and grant a privilege on the collection
resp, err := s.Cluster.MilvusClient.CreateRole(ctx, &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{Name: roleName},
})
s.NoError(err)
s.True(merr.Ok(resp))
grantResp, err := s.Cluster.MilvusClient.OperatePrivilegeV2(ctx, &milvuspb.OperatePrivilegeV2Request{
Role: &milvuspb.RoleEntity{Name: roleName},
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: "GetStatistics"},
},
Type: milvuspb.OperatePrivilegeType_Grant,
DbName: util.AnyWord,
CollectionName: colName,
})
s.NoError(err)
s.True(merr.Ok(grantResp))
// Verify grant exists
selectResp, err := s.Cluster.MilvusClient.SelectGrant(ctx, &milvuspb.SelectGrantRequest{
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: colName,
DbName: util.AnyWord,
},
})
s.NoError(err)
s.True(merr.Ok(selectResp.GetStatus()))
s.NotEmpty(selectResp.GetEntities(), "grant should exist before drop")
// Drop collection
dropResp, err := s.Cluster.MilvusClient.DropCollection(ctx, &milvuspb.DropCollectionRequest{
CollectionName: colName,
})
s.NoError(err)
s.True(merr.Ok(dropResp))
time.Sleep(2 * time.Second)
// Verify grant is cleaned up
selectResp, err = s.Cluster.MilvusClient.SelectGrant(ctx, &milvuspb.SelectGrantRequest{
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: colName,
DbName: util.AnyWord,
},
})
s.NoError(err)
s.True(merr.Ok(selectResp.GetStatus()))
s.Empty(selectResp.GetEntities(), "grants should be cleaned up after drop collection")
// Cleanup
s.Cluster.MilvusClient.DropRole(ctx, &milvuspb.DropRoleRequest{RoleName: roleName, ForceDrop: true}) //nolint
}
// TestRenameCollectionMigratesGrants verifies that renaming a collection
// migrates its RBAC grant entries to the new name.
func (s *RBACBasicTestSuite) TestRenameCollectionMigratesGrants() {
ctx := GetContext(context.Background(), defaultAuth)
oldName := "grant_migrate_old_col"
newName := "grant_migrate_new_col"
roleName := "grant_migrate_role"
// Create collection
schema := integration.ConstructSchema(oldName, dim, true)
marshaledSchema, err := proto.Marshal(schema)
s.NoError(err)
createResp, err := s.Cluster.MilvusClient.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
CollectionName: oldName,
Schema: marshaledSchema,
})
s.NoError(err)
s.True(merr.Ok(createResp))
// Create role and grant a privilege on the collection
resp, err := s.Cluster.MilvusClient.CreateRole(ctx, &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{Name: roleName},
})
s.NoError(err)
s.True(merr.Ok(resp))
grantResp, err := s.Cluster.MilvusClient.OperatePrivilegeV2(ctx, &milvuspb.OperatePrivilegeV2Request{
Role: &milvuspb.RoleEntity{Name: roleName},
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: "GetStatistics"},
},
Type: milvuspb.OperatePrivilegeType_Grant,
DbName: util.AnyWord,
CollectionName: oldName,
})
s.NoError(err)
s.True(merr.Ok(grantResp))
// Rename collection
renameResp, err := s.Cluster.MilvusClient.RenameCollection(ctx, &milvuspb.RenameCollectionRequest{
OldName: oldName,
NewName: newName,
})
s.NoError(err)
s.True(merr.Ok(renameResp))
time.Sleep(2 * time.Second)
// Verify grant no longer exists under old name
selectResp, err := s.Cluster.MilvusClient.SelectGrant(ctx, &milvuspb.SelectGrantRequest{
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: oldName,
DbName: util.AnyWord,
},
})
s.NoError(err)
s.True(merr.Ok(selectResp.GetStatus()))
s.Empty(selectResp.GetEntities(), "grants should not exist under old name after rename")
// Verify grant now exists under new name
selectResp, err = s.Cluster.MilvusClient.SelectGrant(ctx, &milvuspb.SelectGrantRequest{
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: newName,
DbName: util.AnyWord,
},
})
s.NoError(err)
s.True(merr.Ok(selectResp.GetStatus()))
s.NotEmpty(selectResp.GetEntities(), "grants should be migrated to new name after rename")
// Cleanup
s.Cluster.MilvusClient.DropCollection(ctx, &milvuspb.DropCollectionRequest{CollectionName: newName}) //nolint
s.Cluster.MilvusClient.DropRole(ctx, &milvuspb.DropRoleRequest{RoleName: roleName, ForceDrop: true}) //nolint
}