mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
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:
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user