milvus-io/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-10-11 22:49:53 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
87bdd624027d43a4da1d2d00de5e63a959b5f17c
milvus/internal/rootcoord/meta_table_test.go
T
aoiasd 87bdd62402 feat: [RLS8] decouple RLS metadata from enforcement (#54059) 
relate: #50263
design doc: docs/design-docs/design_docs/20250610-rls_design.md
design doc PR: #53173

## Summary

Allow changing `rls.enabled` independently of policy/tag management.
Preserve metadata across toggles and permit configuration before
enabling enforcement; keep schema dependency checks and explicit drop
cleanup.

Recover all saved policies into resident Coord metadata before
readiness, including disabled collections, and maintain the map after
successful persistence. Principal tags remain persisted in etcd and
point-read by Coord without startup preload; Proxy execution caches are
unchanged.

Reuse existing collection-cache invalidation, continue invalidation
after an ignored missing-resource-group result, and simplify redundant
request/enum adapters without removing supported protocol modes.

## Validation

Two fresh independent static reviewers completed the full eligible
production diff at `96688e8484`; primary reconciliation found no new
actionable finding. Static failure-path checks cover recovery failure,
persisted mutation failure, metadata preservation across toggles and
cache invalidation retries. Regression coverage was updated but has not
been executed for this revision.

No new local build or test run has been performed for this revision.
Earlier build and test results apply to earlier commits only; current
runtime validation remains with PR CI. Static review excludes
tests/generated contents and does not verify distributed execution.
Distributed cutover is not claimed to be atomic.

---------

Signed-off-by: aoiasd <[email protected]>
2026-10-10 13:45:39 +08:00

5275 lines
178 KiB
Go
Raw Blame History

// 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 rootcoord
import (
"context"
"math/rand"
"strconv"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/types/known/fieldmaskpb"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/rootcoord"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
mocktso "github.com/milvus-io/milvus/internal/tso/mocks"
kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
"github.com/milvus-io/milvus/internal/util/hookutil"
"github.com/milvus-io/milvus/internal/util/rlsutil"
"github.com/milvus-io/milvus/pkg/v3/common"
pb "github.com/milvus-io/milvus/pkg/v3/proto/etcdpb"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/util"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
func TestMetaTable_DescribeAliasAllowsConcurrentReaders(t *testing.T) {
const (
collectionID = int64(100)
collectionName = "test_metatable_describe_alias"
aliasName = "a_alias"
)
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID: {
CollectionID: collectionID,
Name: collectionName,
},
},
aliases: newNameDb(),
}
meta.aliases.insert(util.DefaultDBName, aliasName, collectionID)
type result struct {
collectionName string
err error
}
resultCh := make(chan result, 1)
meta.ddLock.RLock()
go func() {
collectionName, err := meta.DescribeAlias(context.Background(), util.DefaultDBName, aliasName, 0)
resultCh <- result{collectionName: collectionName, err: err}
}()
select {
case result := <-resultCh:
meta.ddLock.RUnlock()
require.NoError(t, result.err)
assert.Equal(t, collectionName, result.collectionName)
case <-time.After(3 * time.Second):
meta.ddLock.RUnlock()
<-resultCh
t.Fatal("DescribeAlias blocked behind another reader")
}
}
func generateMetaTable(_ *testing.T) *MetaTable {
kv, _ := kvfactory.GetEtcdAndPath()
path := funcutil.RandomString(10)
catalogKV := etcdkv.NewEtcdKV(kv, path)
return &MetaTable{catalog: rootcoord.NewCatalog(catalogKV)}
}
func newRLSMetaTableForTest(t *testing.T) (*MetaTable, *mocks.RootCoordCatalog) {
catalog := mocks.NewRootCoordCatalog(t)
names := newNameDb()
names.insert("db1", "coll1", 20)
return &MetaTable{
ctx: context.Background(),
catalog: catalog,
names: names,
aliases: newNameDb(),
dbName2Meta: map[string]*model.Database{
"db1": model.NewDatabase(10, "db1", pb.DatabaseState_DatabaseCreated, nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
20: {
DBID: 10,
CollectionID: 20,
DBName: "db1",
Name: "coll1",
State: pb.CollectionState_CollectionCreated,
UpdateTimestamp: 1,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "true"},
},
Fields: []*model.Field{
{
FieldID: 100,
Name: "id",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 101,
Name: "dept",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "128"},
},
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 102,
Name: "owner",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "128"},
},
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 103,
Name: "age",
DataType: schemapb.DataType_Int64,
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 104,
Name: "tags",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 105,
Name: "scores",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 106,
Name: "metadata",
DataType: schemapb.DataType_JSON,
State: schemapb.FieldState_FieldCreated,
},
{
FieldID: 107,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4"},
},
State: schemapb.FieldState_FieldCreated,
},
},
Partitions: []*model.Partition{
{PartitionID: 1, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
},
},
}, catalog
}
// These helpers exercise the prepare and callback-apply phases together. The
// production API persists the prepared post-image through a broadcast message
// before invoking the apply phase.
func (mt *MetaTable) CreateRLSPolicy(ctx context.Context, req *rlsutil.CreateRowPolicyRequest, policyID int64) (int64, error) {
policy, err := mt.PrepareCreateRLSPolicy(ctx, req, policyID)
if err != nil {
return 0, err
}
if err := mt.ApplyAlterRLSPolicy(ctx, policy); err != nil {
return 0, err
}
return policy.CollectionID, nil
}
func (mt *MetaTable) UpdateRLSPolicy(ctx context.Context, req *rlsutil.UpdateRowPolicyRequest) (int64, error) {
policy, err := mt.PrepareUpdateRLSPolicy(ctx, req)
if err != nil {
return 0, err
}
if err := mt.ApplyAlterRLSPolicy(ctx, policy); err != nil {
return 0, err
}
return policy.CollectionID, nil
}
func (mt *MetaTable) DropRLSPolicy(ctx context.Context, req *rlsutil.DropRowPolicyRequest) (int64, error) {
policy, err := mt.PrepareDropRLSPolicy(ctx, req)
if err != nil {
return 0, err
}
if err := mt.ApplyDropRLSPolicy(ctx, policy.CollectionID, policy.PolicyName); err != nil {
return 0, err
}
return policy.CollectionID, nil
}
func (mt *MetaTable) SetRLSPrincipalTags(ctx context.Context, req *rlsutil.SetRLSPrincipalTagsRequest) (int64, error) {
principal, err := mt.PrepareSetRLSPrincipalTags(ctx, req)
if err != nil {
return 0, err
}
if err := mt.ApplyAlterRLSPrincipal(ctx, principal); err != nil {
return 0, err
}
return principal.CollectionID, nil
}
func (mt *MetaTable) DeleteRLSPrincipalTags(ctx context.Context, req *rlsutil.DeleteRLSPrincipalTagsRequest) (int64, error) {
principal, drop, err := mt.PrepareDeleteRLSPrincipalTags(ctx, req)
if err != nil {
return 0, err
}
if drop {
err = mt.ApplyDropRLSPrincipal(ctx, principal.CollectionID, principal.PrincipalName)
} else {
err = mt.ApplyAlterRLSPrincipal(ctx, principal)
}
if err != nil {
return 0, err
}
return principal.CollectionID, nil
}
func TestReloadCollectionRLSMetadataLoadsPoliciesOnly(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{catalog: catalog}
collection := &model.Collection{
DBID: 11,
CollectionID: 20,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "true"},
},
}
stored := &model.RLSPolicy{
DBID: 10, CollectionID: 20, PolicyID: 100, PolicyName: "tenant",
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
}
catalog.EXPECT().ListRLSPolicies(mock.Anything, int64(20)).Return([]*model.RLSPolicy{nil, stored}, nil).Once()
require.NoError(t, meta.reloadCollectionRLSMetadata(context.Background(), collection))
require.Len(t, collection.RLSPolicies, 1)
require.Equal(t, int64(11), collection.RLSPolicies["tenant"].DBID)
require.Equal(t, int64(10), stored.DBID)
collection.RLSPolicies["tenant"].Actions[0] = rlsutil.PolicyActionInsert
require.Equal(t, rlsutil.PolicyActionQuery, stored.Actions[0])
}
func TestReloadRLSMetadataRejectsDuplicatePolicyNames(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{catalog: catalog}
existing := &model.RLSPolicy{CollectionID: 20, PolicyID: 99, PolicyName: "existing"}
collection := &model.Collection{CollectionID: 20, RLSPolicies: map[string]*model.RLSPolicy{"existing": existing}}
catalog.EXPECT().ListRLSPolicies(mock.Anything, int64(20)).Return([]*model.RLSPolicy{
{CollectionID: 20, PolicyID: 100, PolicyName: "tenant"},
{CollectionID: 20, PolicyID: 101, PolicyName: "tenant"},
}, nil).Once()
require.ErrorIs(t, meta.reloadCollectionRLSMetadata(context.Background(), collection), merr.ErrDataIntegrity)
require.Len(t, collection.RLSPolicies, 1)
require.Same(t, existing, collection.RLSPolicies["existing"])
}
func TestReloadCollectionsRLSMetadataRestoresDisabledCollection(t *testing.T) {
for _, tc := range []struct {
name string
properties []*commonpb.KeyValuePair
}{
{name: "property absent"},
{
name: "explicitly disabled",
properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "false"},
},
},
} {
t.Run(tc.name, func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
collection := meta.collID2Meta[20]
collection.Properties = tc.properties
collection.RLSPolicies = map[string]*model.RLSPolicy{"stale": {PolicyName: "stale"}}
catalog.EXPECT().ListRLSPolicies(mock.Anything, int64(20)).Return([]*model.RLSPolicy{{
CollectionID: 20, PolicyID: 100, PolicyName: "tenant", UsingExpr: `dept == "sales"`,
}}, nil).Once()
require.NoError(t, meta.reloadCollectionsRLSMetadata(context.Background(), []*model.Collection{collection}))
require.Len(t, collection.RLSPolicies, 1)
require.Equal(t, int64(100), collection.RLSPolicies["tenant"].PolicyID)
require.ErrorContains(t, validateRLSNoReferencedFieldDropped(collection, []int64{101}), "tenant")
require.NoError(t, validateRLSNoReferencedFieldDropped(collection, []int64{102}))
})
}
}
func TestReloadCollectionsRLSMetadataRejectsInvalidProperty(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{catalog: catalog}
collection := &model.Collection{
CollectionID: 20,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "invalid"},
},
}
err := meta.reloadCollectionsRLSMetadata(context.Background(), []*model.Collection{collection})
require.ErrorIs(t, err, merr.ErrDataIntegrity)
}
func TestReloadDisabledCollectionRLSMetadataPropagatesCatalogError(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
collection := meta.collID2Meta[20]
collection.Properties = nil
policy := &model.RLSPolicy{PolicyName: "tenant", PolicyID: 100}
collection.RLSPolicies = map[string]*model.RLSPolicy{"tenant": policy}
catalog.EXPECT().ListRLSPolicies(mock.Anything, int64(20)).Return(
nil, merr.WrapErrServiceUnavailableMsg("catalog unavailable")).Once()
err := meta.reloadCollectionsRLSMetadata(context.Background(), []*model.Collection{collection})
require.ErrorIs(t, err, merr.ErrServiceUnavailable)
require.Same(t, policy, collection.RLSPolicies["tenant"])
}
func TestRLSPoliciesStayResidentWhileTagsLoadOnDemand(t *testing.T) {
ctx := context.Background()
meta, catalog := newRLSMetaTableForTest(t)
collection := meta.collID2Meta[20]
collection.Properties = nil
policy := &model.RLSPolicy{DBID: 10, CollectionID: 20, PolicyID: 100, PolicyName: "tenant"}
catalog.EXPECT().SaveRLSPolicy(mock.Anything, policy).Return(nil).Once()
require.NoError(t, meta.ApplyAlterRLSPolicy(ctx, policy))
require.Equal(t, policy, collection.RLSPolicies["tenant"])
require.NotSame(t, policy, collection.RLSPolicies["tenant"])
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(nil, merr.ErrIoKeyNotFound).Once()
metadata, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "alice")
require.NoError(t, err)
require.Empty(t, metadata.Principals)
catalog.AssertNotCalled(t, "ListRLSPolicies", mock.Anything, mock.Anything)
metadata, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_POLICIES, "")
require.NoError(t, err)
require.Len(t, metadata.Policies, 1)
require.Equal(t, int64(100), metadata.Policies[0].PolicyID)
require.Len(t, collection.RLSPolicies, 1)
catalog.AssertNotCalled(t, "ListRLSPolicies", mock.Anything, mock.Anything)
}
func TestRLSPolicyMetadataChangesAfterPersistence(t *testing.T) {
ctx := context.Background()
meta, catalog := newRLSMetaTableForTest(t)
collection := meta.collID2Meta[20]
collection.Properties = nil
old := &model.RLSPolicy{CollectionID: 20, PolicyID: 100, PolicyName: "tenant", UsingExpr: "true"}
collection.RLSPolicies = map[string]*model.RLSPolicy{"tenant": old}
updated := model.CloneRLSPolicy(old)
updated.UsingExpr = "false"
catalog.EXPECT().SaveRLSPolicy(mock.Anything, updated).Return(merr.ErrServiceUnavailable).Once()
require.ErrorIs(t, meta.ApplyAlterRLSPolicy(ctx, updated), merr.ErrServiceUnavailable)
require.Same(t, old, collection.RLSPolicies["tenant"])
catalog.EXPECT().SaveRLSPolicy(mock.Anything, updated).Return(nil).Once()
require.NoError(t, meta.ApplyAlterRLSPolicy(ctx, updated))
require.Equal(t, updated, collection.RLSPolicies["tenant"])
require.NotSame(t, updated, collection.RLSPolicies["tenant"])
catalog.EXPECT().DropRLSPolicy(mock.Anything, int64(20), int64(100)).Return(merr.ErrServiceUnavailable).Once()
require.ErrorIs(t, meta.ApplyDropRLSPolicy(ctx, 20, "tenant"), merr.ErrServiceUnavailable)
require.Equal(t, updated, collection.RLSPolicies["tenant"])
catalog.EXPECT().DropRLSPolicy(mock.Anything, int64(20), int64(100)).Return(nil).Once()
require.NoError(t, meta.ApplyDropRLSPolicy(ctx, 20, "tenant"))
require.Empty(t, collection.RLSPolicies)
}
func TestNewMetaTableRejectsPolicyRecoveryFailure(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().ListDatabases(mock.Anything, typeutil.MaxTimestamp).
Return([]*model.Database{model.NewDefaultDatabase(nil)}, nil).Once()
catalog.EXPECT().ListCollections(mock.Anything, util.NonDBID, typeutil.MaxTimestamp).
Return([]*model.Collection{{CollectionID: 20}}, nil).Once()
catalog.EXPECT().ListRLSPolicies(mock.Anything, int64(20)).Return(nil, merr.ErrServiceUnavailable).Once()
meta, err := NewMetaTable(context.Background(), catalog, nil)
require.ErrorIs(t, err, merr.ErrServiceUnavailable)
require.Nil(t, meta, "do not publish metadata when policy recovery fails")
}
func TestGetRLSMetadataOwnsSortedPolicySnapshot(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
policies := map[string]*model.RLSPolicy{
"z": {PolicyName: "z", PolicyID: 100, Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery}},
"a": {PolicyName: "a", PolicyID: 200, Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery}},
}
meta.collID2Meta[20].RLSPolicies = policies
metadata, err := meta.GetRLSMetadata(context.Background(), 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_POLICIES, "")
require.NoError(t, err)
require.Len(t, metadata.Policies, 2)
require.Equal(t, "a", metadata.Policies[0].PolicyName)
require.Equal(t, "z", metadata.Policies[1].PolicyName)
metadata.Policies[0].PolicyName = "changed"
metadata.Policies[0].Actions[0] = rlsutil.PolicyActionInsert
require.Equal(t, "a", policies["a"].PolicyName)
require.Equal(t, rlsutil.PolicyActionQuery, policies["a"].Actions[0])
}
func TestAlterCollectionPreservesRLSPolicies(t *testing.T) {
for _, test := range []struct {
name string
before string
after string
properties bool
}{
{"unrelated disabled alter", "false", "true", false},
{"unrelated enabled alter", "true", "false", false},
{"enabled property unchanged", "true", "true", true},
{"enable", "false", "true", true},
{"disable", "true", "false", true},
{"remove enable property", "true", "", true},
} {
t.Run(test.name, func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
collection := meta.collID2Meta[20]
collection.Properties = common.NewKeyValuePairs(map[string]string{common.RLSEnabledKey: test.before})
collection.RLSPolicies = map[string]*model.RLSPolicy{
"tenant": {CollectionID: 20, PolicyID: 100, PolicyName: "tenant"},
}
catalog.EXPECT().AlterCollection(mock.Anything, mock.Anything, mock.Anything, metastore.MODIFY, uint64(2), false).Return(nil).Once()
paths := []string{message.FieldMaskCollectionDescription}
if test.properties {
paths = append(paths, message.FieldMaskCollectionProperties)
}
updates := &message.AlterCollectionMessageUpdates{Description: "updated"}
if test.after != "" {
updates.Properties = common.NewKeyValuePairs(map[string]string{common.RLSEnabledKey: test.after})
}
raw := message.NewAlterCollectionMessageBuilderV2().
WithHeader(&message.AlterCollectionMessageHeader{
CollectionId: 20,
UpdateMask: &fieldmaskpb.FieldMask{Paths: paths},
}).
WithBody(&message.AlterCollectionMessageBody{Updates: updates}).
WithBroadcast([]string{"control"}).MustBuildBroadcast()
require.NoError(t, meta.AlterCollection(context.Background(), message.BroadcastResultAlterCollectionMessageV2{
Message: message.MustAsBroadcastAlterCollectionMessageV2(raw),
Results: map[string]*message.AppendResult{"control": {TimeTick: 2}},
}))
require.Equal(t, collection.RLSPolicies, meta.collID2Meta[20].RLSPolicies)
catalog.AssertNotCalled(t, "ListRLSPolicies", mock.Anything, mock.Anything)
enabled, err := common.IsRLSEnabled(meta.collID2Meta[20].Properties...)
require.NoError(t, err)
expected := test.before
if test.properties {
expected = test.after
}
require.Equal(t, expected == "true", enabled)
})
}
}
func TestResolveRLSCollectionAllowsDisabledCollection(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
for _, properties := range [][]*commonpb.KeyValuePair{
nil,
{{Key: common.RLSEnabledKey, Value: "false"}},
} {
meta.collID2Meta[20].Properties = properties
collection, err := meta.resolveRLSCollection(context.Background(), "db1", "coll1")
require.NoError(t, err)
require.Equal(t, int64(20), collection.CollectionID)
}
meta.collID2Meta[20].Properties = []*commonpb.KeyValuePair{{Key: common.RLSEnabledKey, Value: "invalid"}}
_, err := meta.resolveRLSCollection(context.Background(), "db1", "coll1")
require.ErrorIs(t, err, merr.ErrDataIntegrity)
}
func TestReloadCollectionsRLSMetadataQueriesAllCollections(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{catalog: catalog}
collections := []*model.Collection{
nil,
{CollectionID: 10},
{
CollectionID: 20,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "true"},
},
},
{
CollectionID: 30,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "false"},
},
},
}
for _, id := range []int64{10, 20, 30} {
catalog.EXPECT().ListRLSPolicies(mock.Anything, id).Return(nil, nil).Once()
}
require.NoError(t, meta.reloadCollectionsRLSMetadata(context.Background(), collections))
}
func TestReloadCollectionsRLSMetadataUsesBoundedConcurrency(t *testing.T) {
const collectionCount = rlsRecoveryConcurrency * 2
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{catalog: catalog}
collections := make([]*model.Collection, 0, collectionCount)
for id := int64(1); id <= collectionCount; id++ {
collections = append(collections, &model.Collection{
CollectionID: id,
Properties: []*commonpb.KeyValuePair{
{Key: common.RLSEnabledKey, Value: "true"},
},
})
}
var inFlight atomic.Int32
var maxInFlight atomic.Int32
release := make(chan struct{})
unblock := sync.OnceFunc(func() { close(release) })
trackCall := func() {
current := inFlight.Add(1)
for {
previous := maxInFlight.Load()
if current <= previous || maxInFlight.CompareAndSwap(previous, current) {
break
}
}
<-release
inFlight.Add(-1)
}
catalog.EXPECT().ListRLSPolicies(mock.Anything, mock.Anything).RunAndReturn(
func(context.Context, int64) ([]*model.RLSPolicy, error) {
trackCall()
return nil, nil
}).Times(collectionCount)
done := make(chan struct{})
var reloadErr error
go func() {
reloadErr = meta.reloadCollectionsRLSMetadata(context.Background(), collections)
close(done)
}()
t.Cleanup(func() {
unblock()
<-done
})
require.Eventually(t, func() bool {
return maxInFlight.Load() == rlsRecoveryConcurrency
}, time.Second, time.Millisecond)
select {
case <-done:
t.Fatal("metadata recovery returned before policies loaded")
default:
}
unblock()
<-done
require.NoError(t, reloadErr)
assert.LessOrEqual(t, maxInFlight.Load(), int32(rlsRecoveryConcurrency))
}
func buildAlterUserMessage(credInfo *internalpb.CredentialInfo, timetick uint64) message.BroadcastResultAlterUserMessageV2 {
msg := message.NewAlterUserMessageBuilderV2().
WithHeader(&message.AlterUserMessageHeader{
UserEntity: &milvuspb.UserEntity{
Name: credInfo.Username,
},
}).
WithBody(&message.AlterUserMessageBody{
CredentialInfo: credInfo,
}).
WithBroadcast([]string{funcutil.GetControlChannel("by-dev-rootcoord-dml_1")}).
MustBuildBroadcast()
return message.BroadcastResultAlterUserMessageV2{
Message: message.MustAsBroadcastAlterUserMessageV2(msg),
Results: map[string]*message.AppendResult{
funcutil.GetControlChannel("by-dev-rootcoord-dml_1"): {TimeTick: timetick},
},
}
}
func buildDropUserMessage(credInfo *internalpb.CredentialInfo, timetick uint64) message.BroadcastResultDropUserMessageV2 {
msg := message.NewDropUserMessageBuilderV2().
WithHeader(&message.DropUserMessageHeader{
UserName: credInfo.Username,
}).
WithBody(&message.DropUserMessageBody{}).
WithBroadcast([]string{funcutil.GetControlChannel("by-dev-rootcoord-dml_1")}).
MustBuildBroadcast()
return message.BroadcastResultDropUserMessageV2{
Message: message.MustAsBroadcastDropUserMessageV2(msg),
Results: map[string]*message.AppendResult{
funcutil.GetControlChannel("by-dev-rootcoord-dml_1"): {TimeTick: timetick},
},
}
}
func TestMetaTable_RLSMetadata(t *testing.T) {
ctx := context.Background()
t.Run("create duplicate and list policies", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
createReq := &rlsutil.CreateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "dept_read",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == $current_principal_tags['dept']",
Description: "department read policy",
}
catalog.EXPECT().SaveRLSPolicy(mock.Anything, mock.Anything).RunAndReturn(
func(_ context.Context, policy *model.RLSPolicy) error {
assert.Equal(t, int64(10), policy.DBID)
assert.Equal(t, int64(20), policy.CollectionID)
assert.Equal(t, int64(100), policy.PolicyID)
assert.Equal(t, createReq.GetPolicyName(), policy.PolicyName)
assert.Equal(t, createReq.GetActions(), policy.Actions)
return nil
}).Once()
collectionID, err := meta.CreateRLSPolicy(ctx, createReq, 100)
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
require.Len(t, meta.collID2Meta[20].RLSPolicies, 1)
assert.Equal(t, "dept_read", meta.collID2Meta[20].RLSPolicies["dept_read"].PolicyName)
t.Run("duplicate create ignores lowered creation limits", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyNameLength.Key, "1")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key, "1")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxExpressionLength.Key, "1")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyNameLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxExpressionLength.Key)
})
_, err := meta.PrepareCreateRLSPolicy(ctx, createReq, unallocatedRLSPolicyID)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
require.Contains(t, err.Error(), "already exists")
})
_, err = meta.CreateRLSPolicy(ctx, createReq, 101)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
require.Contains(t, err.Error(), "already exists")
require.Len(t, meta.collID2Meta[20].RLSPolicies, 1)
require.Equal(t, int64(100), meta.collID2Meta[20].RLSPolicies["dept_read"].PolicyID)
_, err = meta.CreateRLSPolicy(ctx, createReq, 102)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
require.Contains(t, err.Error(), "already exists")
updateReq := &rlsutil.UpdateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "dept_read",
PolicyType: rlsutil.PolicyTypeRestrictive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionSearch},
UsingExpr: "owner == $current_principal",
Description: "updated policy",
}
catalog.EXPECT().SaveRLSPolicy(mock.Anything, mock.Anything).RunAndReturn(
func(_ context.Context, policy *model.RLSPolicy) error {
assert.Equal(t, int64(100), policy.PolicyID)
assert.Equal(t, rlsutil.PolicyTypeRestrictive, policy.PolicyType)
assert.Equal(t, []rlsutil.PolicyAction{rlsutil.PolicyActionSearch}, policy.Actions)
assert.Equal(t, "owner == $current_principal", policy.UsingExpr)
assert.Equal(t, "updated policy", policy.Description)
return nil
}).Once()
collectionID, err = meta.UpdateRLSPolicy(ctx, updateReq)
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
require.Len(t, meta.collID2Meta[20].RLSPolicies, 1)
assert.Equal(t, int64(100), meta.collID2Meta[20].RLSPolicies["dept_read"].PolicyID)
assert.Equal(t, updateReq.GetUsingExpr(), meta.collID2Meta[20].RLSPolicies["dept_read"].UsingExpr)
t.Run("update ignores lowered creation name limit", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyNameLength.Key, "1")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyNameLength.Key)
})
policy, err := meta.PrepareUpdateRLSPolicy(ctx, updateReq)
require.NoError(t, err)
require.Equal(t, int64(100), policy.PolicyID)
require.Equal(t, "dept_read", policy.PolicyName)
})
_, err = meta.UpdateRLSPolicy(ctx, &rlsutil.UpdateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "missing_policy",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == $current_principal_tags['dept']",
})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
policies, err := meta.ListRLSPolicies(ctx, &rlsutil.ListRowPoliciesRequest{DbName: "db1", CollectionName: "coll1"})
require.NoError(t, err)
require.Len(t, policies, 1)
assert.Equal(t, "dept_read", policies[0].PolicyName)
assert.Equal(t, updateReq.GetUsingExpr(), policies[0].UsingExpr)
})
t.Run("reject unknown action value", func(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
_, err := meta.CreateRLSPolicy(ctx, &rlsutil.CreateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "unknown_action_value",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyAction(100)},
UsingExpr: "dept == $current_principal_tags['dept']",
}, 101)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "invalid RLS policy action")
})
t.Run("reject whitespace-only required expressions", func(t *testing.T) {
for _, test := range []struct {
name string
action rlsutil.PolicyAction
usingExpr string
checkExpr string
required string
}{
{
name: "using expression",
action: rlsutil.PolicyActionQuery,
usingExpr: " \t ",
checkExpr: "dept == 'sales'",
required: "using_expr is required",
},
{
name: "check expression",
action: rlsutil.PolicyActionInsert,
usingExpr: "dept == 'sales'",
checkExpr: " \n ",
required: "check_expr is required",
},
} {
t.Run(test.name, func(t *testing.T) {
err := rlsutil.ValidatePolicy("policy", rlsutil.PolicyTypePermissive, []rlsutil.PolicyAction{test.action}, test.usingExpr, test.checkExpr)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
require.Contains(t, err.Error(), test.required)
})
}
})
t.Run("validate policy expressions", func(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
coll := meta.collID2Meta[20]
validExprs := []string{
"true",
"false",
"dept == $current_principal_tags['dept']",
"age == $current_principal_tags['age']",
"owner == $current_principal",
`dept == "{literal}"`,
`dept == r"{raw_literal}"`,
`dept == "$current_principal"`,
`dept == "$current_principal_tags['dept']"`,
`dept == r"$current_principal"`,
`dept in ["sales", "engineering"]`,
"array_contains(tags, $current_principal)",
"array_contains(scores, $current_principal_tags['score'])",
`array_contains_all(tags, ["sales", "engineering"])`,
"array_contains_all(scores, [1, 2])",
"array_contains_all(tags, $current_principal_tags['groups'])",
"array_contains_any(scores, $current_principal_tags['scores'])",
`not (dept == "sales")`,
`not array_contains(tags, "blocked")`,
}
for _, expr := range validExprs {
require.NoError(t, validateRLSPolicyExpressions(coll, expr, ""), expr)
}
require.NoError(t, validateRLSPolicyExpressions(coll, "", "false"))
for _, quotedExpr := range []string{
`dept == "$current_principal"`,
`owner == "$current_principal_tags['owner']"`,
} {
templateExpr, _, err := toRLSPolicyTemplateExpr(quotedExpr)
require.NoError(t, err)
assert.Equal(t, quotedExpr, templateExpr)
}
coll.Fields = append(coll.Fields, &model.Field{
FieldID: 108,
Name: "ts",
DataType: schemapb.DataType_Timestamptz,
State: schemapb.FieldState_FieldCreated,
}, &model.Field{
FieldID: 109,
Name: "enabled",
DataType: schemapb.DataType_Bool,
State: schemapb.FieldState_FieldCreated,
}, &model.Field{
FieldID: 110,
Name: "nullable_tags",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
ElementNullable: true,
State: schemapb.FieldState_FieldCreated,
})
coll.Properties = append(coll.Properties, &commonpb.KeyValuePair{Key: common.TimezoneKey, Value: "Asia/Shanghai"})
require.NoError(t, validateRLSPolicyExpressions(coll, "ts == ISO '2025-01-01 00:00:00'", ""))
invalidExprs := []string{
`dept == "sales" and owner == "alice"`,
`dept == "sales" or owner == "alice"`,
"missing == 'sales'",
"dept == $principal.name",
"dept == $current_principal_tags['x'y'] and dept == {raw}",
`dept == $current_principal_tags['x\'] and dept == {raw}`,
"dept == {foo}",
"dept == {__rls_principal}",
"array_contains(tags, {foo})",
"age + $current_principal == id",
"age > 10",
"dept in $current_principal_tags['dept']",
"json_contains(metadata, \"sales\")",
"array_contains(scores, $current_principal)",
"array_contains_all(scores, [1.0])",
"array_contains_any(scores, [1, 1.5])",
"ts == $current_principal_tags['ts']",
"enabled == $current_principal_tags['enabled']",
`array_contains(nullable_tags, "red")`,
}
for _, expr := range invalidExprs {
err := validateRLSPolicyExpressions(coll, expr, "")
require.ErrorIs(t, err, merr.ErrParameterInvalid, expr)
}
for _, expr := range []string{
`dept == "sales" and owner == "alice"`,
`dept == "sales" or owner == "alice"`,
} {
err := validateRLSPolicyExpressions(coll, expr, "")
require.ErrorContains(t, err, "compound RLS expressions are not supported")
}
require.NoError(t, validateRLSPolicyExpressions(coll, "", `array_contains(nullable_tags, "red")`))
require.ErrorIs(t, validateRLSPolicyExpressions(coll, `not array_contains(nullable_tags, "red")`, ""), merr.ErrParameterInvalid)
require.ErrorIs(t, rlsutil.ValidateParsedExpression(&planpb.Expr{Expr: &planpb.Expr_UnaryRangeExpr{UnaryRangeExpr: &planpb.UnaryRangeExpr{
ColumnInfo: &planpb.ColumnInfo{FieldId: 101, DataType: schemapb.DataType_VarChar},
Op: planpb.OpType_Equal,
}}}, nil), merr.ErrParameterInvalid)
malformedTagExpr := "dept == $current_principal_tags['x'y'] and dept == {raw}"
require.ErrorIs(t, validateRLSPolicyExpressionsWithSchema(
coll.ToCollectionSchemaPB(),
malformedTagExpr,
"",
), merr.ErrParameterInvalid)
paramtable.Get().Save(Params.ProxyCfg.RLSMaxArrayLiteralElements.Key, "1")
defer paramtable.Get().Reset(Params.ProxyCfg.RLSMaxArrayLiteralElements.Key)
err := validateRLSPolicyExpressions(coll, `dept in ["sales", "engineering"]`, "")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max array literal elements")
err = validateRLSPolicyExpressions(coll, `not (dept in ["sales", "engineering"])`, "")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max array literal elements")
})
t.Run("rls limits use paramtable", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyNameLength.Key, "3")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key, "3")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxExpressionLength.Key, "4")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPrincipalNameLength.Key, "3")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagKeyLength.Key, "3")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagValueLength.Key, "3")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key, "1")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPoliciesPerCollection.Key, "1")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyNameLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxExpressionLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPrincipalNameLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxTagKeyLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxTagValueLength.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key)
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPoliciesPerCollection.Key)
})
err := rlsutil.ValidatePolicy("long", rlsutil.PolicyTypePermissive, []rlsutil.PolicyAction{rlsutil.PolicyActionQuery}, "a", "")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max length 3")
err = rlsutil.ValidatePolicyDescription("long")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max length 3")
err = rlsutil.ValidatePolicy("p", rlsutil.PolicyTypePermissive, []rlsutil.PolicyAction{rlsutil.PolicyActionQuery}, "owner == 'alice'", "")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max length 4")
err = rlsutil.ValidatePrincipalNameWithLimit("alice")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "max length 3")
require.NoError(t, rlsutil.ValidatePrincipalName("alice"))
err = rlsutil.ValidateTags(map[string]rlsutil.TagValue{"team": rlsutil.NewStringTagValue("abc")})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "tag key exceeds max length 3")
err = rlsutil.ValidateTags(map[string]rlsutil.TagValue{"a": rlsutil.NewStringTagValue("abcd")})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "tag value exceeds max length 3")
err = rlsutil.ValidateTags(map[string]rlsutil.TagValue{
"a": rlsutil.NewInt64TagValue(1),
"b": rlsutil.NewInt64TagValue(2),
})
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
err = rlsutil.ValidateTags(map[string]rlsutil.TagValue{"a']": rlsutil.NewInt64TagValue(1)})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "reserved character")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxExpressionLength.Key, "100")
meta, _ := newRLSMetaTableForTest(t)
err = validateRLSPolicyExpressions(meta.collID2Meta[20], "dept == $current_principal_tags['team']", "")
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "tag key exceeds max length 3")
require.NoError(t, validateRLSPolicyExpressionsWithSchema(
meta.collID2Meta[20].ToCollectionSchemaPB(),
"dept == $current_principal_tags['team']",
"",
))
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{"existing": {PolicyName: "existing"}}
_, err = meta.CreateRLSPolicy(ctx, &rlsutil.CreateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "p",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == 'sales'",
}, 101)
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
})
t.Run("schema validation grandfathers existing array literal quota", func(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
coll := meta.collID2Meta[20]
policy := &model.RLSPolicy{
PolicyName: "existing",
UsingExpr: `dept in ["sales", "engineering"]`,
}
paramtable.Get().Save(Params.ProxyCfg.RLSMaxArrayLiteralElements.Key, "1")
defer paramtable.Get().Reset(Params.ProxyCfg.RLSMaxArrayLiteralElements.Key)
err := validateRLSPolicyExpressionsWithSchema(coll.ToCollectionSchemaPB(), policy.UsingExpr, "")
require.ErrorContains(t, err, "max array literal elements")
require.NoError(t, validateRLSPoliciesWithSchema(map[string]*model.RLSPolicy{"existing": policy}, coll.ToCollectionSchemaPB()))
})
t.Run("reject deprecated string fields", func(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{
FieldID: 100,
Name: "legacy",
DataType: schemapb.DataType_String,
}}}
require.ErrorIs(t, validateRLSPolicyExpressionsWithSchema(schema, "", `legacy == "value"`), merr.ErrParameterInvalid)
require.ErrorIs(t, validateRLSPolicyExpressionsWithSchema(schema, `legacy in ["value"]`, ""), merr.ErrParameterInvalid)
require.ErrorIs(t, validateRLSPoliciesWithSchema(map[string]*model.RLSPolicy{
"existing": {PolicyName: "existing", CheckExpr: `legacy == "value"`},
}, schema), merr.ErrParameterInvalid)
})
t.Run("reject policy descriptions before persistence", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key, "3")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPolicyDescriptionLength.Key)
})
meta, _ := newRLSMetaTableForTest(t)
_, err := meta.CreateRLSPolicy(ctx, &rlsutil.CreateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "policy",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "true",
Description: "long",
}, 101)
require.ErrorIs(t, err, merr.ErrParameterInvalid)
_, err = meta.UpdateRLSPolicy(ctx, &rlsutil.UpdateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "policy",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "true",
Description: "long",
})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("reject combined expression length on create but allow update", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxCombinedExpressionLength.Key, "24")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxCombinedExpressionLength.Key)
})
meta, _ := newRLSMetaTableForTest(t)
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{
"existing": {
PolicyName: "existing",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "id > 0",
},
}
_, err := meta.CreateRLSPolicy(ctx, &rlsutil.CreateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "too_long",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == 'sales'",
}, 101)
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
assert.Contains(t, err.Error(), "combined RLS using expression")
assert.Contains(t, err.Error(), "max length 24")
meta, _ = newRLSMetaTableForTest(t)
otherPolicy := &model.RLSPolicy{
PolicyID: 100,
PolicyName: "other",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "id > 0",
DBID: 10,
CollectionID: 20,
}
oldPolicy := &model.RLSPolicy{
PolicyID: 101,
PolicyName: "existing",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "id < 100",
DBID: 10,
CollectionID: 20,
}
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{
"other": otherPolicy,
"existing": oldPolicy,
}
policy, err := meta.PrepareUpdateRLSPolicy(ctx, &rlsutil.UpdateRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "existing",
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == 'sales'",
})
require.NoError(t, err)
assert.Equal(t, "dept == 'sales'", policy.UsingExpr)
})
t.Run("combine policy expressions with runtime semantics", func(t *testing.T) {
policies := []*model.RLSPolicy{
{
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "dept == $current_principal_tags['dept']",
},
{
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "owner == $current_principal",
},
{
PolicyType: rlsutil.PolicyTypeRestrictive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "active == true",
},
}
const expected = "((dept == {__rls_tag_0}) or (owner == {__rls_principal})) and ((active == true))"
oldLimit := Params.ProxyCfg.RLSMaxCombinedExpressionLength.SwapTempValue(strconv.Itoa(len(expected)))
defer Params.ProxyCfg.RLSMaxCombinedExpressionLength.SwapTempValue(oldLimit)
require.NoError(t, validateRLSCombinedExpressionLength(policies))
Params.ProxyCfg.RLSMaxCombinedExpressionLength.SwapTempValue(strconv.Itoa(len(expected) - 1))
require.ErrorIs(t, validateRLSCombinedExpressionLength(policies), merr.ErrServiceQuotaExceeded)
})
t.Run("combined limit ignores unused expression kinds", func(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.RLSMaxCombinedExpressionLength.Key, "10")
t.Cleanup(func() {
paramtable.Get().Reset(Params.ProxyCfg.RLSMaxCombinedExpressionLength.Key)
})
err := validateRLSCombinedExpressionLength([]*model.RLSPolicy{
{
PolicyType: rlsutil.PolicyTypePermissive,
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
UsingExpr: "id > 0",
CheckExpr: "this optional check expression is not used by query",
},
})
require.NoError(t, err)
})
t.Run("get bulk metadata by collection id", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
meta.collID2Meta[20].DBName = ""
meta.collID2Meta[20].DBID = 999
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{
"dept_read": {
DBID: 10,
CollectionID: 20,
PolicyID: 100,
PolicyName: "dept_read",
Actions: []rlsutil.PolicyAction{rlsutil.PolicyActionQuery},
},
}
catalogPrincipal := &model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
},
}
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return([]*model.RLSPrincipal{catalogPrincipal}, nil).Once()
metadata, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_ALL, "")
require.NoError(t, err)
require.Equal(t, int64(20), metadata.CollectionID)
require.Len(t, metadata.Policies, 1)
require.Equal(t, "dept_read", metadata.Policies[0].PolicyName)
require.Len(t, metadata.Principals, 1)
require.Equal(t, map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
}, metadata.Principals[0].Tags)
require.Equal(t, int64(999), metadata.Principals[0].DBID)
require.Equal(t, int64(20), metadata.Principals[0].CollectionID)
require.Same(t, catalogPrincipal, metadata.Principals[0])
metadata.Policies[0].Actions[0] = rlsutil.PolicyActionInsert
require.Equal(t, rlsutil.PolicyActionQuery, meta.collID2Meta[20].RLSPolicies["dept_read"].Actions[0])
policyMetadata, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_POLICIES, "")
require.NoError(t, err)
require.Len(t, policyMetadata.Policies, 1)
require.Empty(t, policyMetadata.Principals)
policyMetadata.Policies[0].Actions[0] = rlsutil.PolicyActionInsert
require.Equal(t, rlsutil.PolicyActionQuery, meta.collID2Meta[20].RLSPolicies["dept_read"].Actions[0])
secondCatalogPrincipal := model.CloneRLSPrincipal(catalogPrincipal)
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return([]*model.RLSPrincipal{secondCatalogPrincipal}, nil).Once()
principalMetadata, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "")
require.NoError(t, err)
require.Empty(t, principalMetadata.Policies)
require.Len(t, principalMetadata.Principals, 1)
require.Same(t, secondCatalogPrincipal, principalMetadata.Principals[0])
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return(nil, merr.WrapErrIoFailedReason("list failed")).Once()
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "")
require.ErrorIs(t, err, merr.ErrIoFailed)
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return(nil, merr.WrapErrServiceQuotaExceededMsg("list too large")).Once()
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_ALL, "")
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return([]*model.RLSPrincipal{nil}, nil).Once()
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "")
require.ErrorIs(t, err, merr.ErrDataIntegrity)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
},
}, nil).Once()
filtered, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "alice")
require.NoError(t, err)
require.Len(t, filtered.Principals, 1)
require.Equal(t, "alice", filtered.Principals[0].PrincipalName)
require.Equal(t, int64(999), filtered.Principals[0].DBID)
require.Equal(t, int64(20), filtered.Principals[0].CollectionID)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "bob").Return(nil, merr.ErrIoKeyNotFound).Once()
missing, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "bob")
require.NoError(t, err)
require.Empty(t, missing.Principals)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "carol").Return(nil, merr.WrapErrIoFailedReason("lookup failed")).Once()
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "carol")
require.ErrorIs(t, err, merr.ErrIoFailed)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "dave").Return(nil, nil).Once()
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "dave")
require.ErrorIs(t, err, merr.ErrDataIntegrity)
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_POLICIES, "alice")
require.ErrorIs(t, err, merr.ErrServiceInternal)
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind(100), "")
require.ErrorIs(t, err, merr.ErrServiceInternal)
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_UNSPECIFIED, "")
require.ErrorIs(t, err, merr.ErrServiceInternal)
_, err = meta.GetRLSMetadata(ctx, 0, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_ALL, "")
require.ErrorIs(t, err, merr.ErrServiceInternal)
meta.collID2Meta[20].State = pb.CollectionState_CollectionDropping
_, err = meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_ALL, "")
require.ErrorIs(t, err, merr.ErrCollectionNotFound)
})
t.Run("drop policy by name", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{"dept_read": {
PolicyID: 100,
PolicyName: "dept_read",
}}
catalog.EXPECT().DropRLSPolicy(mock.Anything, int64(20), int64(100)).Return(nil).Once()
collectionID, err := meta.DropRLSPolicy(ctx, &rlsutil.DropRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "dept_read",
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
require.Empty(t, meta.collID2Meta[20].RLSPolicies)
collectionID, err = meta.DropRLSPolicy(ctx, &rlsutil.DropRowPolicyRequest{
DbName: "db1",
CollectionName: "coll1",
PolicyName: "dept_read",
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
})
t.Run("reject empty principal tags", func(t *testing.T) {
meta, _ := newRLSMetaTableForTest(t)
for _, tags := range []map[string]rlsutil.TagValue{nil, {}} {
_, err := meta.SetRLSPrincipalTags(ctx, &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
Tags: tags,
})
require.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "RLS principal tags are empty")
}
})
t.Run("set principal tags incrementally", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
existing := &model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
},
}
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(existing, nil).Once()
prepared, err := meta.PrepareSetRLSPrincipalTags(ctx, &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("engineering"),
"region": rlsutil.NewStringTagValue("west"),
},
})
require.NoError(t, err)
assert.Equal(t, map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("engineering"),
"tier": rlsutil.NewStringTagValue("gold"),
"region": rlsutil.NewStringTagValue("west"),
}, prepared.Tags)
assert.Equal(t, rlsutil.NewStringTagValue("sales"), existing.Tags["dept"])
assert.NotContains(t, existing.Tags, "region")
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key, "2")
defer paramtable.Get().Reset(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(existing, nil).Once()
_, err = meta.PrepareSetRLSPrincipalTags(ctx, &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{"region": rlsutil.NewStringTagValue("west")},
})
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
})
t.Run("deleting the last selected tag drops the principal", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
},
}, nil).Once()
principal, drop, err := meta.PrepareDeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
TagKeys: []string{"dept"},
})
require.NoError(t, err)
assert.True(t, drop)
assert.Equal(t, "alice", principal.PrincipalName)
})
t.Run("principal tags", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
setReq := &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
},
}
catalog.EXPECT().SaveRLSPrincipal(mock.Anything, mock.Anything).RunAndReturn(
func(_ context.Context, principal *model.RLSPrincipal) error {
assert.Equal(t, int64(10), principal.DBID)
assert.Equal(t, int64(20), principal.CollectionID)
assert.Equal(t, "alice", principal.PrincipalName)
assert.Equal(t, setReq.GetTags(), principal.Tags)
return nil
}).Once()
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(nil, merr.ErrIoKeyNotFound).Once()
prepared, err := meta.PrepareSetRLSPrincipalTags(ctx, setReq)
require.NoError(t, err)
require.NoError(t, meta.ApplyAlterRLSPrincipal(ctx, prepared))
collectionID := prepared.CollectionID
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
catalog.EXPECT().SaveRLSPrincipal(mock.Anything, mock.MatchedBy(func(principal *model.RLSPrincipal) bool {
return principal.PrincipalName == "alice" && assert.ObjectsAreEqual(principal.Tags["dept"], rlsutil.NewStringTagValue("engineering"))
})).Return(nil).Once()
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{PrincipalName: "alice"}, nil).Once()
setReq.Tags["dept"] = rlsutil.NewStringTagValue("engineering")
prepared, err = meta.PrepareSetRLSPrincipalTags(ctx, setReq)
require.NoError(t, err)
require.NoError(t, meta.ApplyAlterRLSPrincipal(ctx, prepared))
collectionID = prepared.CollectionID
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "bob").Return(nil, merr.ErrIoKeyNotFound).Once()
prepared, err = meta.PrepareSetRLSPrincipalTags(ctx, &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "bob",
Tags: map[string]rlsutil.TagValue{"dept": rlsutil.NewStringTagValue("sales")},
})
require.NoError(t, err)
require.Equal(t, "bob", prepared.PrincipalName)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "carol").Return(nil, merr.WrapErrIoFailedReason("lookup failed")).Once()
_, err = meta.PrepareSetRLSPrincipalTags(ctx, &rlsutil.SetRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "carol",
Tags: map[string]rlsutil.TagValue{"dept": rlsutil.NewStringTagValue("sales")},
})
require.ErrorIs(t, err, merr.ErrIoFailed)
assert.Contains(t, err.Error(), "failed to get RLS principal")
// Existing principals remain addressable if the creation limit is
// lowered after they were stored.
paramtable.Get().Save(Params.ProxyCfg.RLSMaxPrincipalNameLength.Key, "3")
defer paramtable.Get().Reset(Params.ProxyCfg.RLSMaxPrincipalNameLength.Key)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{PrincipalName: "alice"}, nil).Once()
catalog.EXPECT().SaveRLSPrincipal(mock.Anything, mock.MatchedBy(func(principal *model.RLSPrincipal) bool {
return principal.PrincipalName == "alice" && assert.ObjectsAreEqual(principal.Tags["dept"], rlsutil.NewStringTagValue("support"))
})).Return(nil).Once()
setReq.Tags["dept"] = rlsutil.NewStringTagValue("support")
collectionID, err = meta.SetRLSPrincipalTags(ctx, setReq)
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
},
}, nil).Once()
tags, err := meta.GetRLSPrincipalTags(ctx, &rlsutil.GetRLSPrincipalTagsRequest{DbName: "db1", CollectionName: "coll1", PrincipalName: "alice"})
require.NoError(t, err)
assert.Equal(t, map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
}, tags)
listMeta, listCatalog := newRLSMetaTableForTest(t)
listCatalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return([]*model.RLSPrincipal{
{PrincipalName: "bob"},
{PrincipalName: "alice"},
}, nil).Once()
principals, err := listMeta.ListRLSPrincipals(ctx, &rlsutil.ListRLSPrincipalsRequest{DbName: "db1", CollectionName: "coll1"})
require.NoError(t, err)
assert.Equal(t, []string{"alice", "bob"}, principals)
listCatalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).Return(nil, merr.WrapErrServiceQuotaExceededMsg("list too large")).Once()
_, err = listMeta.ListRLSPrincipals(ctx, &rlsutil.ListRLSPrincipalsRequest{DbName: "db1", CollectionName: "coll1"})
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
},
}, nil).Once()
catalog.EXPECT().SaveRLSPrincipal(mock.Anything, mock.Anything).RunAndReturn(
func(_ context.Context, principal *model.RLSPrincipal) error {
assert.Equal(t, map[string]rlsutil.TagValue{
"tier": rlsutil.NewStringTagValue("gold"),
}, principal.Tags)
return nil
}).Once()
collectionID, err = meta.DeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
TagKeys: []string{"dept"},
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"tier": rlsutil.NewStringTagValue("gold"),
},
}, nil).Once()
catalog.EXPECT().DropRLSPrincipal(mock.Anything, int64(20), "alice").Return(nil).Once()
collectionID, err = meta.DeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key, "2")
defer paramtable.Get().Reset(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(&model.RLSPrincipal{
DBID: 10,
CollectionID: 20,
PrincipalName: "alice",
Tags: map[string]rlsutil.TagValue{
"dept": rlsutil.NewStringTagValue("sales"),
"tier": rlsutil.NewStringTagValue("gold"),
},
}, nil).Once()
catalog.EXPECT().SaveRLSPrincipal(mock.Anything, mock.MatchedBy(func(principal *model.RLSPrincipal) bool {
return len(principal.Tags) == 1 && assert.ObjectsAreEqual(principal.Tags["tier"], rlsutil.NewStringTagValue("gold"))
})).Return(nil).Once()
collectionID, err = meta.DeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
TagKeys: []string{"dept", "dept"},
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
paramtable.Get().Save(Params.ProxyCfg.RLSMaxTagsPerPrincipal.Key, "1")
_, err = meta.DeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
TagKeys: []string{"dept", "tier"},
})
require.ErrorIs(t, err, merr.ErrServiceQuotaExceeded)
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").Return(nil, merr.ErrIoKeyNotFound).Once()
catalog.EXPECT().DropRLSPrincipal(mock.Anything, int64(20), "alice").Return(nil).Once()
collectionID, err = meta.DeleteRLSPrincipalTags(ctx, &rlsutil.DeleteRLSPrincipalTagsRequest{
DbName: "db1",
CollectionName: "coll1",
PrincipalName: "alice",
})
require.NoError(t, err)
require.Equal(t, int64(20), collectionID)
})
}
func TestMetaTable_RLSCatalogIODoesNotHoldGlobalDDLock(t *testing.T) {
ctx := context.Background()
assertCollectionReadCompletes := func(t *testing.T, meta *MetaTable) {
t.Helper()
done := make(chan error, 1)
go func() {
_, err := meta.GetCollectionByName(ctx, "db1", "coll1", typeutil.MaxTimestamp, false)
done <- err
}()
select {
case err := <-done:
require.NoError(t, err)
case <-time.After(time.Second):
assert.Fail(t, "collection read blocked by RLS catalog I/O")
}
}
t.Run("drop", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{
"policy": {
CollectionID: 20,
PolicyID: 100,
PolicyName: "policy",
},
}
catalogStarted := make(chan struct{})
releaseCatalog := make(chan struct{})
catalog.EXPECT().DropRLSPolicy(mock.Anything, int64(20), int64(100)).RunAndReturn(
func(context.Context, int64, int64) error {
close(catalogStarted)
<-releaseCatalog
return nil
}).Once()
done := make(chan error, 1)
go func() {
done <- meta.ApplyDropRLSPolicy(ctx, 20, "policy")
}()
<-catalogStarted
assertCollectionReadCompletes(t, meta)
close(releaseCatalog)
require.NoError(t, <-done)
})
t.Run("apply", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
catalogStarted := make(chan struct{})
releaseCatalog := make(chan struct{})
catalog.EXPECT().SaveRLSPolicy(mock.Anything, mock.Anything).RunAndReturn(
func(context.Context, *model.RLSPolicy) error {
close(catalogStarted)
<-releaseCatalog
return nil
}).Once()
done := make(chan error, 1)
go func() {
done <- meta.ApplyAlterRLSPolicy(ctx, &model.RLSPolicy{
DBID: 10,
CollectionID: 20,
PolicyID: 100,
PolicyName: "policy",
})
}()
<-catalogStarted
assertCollectionReadCompletes(t, meta)
close(releaseCatalog)
require.NoError(t, <-done)
})
t.Run("get filtered principal metadata", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
catalogStarted := make(chan struct{})
releaseCatalog := make(chan struct{})
catalog.EXPECT().GetRLSPrincipal(mock.Anything, int64(20), "alice").RunAndReturn(
func(context.Context, int64, string) (*model.RLSPrincipal, error) {
close(catalogStarted)
<-releaseCatalog
return &model.RLSPrincipal{CollectionID: 20, PrincipalName: "alice"}, nil
}).Once()
done := make(chan error, 1)
go func() {
_, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_PRINCIPALS, "alice")
done <- err
}()
<-catalogStarted
assertCollectionReadCompletes(t, meta)
meta.ddLock.Lock()
meta.collID2Meta[20].State = pb.CollectionState_CollectionDropping
meta.ddLock.Unlock()
close(releaseCatalog)
require.ErrorIs(t, <-done, merr.ErrCollectionNotFound)
})
t.Run("get all metadata uses current collection identity", func(t *testing.T) {
meta, catalog := newRLSMetaTableForTest(t)
meta.collID2Meta[20].RLSPolicies = map[string]*model.RLSPolicy{
"policy": {DBID: 10, CollectionID: 20, PolicyName: "policy"},
}
catalogStarted := make(chan struct{})
releaseCatalog := make(chan struct{})
catalog.EXPECT().ListRLSPrincipals(mock.Anything, int64(20)).RunAndReturn(
func(context.Context, int64) ([]*model.RLSPrincipal, error) {
close(catalogStarted)
<-releaseCatalog
return []*model.RLSPrincipal{{DBID: 10, CollectionID: 20, PrincipalName: "alice"}}, nil
}).Once()
result := make(chan struct {
metadata *model.RLSMetadata
err error
}, 1)
go func() {
metadata, err := meta.GetRLSMetadata(ctx, 20, rootcoordpb.RLSMetadataKind_RLS_METADATA_KIND_ALL, "")
result <- struct {
metadata *model.RLSMetadata
err error
}{metadata: metadata, err: err}
}()
<-catalogStarted
meta.ddLock.Lock()
meta.collID2Meta[20].DBID = 11
meta.collID2Meta[20].RLSPolicies["policy"].DBID = 11
meta.ddLock.Unlock()
close(releaseCatalog)
got := <-result
require.NoError(t, got.err)
require.Equal(t, int64(11), got.metadata.Policies[0].DBID)
require.Equal(t, int64(11), got.metadata.Principals[0].DBID)
})
}
func TestRbacCredential(t *testing.T) {
mt := generateMetaTable(t)
username := "user" + funcutil.RandomString(10)
credInfo := &internalpb.CredentialInfo{
Username: username,
Tenant: util.DefaultTenant,
}
err := mt.CheckIfAddCredential(context.TODO(), credInfo)
require.NoError(t, err)
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(credInfo, 1))
require.NoError(t, err)
// idempotency
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(credInfo, 1))
require.NoError(t, err)
tests := []struct {
description string
maxUser bool
info *internalpb.CredentialInfo
}{
{"Empty username", false, &internalpb.CredentialInfo{Username: ""}},
{"exceed MaxUserNum", true, &internalpb.CredentialInfo{Username: "user3", Tenant: util.DefaultTenant}},
{"user exist", false, &internalpb.CredentialInfo{Username: username, Tenant: util.DefaultTenant}},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
if test.maxUser {
paramtable.Get().Save(Params.ProxyCfg.MaxUserNum.Key, "1")
} else {
paramtable.Get().Save(Params.ProxyCfg.MaxUserNum.Key, "3")
}
defer paramtable.Get().Reset(Params.ProxyCfg.MaxUserNum.Key)
err := mt.CheckIfAddCredential(context.TODO(), test.info)
assert.Error(t, err)
})
}
// should be ignored if timetick is too low.
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(credInfo, 0))
require.NoError(t, err)
newCred, err := mt.GetCredential(context.TODO(), credInfo.Username)
require.NoError(t, err)
assert.Equal(t, newCred.TimeTick, uint64(1))
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(credInfo, 2))
require.NoError(t, err)
newCred, err = mt.GetCredential(context.TODO(), credInfo.Username)
require.NoError(t, err)
assert.Equal(t, newCred.TimeTick, uint64(2))
// should be ignored if timetick is too low.
err = mt.DeleteCredential(context.TODO(), buildDropUserMessage(credInfo, 2))
require.NoError(t, err)
newCred, err = mt.GetCredential(context.TODO(), credInfo.Username)
require.NoError(t, err)
assert.Equal(t, newCred.TimeTick, uint64(2))
err = mt.DeleteCredential(context.TODO(), buildDropUserMessage(credInfo, 3))
require.NoError(t, err)
newCred, err = mt.GetCredential(context.TODO(), credInfo.Username)
require.ErrorIs(t, err, merr.ErrIoKeyNotFound)
require.Nil(t, newCred)
err = mt.DeleteCredential(context.TODO(), buildDropUserMessage(credInfo, 0))
require.NoError(t, err)
}
func TestRbacCredentialAlterCredentialMergesPartialUpdates(t *testing.T) {
mt := generateMetaTable(t)
ptr := func(s string) *string {
return &s
}
username := "user" + funcutil.RandomString(10)
err := mt.AlterCredential(context.TODO(), buildAlterUserMessage(&internalpb.CredentialInfo{
Username: username,
EncryptedPassword: "old-password",
Description: ptr("initial description"),
}, 1))
require.NoError(t, err)
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(&internalpb.CredentialInfo{
Username: username,
EncryptedPassword: "new-password",
}, 2))
require.NoError(t, err)
cred, err := mt.GetCredential(context.TODO(), username)
require.NoError(t, err)
assert.Equal(t, "new-password", cred.GetEncryptedPassword())
assert.Equal(t, "initial description", cred.GetDescription())
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(&internalpb.CredentialInfo{
Username: username,
Description: ptr("updated description"),
}, 3))
require.NoError(t, err)
cred, err = mt.GetCredential(context.TODO(), username)
require.NoError(t, err)
assert.Equal(t, "new-password", cred.GetEncryptedPassword())
assert.Equal(t, "updated description", cred.GetDescription())
err = mt.AlterCredential(context.TODO(), buildAlterUserMessage(&internalpb.CredentialInfo{
Username: username,
Description: ptr(""),
}, 4))
require.NoError(t, err)
cred, err = mt.GetCredential(context.TODO(), username)
require.NoError(t, err)
assert.Equal(t, "new-password", cred.GetEncryptedPassword())
assert.Equal(t, "", cred.GetDescription())
}
func TestRbacCredentialRejectsInconsistentPasswordUpdate(t *testing.T) {
mt := generateMetaTable(t)
username := "user" + funcutil.RandomString(10)
err := mt.AlterCredential(context.TODO(), buildAlterUserMessage(&internalpb.CredentialInfo{
Username: username,
EncryptedPassword: "old-password",
}, 1))
require.NoError(t, err)
err = mt.CheckIfUpdateCredential(context.TODO(), &internalpb.CredentialInfo{
Username: username,
EncryptedPassword: "new-password",
})
require.Error(t, err)
assert.ErrorContains(t, err, "must include both encrypted and sha256 password")
err = mt.CheckIfUpdateCredential(context.TODO(), &internalpb.CredentialInfo{
Username: username,
Sha256Password: "sha256",
})
require.Error(t, err)
assert.ErrorContains(t, err, "must include both encrypted and sha256 password")
description := "description-only update"
err = mt.CheckIfUpdateCredential(context.TODO(), &internalpb.CredentialInfo{
Username: username,
Description: &description,
})
require.NoError(t, err)
err = mt.CheckIfUpdateCredential(context.TODO(), &internalpb.CredentialInfo{
Username: username,
})
require.Error(t, err)
assert.ErrorContains(t, err, "credential update must change password or description")
}
func TestRbacCreateRole(t *testing.T) {
mt := generateMetaTable(t)
paramtable.Get().Save(Params.ProxyCfg.MaxRoleNum.Key, "2")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleNum.Key)
err := mt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{Entity: &milvuspb.RoleEntity{Name: "role1"}})
require.NoError(t, err)
err = mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
require.NoError(t, err)
err = mt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{Entity: &milvuspb.RoleEntity{Name: "role2"}})
require.NoError(t, err)
err = mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role2"})
require.NoError(t, err)
tests := []struct {
inEntity *milvuspb.RoleEntity
description string
}{
{&milvuspb.RoleEntity{Name: ""}, "empty string"},
{&milvuspb.RoleEntity{Name: "role3"}, "role number reached the limit"},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{Entity: test.inEntity})
assert.Error(t, err)
})
}
t.Run("role has existed", func(t *testing.T) {
err := mt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{Entity: &milvuspb.RoleEntity{Name: "role1"}})
assert.Error(t, err)
assert.True(t, errors.Is(err, errRoleAlreadyExists))
})
{
mockCata := mocks.NewRootCoordCatalog(t)
mockCata.On("ListRole",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock list role"))
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{Entity: &milvuspb.RoleEntity{Name: "role1"}})
assert.Error(t, err)
}
}
func TestRbacAlterRoleDescription(t *testing.T) {
mt := generateMetaTable(t)
roleName := "role" + funcutil.RandomString(10)
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "old description",
})
require.NoError(t, err)
err = mt.AlterRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "new description",
})
require.NoError(t, err)
roles, err := mt.SelectRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: roleName}, false)
require.NoError(t, err)
require.Len(t, roles, 1)
assert.Equal(t, "new description", roles[0].GetRole().GetDescription())
err = mt.AlterRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: "role_not_exist",
Description: "ignored",
})
require.ErrorIs(t, err, errRoleNotExists)
err = mt.AlterRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: util.RoleAdmin,
Description: "ignored",
})
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrPrivilegeNotPermitted)
err = mt.CheckIfAlterRole(context.TODO(), &milvuspb.AlterRoleRequest{
RoleName: util.RolePublic,
Description: "ignored",
})
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrPrivilegeNotPermitted)
}
func TestRbacAlterRoleDescriptionErrors(t *testing.T) {
ctx := context.TODO()
t.Run("check empty role name", func(t *testing.T) {
mockCata := mocks.NewRootCoordCatalog(t)
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.CheckIfAlterRole(ctx, &milvuspb.AlterRoleRequest{
RoleName: "",
Description: "description",
})
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("check list role error", func(t *testing.T) {
targetErr := errors.New("mock list role error")
roleName := "role_check_list_error"
mockCata := mocks.NewRootCoordCatalog(t)
mockCata.EXPECT().ListRole(
mock.Anything,
util.DefaultTenant,
mock.MatchedBy(func(entity *milvuspb.RoleEntity) bool {
return entity.GetName() == roleName
}),
false,
).Return(nil, targetErr)
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.CheckIfAlterRole(ctx, &milvuspb.AlterRoleRequest{
RoleName: roleName,
Description: "description",
})
assert.ErrorIs(t, err, targetErr)
})
t.Run("alter empty role name", func(t *testing.T) {
mockCata := mocks.NewRootCoordCatalog(t)
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.AlterRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{
Name: "",
Description: "description",
})
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("alter list role error", func(t *testing.T) {
targetErr := errors.New("mock list role error")
roleName := "role_alter_list_error"
mockCata := mocks.NewRootCoordCatalog(t)
mockCata.EXPECT().ListRole(
mock.Anything,
util.DefaultTenant,
mock.MatchedBy(func(entity *milvuspb.RoleEntity) bool {
return entity.GetName() == roleName
}),
false,
).Return(nil, targetErr)
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.AlterRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "description",
})
assert.ErrorIs(t, err, targetErr)
})
t.Run("alter catalog error", func(t *testing.T) {
targetErr := errors.New("mock alter role error")
roleName := "role_alter_catalog_error"
mockCata := mocks.NewRootCoordCatalog(t)
mockCata.EXPECT().ListRole(
mock.Anything,
util.DefaultTenant,
mock.MatchedBy(func(entity *milvuspb.RoleEntity) bool {
return entity.GetName() == roleName
}),
false,
).Return([]*milvuspb.RoleResult{{Role: &milvuspb.RoleEntity{Name: roleName}}}, nil)
mockCata.EXPECT().AlterRole(
mock.Anything,
util.DefaultTenant,
mock.MatchedBy(func(entity *milvuspb.RoleEntity) bool {
return entity.GetName() == roleName && entity.GetDescription() == "description"
}),
).Return(targetErr)
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.AlterRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "description",
})
assert.ErrorIs(t, err, targetErr)
})
}
func TestRbacCreateRoleToleratesMalformedStoredRoleValue(t *testing.T) {
ctx := context.TODO()
mt := generateMetaTable(t)
catalog := mt.catalog.(*rootcoord.Catalog)
require.NoError(t, catalog.CreateRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: "existing_role"}))
require.NoError(t, catalog.Txn.Save(ctx, rootcoord.RolePrefix+"/malformed_role", "{"))
err := mt.CheckIfCreateRole(ctx, &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{Name: "new_role"},
})
require.NoError(t, err)
}
func TestRbacRoleDescriptionLengthLimit(t *testing.T) {
mt := generateMetaTable(t)
paramtable.Get().Save(Params.ProxyCfg.MaxRoleDescriptionLength.Key, "4")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleDescriptionLength.Key)
err := mt.CheckIfCreateRole(context.TODO(), &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{
Name: "role_desc_limit_create",
Description: "12345",
},
})
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
err = mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role_desc_limit_alter"})
require.NoError(t, err)
err = mt.CheckIfAlterRole(context.TODO(), &milvuspb.AlterRoleRequest{
RoleName: "role_desc_limit_alter",
Description: "12345",
})
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
err = mt.CheckIfRBACRestorable(context.TODO(), &milvuspb.RestoreRBACMetaRequest{
RBACMeta: &milvuspb.RBACMeta{
Roles: []*milvuspb.RoleEntity{
{
Name: "role_desc_limit_restore",
Description: "12345",
},
},
},
})
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
}
func TestRbacRoleDescriptionApplyPathSkipsLengthLimit(t *testing.T) {
t.Run("create role apply path", func(t *testing.T) {
mt := generateMetaTable(t)
roleName := "role_desc_apply_create"
paramtable.Get().Save(Params.ProxyCfg.MaxRoleDescriptionLength.Key, "4")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleDescriptionLength.Key)
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "12345",
})
require.NoError(t, err)
roles, err := mt.SelectRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: roleName}, false)
require.NoError(t, err)
require.Len(t, roles, 1)
assert.Equal(t, "12345", roles[0].GetRole().GetDescription())
})
t.Run("alter role apply path", func(t *testing.T) {
mt := generateMetaTable(t)
roleName := "role_desc_apply_alter"
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: roleName})
require.NoError(t, err)
paramtable.Get().Save(Params.ProxyCfg.MaxRoleDescriptionLength.Key, "4")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleDescriptionLength.Key)
err = mt.AlterRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{
Name: roleName,
Description: "12345",
})
require.NoError(t, err)
roles, err := mt.SelectRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: roleName}, false)
require.NoError(t, err)
require.Len(t, roles, 1)
assert.Equal(t, "12345", roles[0].GetRole().GetDescription())
})
}
func TestRbacDropRole(t *testing.T) {
mt := generateMetaTable(t)
// drop a exist role
roleExist := "role" + funcutil.RandomString(10)
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: roleExist})
require.NoError(t, err)
err = mt.CheckIfDropRole(context.TODO(), &milvuspb.DropRoleRequest{RoleName: roleExist})
require.NoError(t, err)
err = mt.DropRole(context.TODO(), util.DefaultTenant, roleExist)
require.NoError(t, err)
// idempotency
mt.DropRole(context.TODO(), util.DefaultTenant, roleExist)
require.NoError(t, err)
// drop a not exist role
err = mt.CheckIfDropRole(context.TODO(), &milvuspb.DropRoleRequest{RoleName: "role_not_exist"})
require.ErrorIs(t, err, errRoleNotExists)
}
func TestRbacOperateRole(t *testing.T) {
mt := generateMetaTable(t)
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
require.NoError(t, err)
tests := []struct {
description string
user string
role string
oType milvuspb.OperateUserRoleType
}{
{"empty user", "", "role1", milvuspb.OperateUserRoleType_AddUserToRole},
{"empty role", "user1", "", milvuspb.OperateUserRoleType_AddUserToRole},
{"invalid type", "user1", "role1", milvuspb.OperateUserRoleType(100)},
{"remove not exist pair", "user1", "role2", milvuspb.OperateUserRoleType_RemoveUserFromRole},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.CheckIfOperateUserRole(context.TODO(), &milvuspb.OperateUserRoleRequest{
Username: test.user,
RoleName: test.role,
Type: test.oType,
})
assert.Error(t, err)
})
}
}
func TestRbacSelect(t *testing.T) {
mt := generateMetaTable(t)
roles := []string{"role1", "role2", "role3"}
userRoles := map[string][]string{
"user1": {"role1"},
"user2": {"role1", "role2"},
"user3": {"role1", "role3"},
"user4": {},
}
for _, role := range roles {
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
require.NoError(t, err)
}
for user, rs := range userRoles {
err := mt.catalog.AlterCredential(context.TODO(), &model.Credential{
Username: user,
Tenant: util.DefaultTenant,
})
require.NoError(t, err)
for _, r := range rs {
err := mt.OperateUserRole(
context.TODO(),
util.DefaultTenant,
&milvuspb.UserEntity{Name: user},
&milvuspb.RoleEntity{Name: r},
milvuspb.OperateUserRoleType_AddUserToRole)
require.NoError(t, err)
}
}
tests := []struct {
isValid bool
description string
inEntity *milvuspb.UserEntity
includeRoleInfo bool
expectedOutLength int
}{
{true, "no user entitiy, no role info", nil, false, 4},
{true, "no user entitiy, with role info", nil, true, 4},
{false, "not exist user", &milvuspb.UserEntity{Name: "not_exist"}, false, 0},
{true, "user1, no role info", &milvuspb.UserEntity{Name: "user1"}, false, 1},
{true, "user1, with role info", &milvuspb.UserEntity{Name: "user1"}, true, 1},
{true, "user2, no role info", &milvuspb.UserEntity{Name: "user2"}, false, 1},
{true, "user2, with role info", &milvuspb.UserEntity{Name: "user2"}, true, 1},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
res, err := mt.SelectUser(context.TODO(), util.DefaultTenant, test.inEntity, test.includeRoleInfo)
if test.isValid {
assert.NoError(t, err)
assert.Equal(t, test.expectedOutLength, len(res))
if test.includeRoleInfo {
u := res[0].GetUser().GetName()
roles, ok := userRoles[u]
assert.True(t, ok)
assert.Equal(t, len(roles), len(res[0].GetRoles()))
}
} else {
assert.Error(t, err)
}
})
}
testRoles := []struct {
isValid bool
description string
inEntity *milvuspb.RoleEntity
includeUserInfo bool
expectedOutLength int
}{
{true, "no role entitiy, no user info", nil, false, 3},
{true, "no role entitiy, with user info", nil, true, 3},
{false, "not exist role", &milvuspb.RoleEntity{Name: "not_exist"}, false, 0},
{true, "role1, no user info", &milvuspb.RoleEntity{Name: "role1"}, false, 1},
{true, "role1, with user info", &milvuspb.RoleEntity{Name: "role1"}, true, 1},
{true, "role2, no user info", &milvuspb.RoleEntity{Name: "role2"}, false, 1},
{true, "role2, with user info", &milvuspb.RoleEntity{Name: "role2"}, true, 1},
}
for _, test := range testRoles {
t.Run(test.description, func(t *testing.T) {
res, err := mt.SelectRole(context.TODO(), util.DefaultTenant, test.inEntity, test.includeUserInfo)
if test.isValid {
assert.NoError(t, err)
assert.Equal(t, test.expectedOutLength, len(res))
} else {
assert.Error(t, err)
}
})
}
}
func TestRbacOperatePrivilege(t *testing.T) {
mt := generateMetaTable(t)
tests := []struct {
description string
entity *milvuspb.GrantEntity
oType milvuspb.OperatePrivilegeType
}{
{"empty objectName", &milvuspb.GrantEntity{ObjectName: ""}, milvuspb.OperatePrivilegeType_Grant},
{"nil Object", &milvuspb.GrantEntity{
Object: nil,
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"empty Object name", &milvuspb.GrantEntity{
Object: &milvuspb.ObjectEntity{Name: ""},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"nil Role", &milvuspb.GrantEntity{
Role: nil,
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"empty Role name", &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: ""},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"nil grantor", &milvuspb.GrantEntity{
Grantor: nil,
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"nil grantor privilege", &milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
Privilege: nil,
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"empty grantor privilege name", &milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
Privilege: &milvuspb.PrivilegeEntity{Name: ""},
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"nil grantor user", &milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
User: nil,
Privilege: &milvuspb.PrivilegeEntity{Name: "privilege_name"},
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"empty grantor user name", &milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: ""},
Privilege: &milvuspb.PrivilegeEntity{Name: "privilege_name"},
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType_Grant},
{"invalid operateType", &milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: "user_name"},
Privilege: &milvuspb.PrivilegeEntity{Name: "privilege_name"},
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}, milvuspb.OperatePrivilegeType(-1)},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.OperatePrivilege(context.TODO(), util.DefaultTenant, test.entity, test.oType)
assert.Error(t, err)
})
}
validEntity := milvuspb.GrantEntity{
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: "user_name"},
Privilege: &milvuspb.PrivilegeEntity{Name: "privilege_name"},
},
Role: &milvuspb.RoleEntity{Name: "role_name"},
Object: &milvuspb.ObjectEntity{Name: "obj_name"},
ObjectName: "obj_name",
}
err := mt.OperatePrivilege(context.TODO(), util.DefaultTenant, &validEntity, milvuspb.OperatePrivilegeType_Grant)
assert.NoError(t, err)
}
func TestRbacSelectGrant(t *testing.T) {
mt := generateMetaTable(t)
tests := []struct {
description string
isValid bool
entity *milvuspb.GrantEntity
}{
{"nil Entity", false, nil},
{"nil entity Role", false, &milvuspb.GrantEntity{
Role: nil,
}},
{"empty entity Role name", false, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: ""},
}},
{"valid", true, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: "role"},
}},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
entities, err := mt.SelectGrant(context.TODO(), util.DefaultTenant, test.entity)
if test.isValid {
assert.NoError(t, err)
assert.Equal(t, 0, len(entities))
} else {
assert.Error(t, err)
}
})
}
}
func TestRbacDropGrant(t *testing.T) {
mt := generateMetaTable(t)
tests := []struct {
description string
isValid bool
role *milvuspb.RoleEntity
}{
{"nil role", false, nil},
{"empty Role name", false, &milvuspb.RoleEntity{Name: ""}},
{"valid", true, &milvuspb.RoleEntity{Name: "role"}},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.DropGrant(context.TODO(), util.DefaultTenant, test.role)
if test.isValid {
assert.NoError(t, err)
} else {
assert.Error(t, err)
}
})
}
}
func TestRbacListPolicy(t *testing.T) {
mt := generateMetaTable(t)
policies, err := mt.ListPolicy(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
assert.Empty(t, policies)
userRoles, err := mt.ListUserRole(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
assert.Equal(t, 0, len(userRoles))
}
func TestMetaTable_getCollectionByIDInternal(t *testing.T) {
t.Run("failed to get from catalog", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByID",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock GetCollectionByID"))
meta := &MetaTable{
catalog: catalog,
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
}
ctx := context.Background()
_, err := meta.getCollectionByIDInternal(ctx, util.DefaultDBName, 100, 101, false)
assert.Error(t, err)
})
t.Run("collection not available", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByID",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(&model.Collection{State: pb.CollectionState_CollectionDropped}, nil)
meta := &MetaTable{
catalog: catalog,
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
}
ctx := context.Background()
_, err := meta.getCollectionByIDInternal(ctx, util.DefaultDBName, 100, 101, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
coll, err := meta.getCollectionByIDInternal(ctx, util.DefaultDBName, 100, 101, true)
assert.NoError(t, err)
assert.False(t, coll.Available())
})
t.Run("normal case, filter unavailable partitions", func(t *testing.T) {
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
},
},
},
}
ctx := context.Background()
coll, err := meta.getCollectionByIDInternal(ctx, "", 100, 101, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
assert.Equal(t, Params.CommonCfg.DefaultPartitionName.GetValue(), coll.Partitions[0].PartitionName)
})
t.Run("get latest version", func(t *testing.T) {
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
},
},
},
}
ctx := context.Background()
coll, err := meta.getCollectionByIDInternal(ctx, "", 100, typeutil.MaxTimestamp, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
assert.Equal(t, Params.CommonCfg.DefaultPartitionName.GetValue(), coll.Partitions[0].PartitionName)
})
t.Run("UpdateTimestamp > ts triggers catalog fallback (time-travel correctness)", func(t *testing.T) {
// Regression test for the bug fix in getCollectionByIDInternal:
// cache invalidation was changed from CreateTime to UpdateTimestamp.
// Scenario: collection created at T=50, schema altered at T=100.
// A time-travel query at ts=80 (50 < 80 < 100) must NOT use the in-memory
// cache (which holds the post-alteration schema) — it must fall back to catalog.
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByID",
mock.Anything,
mock.Anything,
uint64(80),
int64(100),
).Return(&model.Collection{
State: pb.CollectionState_CollectionCreated,
CreateTime: 50,
UpdateTimestamp: 100,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
}, nil)
meta := &MetaTable{
catalog: catalog,
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 50,
UpdateTimestamp: 100, // schema was altered at T=100
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
},
},
}
ctx := context.Background()
// ts=80 is between CreateTime(50) and UpdateTimestamp(100):
// UpdateTimestamp(100) > ts(80) → cache bypass → catalog must be called.
coll, err := meta.getCollectionByIDInternal(ctx, util.DefaultDBName, 100, 80, false)
assert.NoError(t, err)
assert.NotNil(t, coll)
catalog.AssertCalled(t, "GetCollectionByID", mock.Anything, mock.Anything, uint64(80), int64(100))
})
t.Run("UpdateTimestamp <= ts uses in-memory cache (no catalog call)", func(t *testing.T) {
// ts=150 >= UpdateTimestamp(100) → the in-memory cache is fresh enough → no catalog call.
catalog := mocks.NewRootCoordCatalog(t)
// No expectations set — testify mock will fail if GetCollectionByID is called.
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 50,
UpdateTimestamp: 100,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
},
},
}
ctx := context.Background()
coll, err := meta.getCollectionByIDInternal(ctx, "", 100, 150, false)
assert.NoError(t, err)
assert.NotNil(t, coll)
catalog.AssertNotCalled(t, "GetCollectionByID")
})
}
func TestMetaTable_GetCollectionByName(t *testing.T) {
t.Run("db not found", func(t *testing.T) {
meta := &MetaTable{
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{},
},
},
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
}
ctx := context.Background()
_, err := meta.GetCollectionByName(ctx, "not_exist", "name", 101, false)
assert.Error(t, err)
})
t.Run("get by alias", func(t *testing.T) {
meta := &MetaTable{
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
},
},
},
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
}
meta.aliases.insert(util.DefaultDBName, "alias", 100)
ctx := context.Background()
coll, err := meta.GetCollectionByName(ctx, "", "alias", 101, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
assert.Equal(t, Params.CommonCfg.DefaultPartitionName.GetValue(), coll.Partitions[0].PartitionName)
})
t.Run("get by name", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
},
},
},
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
}
meta.names.insert(util.DefaultDBName, "name", 100)
ctx := context.Background()
coll, err := meta.GetCollectionByName(ctx, "", "name", 101, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
assert.Equal(t, Params.CommonCfg.DefaultPartitionName.GetValue(), coll.Partitions[0].PartitionName)
})
t.Run("failed to get from catalog", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock GetCollectionByName"))
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
ctx := context.Background()
_, err := meta.GetCollectionByName(ctx, util.DefaultDBName, "name", 101, false)
assert.Error(t, err)
})
t.Run("collection not available", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(&model.Collection{State: pb.CollectionState_CollectionDropped}, nil)
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
ctx := context.Background()
_, err := meta.GetCollectionByName(ctx, util.DefaultDBName, "name", 101, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
})
t.Run("normal case, filter unavailable partitions", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(&model.Collection{
State: pb.CollectionState_CollectionCreated,
CreateTime: 99,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
{PartitionID: 22, PartitionName: "dropped", State: pb.PartitionState_PartitionDropped},
},
}, nil)
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
ctx := context.Background()
coll, err := meta.GetCollectionByName(ctx, util.DefaultDBName, "name", 101, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
assert.Equal(t, UniqueID(11), coll.Partitions[0].PartitionID)
assert.Equal(t, Params.CommonCfg.DefaultPartitionName.GetValue(), coll.Partitions[0].PartitionName)
})
t.Run("get latest version", func(t *testing.T) {
ctx := context.Background()
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
}
_, err := meta.GetCollectionByName(ctx, "", "not_exist", typeutil.MaxTimestamp, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
})
}
func TestMetaTable_AddCollectionDoesNotBlockReadersDuringCatalogCreate(t *testing.T) {
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
t.Cleanup(channel.ResetStaticPChannelStatsManager)
catalog := mocks.NewRootCoordCatalog(t)
enteredCatalog := make(chan struct{})
unblockCatalog := make(chan struct{})
catalog.EXPECT().
CreateCollection(mock.Anything, mock.Anything, mock.Anything).
RunAndReturn(func(context.Context, *model.Collection, uint64) error {
close(enteredCatalog)
<-unblockCatalog
return nil
}).
Once()
existingColl := &model.Collection{
CollectionID: 100,
DBID: util.DefaultDBID,
DBName: util.DefaultDBName,
Name: "existing",
State: pb.CollectionState_CollectionCreated,
Partitions: []*model.Partition{
{PartitionID: 10, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
}
meta := &MetaTable{
catalog: catalog,
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
existingColl.CollectionID: existingColl,
},
partitionName2ID: map[int64]map[string]int64{},
names: newNameDb(),
aliases: newNameDb(),
fileResourceRefCnt: map[int64]int{},
}
meta.names.insert(util.DefaultDBName, existingColl.Name, existingColl.CollectionID)
addErr := make(chan error, 1)
go func() {
addErr <- meta.AddCollection(context.Background(), &model.Collection{
CollectionID: 101,
DBID: util.DefaultDBID,
DBName: util.DefaultDBName,
Name: "creating",
State: pb.CollectionState_CollectionCreated,
ShardsNum: 1,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
})
}()
select {
case <-enteredCatalog:
case <-time.After(time.Second):
require.FailNow(t, "AddCollection did not enter catalog CreateCollection")
}
type readResponse struct {
collection *model.Collection
err error
}
readResult := make(chan readResponse, 1)
go func() {
coll, err := meta.GetCollectionByName(context.Background(), util.DefaultDBName, existingColl.Name, typeutil.MaxTimestamp, false)
readResult <- readResponse{collection: coll, err: err}
}()
select {
case result := <-readResult:
require.NoError(t, result.err)
require.Equal(t, existingColl.CollectionID, result.collection.CollectionID)
case <-time.After(time.Second):
close(unblockCatalog)
require.NoError(t, <-addErr)
require.FailNow(t, "GetCollectionByName blocked while catalog CreateCollection was in progress")
}
close(unblockCatalog)
require.NoError(t, <-addErr)
}
func TestMetaTable_AddCollectionRechecksAfterCatalogCreate(t *testing.T) {
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
t.Cleanup(channel.ResetStaticPChannelStatsManager)
catalog := mocks.NewRootCoordCatalog(t)
enteredCatalog := make(chan struct{})
unblockCatalog := make(chan struct{})
catalog.EXPECT().
CreateCollection(mock.Anything, mock.Anything, mock.Anything).
RunAndReturn(func(context.Context, *model.Collection, uint64) error {
close(enteredCatalog)
<-unblockCatalog
return nil
}).
Once()
coll := &model.Collection{
CollectionID: 101,
DBID: util.DefaultDBID,
DBName: util.DefaultDBName,
Name: "creating",
State: pb.CollectionState_CollectionCreated,
ShardsNum: 1,
Partitions: []*model.Partition{
{PartitionID: 11, PartitionName: Params.CommonCfg.DefaultPartitionName.GetValue(), State: pb.PartitionState_PartitionCreated},
},
}
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
partitionName2ID: map[int64]map[string]int64{},
names: newNameDb(),
aliases: newNameDb(),
fileResourceRefCnt: map[int64]int{},
fileResourceRefHolds: map[int64]map[int64]int{},
generalCnt: 7,
}
addErr := make(chan error, 1)
go func() {
addErr <- meta.AddCollection(context.Background(), coll)
}()
select {
case <-enteredCatalog:
case <-time.After(time.Second):
require.FailNow(t, "AddCollection did not enter catalog CreateCollection")
}
meta.ddLock.Lock()
meta.collID2Meta[coll.CollectionID] = coll.Clone()
meta.names.insert(coll.DBName, coll.Name, coll.CollectionID)
meta.ddLock.Unlock()
close(unblockCatalog)
require.NoError(t, <-addErr)
require.Equal(t, 7, meta.generalCnt)
require.NotContains(t, meta.partitionName2ID, coll.CollectionID)
}
/*
func TestMetaTable_AlterCollection(t *testing.T) {
t.Run("alter metastore fail", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything, // context.Context
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("error"))
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
}
ctx := context.Background()
err := meta.AlterCollection(ctx, nil, nil, 0, false, false)
assert.Error(t, err)
})
t.Run("new name is already an alias", func(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
Name: "old",
DBID: util.DefaultDBID,
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
meta.aliases.insert(util.DefaultDBName, "new", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("alter collection ok", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
}
ctx := context.Background()
oldColl := &model.Collection{CollectionID: 1}
newColl := &model.Collection{CollectionID: 1}
err := meta.AlterCollection(ctx, oldColl, newColl, 0, false, false)
assert.NoError(t, err)
assert.Equal(t, meta.collID2Meta[1], newColl)
})
}
*/
func TestMetaTable_DescribeAlias(t *testing.T) {
t.Run("metatable describe alias ok", func(t *testing.T) {
var collectionID int64 = 100
collectionName := "test_metatable_describe_alias"
aliasName := "a_alias"
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID: {
CollectionID: collectionID,
Name: collectionName,
},
},
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", collectionName, collectionID)
meta.aliases.insert("", aliasName, collectionID)
ctx := context.Background()
descCollectionName, err := meta.DescribeAlias(ctx, "", aliasName, 0)
assert.NoError(t, err)
assert.Equal(t, collectionName, descCollectionName)
})
t.Run("metatable describe not exist alias", func(t *testing.T) {
var collectionID int64 = 100
aliasName1 := "a_alias"
aliasName2 := "a_alias2"
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
meta.aliases.insert("", aliasName1, collectionID)
ctx := context.Background()
descCollectionName, err := meta.DescribeAlias(ctx, "", aliasName2, 0)
assert.Error(t, err)
assert.Equal(t, "", descCollectionName)
})
t.Run("metatable describe not exist database", func(t *testing.T) {
aliasName := "a_alias"
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
ctx := context.Background()
descCollectionName, err := meta.DescribeAlias(ctx, "", aliasName, 0)
assert.Error(t, err)
assert.Equal(t, "", descCollectionName)
})
t.Run("metatable describe alias fail", func(t *testing.T) {
var collectionID int64 = 100
collectionName := "test_metatable_describe_alias"
aliasName := "a_alias"
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", collectionName, collectionID)
meta.aliases.insert("", aliasName, collectionID)
ctx := context.Background()
_, err := meta.DescribeAlias(ctx, "", aliasName, 0)
assert.Error(t, err)
})
t.Run("metatable describe alias dropped collection", func(t *testing.T) {
var collectionID int64 = 100
collectionName := "test_metatable_describe_alias"
aliasName := "a_alias"
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID: {
CollectionID: collectionID,
Name: collectionName,
},
},
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", collectionName, collectionID)
meta.aliases.insert("", aliasName, collectionID)
ctx := context.Background()
meta.collID2Meta[collectionID] = &model.Collection{State: pb.CollectionState_CollectionDropped}
alias, err := meta.DescribeAlias(ctx, "", aliasName, 0)
assert.Equal(t, "", alias)
assert.Error(t, err)
})
}
func TestMetaTable_ListAliases(t *testing.T) {
t.Run("metatable list alias ok", func(t *testing.T) {
var collectionID1 int64 = 101
collectionName1 := "test_metatable_list_alias1"
aliasName1 := "a_alias"
var collectionID2 int64 = 102
collectionName2 := "test_metatable_list_alias2"
aliasName2 := "a_alias2"
var collectionID3 int64 = 103
collectionName3 := "test_metatable_list_alias3"
aliasName3 := "a_alias3"
aliasName4 := "a_alias4"
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID1: {
CollectionID: collectionID1,
Name: collectionName1,
},
collectionID1: {
CollectionID: collectionID2,
Name: collectionName2,
},
},
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", collectionName1, collectionID1)
meta.names.insert("", collectionName2, collectionID2)
meta.names.insert("db2", collectionName3, collectionID3)
meta.aliases.insert("", aliasName1, collectionID1)
meta.aliases.insert("", aliasName2, collectionID2)
meta.aliases.insert("db2", aliasName3, collectionID3)
meta.aliases.insert("db2", aliasName4, collectionID3)
meta.collID2Meta[collectionID1] = &model.Collection{State: pb.CollectionState_CollectionCreated}
meta.collID2Meta[collectionID2] = &model.Collection{State: pb.CollectionState_CollectionCreated}
meta.collID2Meta[collectionID3] = &model.Collection{State: pb.CollectionState_CollectionCreated}
ctx := context.Background()
aliases, err := meta.ListAliases(ctx, "", "", 0)
assert.NoError(t, err)
assert.Equal(t, 2, len(aliases))
aliases2, err := meta.ListAliases(ctx, "", collectionName1, 0)
assert.NoError(t, err)
assert.Equal(t, 1, len(aliases2))
aliases3, err := meta.ListAliases(ctx, "db2", "", 0)
assert.NoError(t, err)
assert.Equal(t, 2, len(aliases3))
aliases4, err := meta.ListAliases(ctx, "db2", collectionName3, 0)
assert.NoError(t, err)
assert.Equal(t, 2, len(aliases4))
})
t.Run("metatable list alias in not exist database", func(t *testing.T) {
aliasName := "a_alias"
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
ctx := context.Background()
aliases, err := meta.ListAliases(ctx, "", aliasName, 0)
assert.Error(t, err)
assert.Equal(t, 0, len(aliases))
})
t.Run("metatable list alias error", func(t *testing.T) {
var collectionID1 int64 = 101
collectionName1 := "test_metatable_list_alias1"
aliasName1 := "a_alias"
var collectionID2 int64 = 102
collectionName2 := "test_metatable_list_alias2"
aliasName2 := "a_alias2"
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID1: {
CollectionID: collectionID1,
Name: collectionName1,
},
collectionID1: {
CollectionID: collectionID2,
Name: collectionName2,
},
},
names: newNameDb(),
aliases: newNameDb(),
}
meta.aliases.insert("", aliasName1, collectionID1)
meta.aliases.insert("", aliasName2, collectionID2)
ctx := context.Background()
_, err := meta.ListAliases(ctx, "", collectionName1, 0)
assert.Error(t, err)
})
t.Run("metatable list alias Dropping collection", func(t *testing.T) {
ctx := context.Background()
var collectionID1 int64 = 101
collectionName1 := "test_metatable_list_alias1"
aliasName1 := "a_alias"
var collectionID2 int64 = 102
collectionName2 := "test_metatable_list_alias2"
aliasName2 := "a_alias2"
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID1: {
CollectionID: collectionID1,
Name: collectionName1,
},
collectionID1: {
CollectionID: collectionID2,
Name: collectionName2,
},
},
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", collectionName1, collectionID1)
meta.names.insert("", collectionName2, collectionID2)
meta.aliases.insert("", aliasName1, collectionID1)
meta.aliases.insert("", aliasName2, collectionID2)
meta.collID2Meta[collectionID1] = &model.Collection{State: pb.CollectionState_CollectionCreated}
meta.collID2Meta[collectionID2] = &model.Collection{State: pb.CollectionState_CollectionDropped}
aliases, err := meta.ListAliases(ctx, "", "", 0)
assert.NoError(t, err)
assert.Equal(t, 1, len(aliases))
aliases2, err := meta.ListAliases(ctx, "", collectionName1, 0)
assert.NoError(t, err)
assert.Equal(t, 1, len(aliases2))
})
}
func Test_filterUnavailable(t *testing.T) {
coll := &model.Collection{}
nPartition := 10
nAvailablePartition := 0
for i := 0; i < nPartition; i++ {
partition := &model.Partition{
State: pb.PartitionState_PartitionDropping,
}
if rand.Int()%2 == 0 {
partition.State = pb.PartitionState_PartitionCreated
nAvailablePartition++
}
coll.Partitions = append(coll.Partitions, partition)
}
clone := filterUnavailablePartition(coll)
assert.Equal(t, nAvailablePartition, len(clone.Partitions))
for _, p := range clone.Partitions {
assert.True(t, p.Available())
}
}
func TestMetaTable_getLatestCollectionByIDInternal(t *testing.T) {
t.Run("not exist", func(t *testing.T) {
ctx := context.Background()
mt := &MetaTable{collID2Meta: nil}
_, err := mt.getLatestCollectionByIDInternal(ctx, 100, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
})
t.Run("nil case", func(t *testing.T) {
ctx := context.Background()
mt := &MetaTable{collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: nil,
}}
_, err := mt.getLatestCollectionByIDInternal(ctx, 100, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
})
t.Run("unavailable", func(t *testing.T) {
ctx := context.Background()
mt := &MetaTable{collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {State: pb.CollectionState_CollectionDropping},
}}
_, err := mt.getLatestCollectionByIDInternal(ctx, 100, false)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrCollectionNotFound)
coll, err := mt.getLatestCollectionByIDInternal(ctx, 100, true)
assert.NoError(t, err)
assert.False(t, coll.Available())
})
t.Run("normal case", func(t *testing.T) {
ctx := context.Background()
mt := &MetaTable{collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
State: pb.CollectionState_CollectionCreated,
Partitions: []*model.Partition{
{State: pb.PartitionState_PartitionCreated},
{State: pb.PartitionState_PartitionDropping},
},
},
}}
coll, err := mt.getLatestCollectionByIDInternal(ctx, 100, false)
assert.NoError(t, err)
assert.Equal(t, 1, len(coll.Partitions))
})
}
func TestMetaTable_RemoveCollection(t *testing.T) {
t.Run("catalog error", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropCollection",
mock.Anything, // context.Context
mock.Anything, // model.Collection
mock.AnythingOfType("uint64"),
).Return(errors.New("error mock DropCollection"))
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
CollectionID: 100,
DBID: int64(100),
},
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
ctx := context.Background()
err := meta.RemoveCollection(ctx, 100, 9999)
assert.Error(t, err)
meta.collID2Meta[100].State = pb.CollectionState_CollectionDropping
err = meta.RemoveCollection(ctx, 100, 9999)
assert.Error(t, err)
})
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropCollection",
mock.Anything, // context.Context
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(),
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)
meta.names.insert("", "alias1", 100)
meta.names.insert("", "alias2", 100)
ctx := context.Background()
err := meta.RemoveCollection(ctx, 100, 9999)
assert.NoError(t, err)
})
}
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(),
fileResourceRefCnt: make(map[int64]int),
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_DropCollection_GrantCleanup(t *testing.T) {
t.Run("grant cleanup on drop", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything,
).Return(nil)
catalog.On("DeleteGrantByCollectionName",
mock.Anything, mock.Anything, "testdb", "collection",
).Return(nil)
meta := &MetaTable{
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {Name: "collection", DBID: 1, State: pb.CollectionState_CollectionCreated},
},
dbName2Meta: map[string]*model.Database{
"testdb": {ID: 1, Name: "testdb"},
},
fileResourceRefCnt: make(map[int64]int),
}
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
meta.names.insert("testdb", "collection", 100)
ctx := context.Background()
err := meta.DropCollection(ctx, 100, 9999)
assert.NoError(t, err)
catalog.AssertCalled(t, "DeleteGrantByCollectionName", mock.Anything, mock.Anything, "testdb", "collection")
})
t.Run("grant cleanup best-effort on drop", func(t *testing.T) {
// When DeleteGrantByCollectionName fails, DropCollection should still succeed
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything,
).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", DBID: 1, State: pb.CollectionState_CollectionCreated},
},
dbName2Meta: map[string]*model.Database{
"default": {ID: 1, Name: "default"},
},
fileResourceRefCnt: make(map[int64]int),
}
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
meta.names.insert("default", "collection", 100)
ctx := context.Background()
err := meta.DropCollection(ctx, 100, 9999)
assert.NoError(t, err)
})
}
func buildAlterCollectionSchemaResult(collectionID int64, schema *schemapb.CollectionSchema, timetick uint64) message.BroadcastResultAlterCollectionMessageV2 {
controlChannel := funcutil.GetControlChannel("test")
raw := message.NewAlterCollectionMessageBuilderV2().
WithHeader(&messagespb.AlterCollectionMessageHeader{
CollectionId: collectionID,
UpdateMask: &fieldmaskpb.FieldMask{
Paths: []string{message.FieldMaskCollectionSchema},
},
}).
WithBody(&messagespb.AlterCollectionMessageBody{
Updates: &messagespb.AlterCollectionMessageUpdates{
Schema: schema,
},
}).
WithBroadcast([]string{controlChannel}).
MustBuildBroadcast()
return message.BroadcastResultAlterCollectionMessageV2{
Message: message.MustAsBroadcastAlterCollectionMessageV2(raw),
Results: map[string]*message.AppendResult{
controlChannel: {TimeTick: timetick},
},
}
}
func TestMetaTableAlterCollectionFileResourceRefCnt(t *testing.T) {
const (
collectionID = int64(100)
resourceID = int64(10)
oldResource = int64(20)
)
newMeta := func(oldIDs []int64, refCnt map[int64]int) (*MetaTable, *mocks.RootCoordCatalog) {
catalog := mocks.NewRootCoordCatalog(t)
meta := &MetaTable{
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
collectionID: {
CollectionID: collectionID,
Name: "collection",
DBName: "db",
DBID: 1,
State: pb.CollectionState_CollectionCreated,
FileResourceIds: oldIDs,
},
},
fileResourceID2Meta: map[int64]*internalpb.FileResourceInfo{
resourceID: {Id: resourceID, Name: "dict", Path: "dict.txt"},
oldResource: {Id: oldResource, Name: "old_dict", Path: "old_dict.txt"},
},
fileResourceRefCnt: refCnt,
fileResourceRefHolds: map[int64]map[int64]int{
999: {resourceID: 1},
},
}
meta.names.insert("db", "collection", collectionID)
return meta, catalog
}
buildSchema := func(ids ...int64) *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: "collection",
Version: 2,
FileResourceIds: ids,
}
}
t.Run("consume request path reservation", func(t *testing.T) {
meta, catalog := newMeta(nil, map[int64]int{})
require.NoError(t, reserveAlterCollectionFileResourceRefs(meta, collectionID, []int64{resourceID}))
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, true).Return(nil).Once()
err := meta.AlterCollection(context.Background(), buildAlterCollectionSchemaResult(collectionID, buildSchema(resourceID), 100))
require.NoError(t, err)
require.Equal(t, 1, meta.fileResourceRefCnt[resourceID])
require.ElementsMatch(t, []int64{resourceID}, meta.collID2Meta[collectionID].FileResourceIds)
})
t.Run("rollback request path reservation", func(t *testing.T) {
meta, _ := newMeta(nil, map[int64]int{})
require.NoError(t, reserveAlterCollectionFileResourceRefs(meta, collectionID, []int64{resourceID}))
rollbackAlterCollectionFileResourceRefs(context.Background(), meta, collectionID, []int64{resourceID})
require.Equal(t, 0, meta.fileResourceRefCnt[resourceID])
require.NotContains(t, meta.fileResourceRefHolds, collectionID)
})
t.Run("replay or replicated task adds refCnt without reservation", func(t *testing.T) {
meta, catalog := newMeta(nil, map[int64]int{resourceID: 1})
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, true).Return(nil).Once()
err := meta.AlterCollection(context.Background(), buildAlterCollectionSchemaResult(collectionID, buildSchema(resourceID), 100))
require.NoError(t, err)
require.Equal(t, 2, meta.fileResourceRefCnt[resourceID])
require.ElementsMatch(t, []int64{resourceID}, meta.collID2Meta[collectionID].FileResourceIds)
require.Equal(t, 1, meta.fileResourceRefHolds[999][resourceID])
})
t.Run("recovery hold prevents double increment for pending alter", func(t *testing.T) {
meta, catalog := newMeta(nil, map[int64]int{})
meta.RecoverFileResourceRefCnt(map[int64][]int64{collectionID: {resourceID}})
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, true).Return(nil).Once()
err := meta.AlterCollection(context.Background(), buildAlterCollectionSchemaResult(collectionID, buildSchema(resourceID), 100))
require.NoError(t, err)
require.Equal(t, 1, meta.fileResourceRefCnt[resourceID])
require.NotContains(t, meta.fileResourceRefHolds, collectionID)
require.ElementsMatch(t, []int64{resourceID}, meta.collID2Meta[collectionID].FileResourceIds)
})
t.Run("replace resource decrements removed and consumes added reservation", func(t *testing.T) {
meta, catalog := newMeta([]int64{oldResource}, map[int64]int{oldResource: 1})
require.NoError(t, reserveAlterCollectionFileResourceRefs(meta, collectionID, []int64{resourceID}))
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, true).Return(nil).Once()
err := meta.AlterCollection(context.Background(), buildAlterCollectionSchemaResult(collectionID, buildSchema(resourceID), 100))
require.NoError(t, err)
require.Equal(t, 0, meta.fileResourceRefCnt[oldResource])
require.Equal(t, 1, meta.fileResourceRefCnt[resourceID])
require.ElementsMatch(t, []int64{resourceID}, meta.collID2Meta[collectionID].FileResourceIds)
})
t.Run("missing added resource fails before catalog update", func(t *testing.T) {
meta, catalog := newMeta(nil, map[int64]int{})
delete(meta.fileResourceID2Meta, resourceID)
err := meta.AlterCollection(context.Background(), buildAlterCollectionSchemaResult(collectionID, buildSchema(resourceID), 100))
require.Error(t, err)
catalog.AssertNotCalled(t, "AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
require.Empty(t, meta.collID2Meta[collectionID].FileResourceIds)
})
}
func TestMetaTableGetFileResources(t *testing.T) {
resourcesByID := map[int64]*internalpb.FileResourceInfo{
1: {Id: 1, Name: "resource-1", Path: "files/1"},
2: {Id: 2, Name: "resource-2", Path: "files/2"},
}
meta := &MetaTable{fileResourceID2Meta: resourcesByID}
t.Run("return requested resources in order", func(t *testing.T) {
resources, err := meta.GetFileResources(context.Background(), 2, 1)
require.NoError(t, err)
require.Equal(t, []*internalpb.FileResourceInfo{resourcesByID[2], resourcesByID[1]}, resources)
require.NotSame(t, resourcesByID[2], resources[0])
require.NotSame(t, resourcesByID[1], resources[1])
})
t.Run("missing resource is internal error", func(t *testing.T) {
resources, err := meta.GetFileResources(context.Background(), 3)
require.Nil(t, resources)
require.ErrorIs(t, err, merr.ErrServiceInternal)
})
}
func TestMetaTableHasFileResource(t *testing.T) {
meta := &MetaTable{fileResourceID2Meta: make(map[int64]*internalpb.FileResourceInfo)}
require.False(t, meta.HasFileResource(context.Background()))
meta.fileResourceID2Meta[1] = &internalpb.FileResourceInfo{Id: 1}
require.True(t, meta.HasFileResource(context.Background()))
}
func TestMetaTable_DropPartition_CopyOnWrite(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
originalPart := &model.Partition{
PartitionID: 100,
PartitionName: "p1",
CollectionID: 100,
State: pb.PartitionState_PartitionCreated,
}
catalog.On("AlterPartition",
mock.Anything,
int64(10),
originalPart,
mock.MatchedBy(func(newPart *model.Partition) bool {
return newPart != nil &&
newPart != originalPart &&
newPart.PartitionID == originalPart.PartitionID &&
newPart.PartitionName == originalPart.PartitionName &&
newPart.CollectionID == originalPart.CollectionID &&
newPart.State == pb.PartitionState_PartitionDropping
}),
metastore.MODIFY,
uint64(9999),
).Return(nil).Once()
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
CollectionID: 100,
DBID: 10,
State: pb.CollectionState_CollectionCreated,
Partitions: []*model.Partition{originalPart},
},
},
partitionName2ID: map[int64]map[string]int64{
100: {"p1": 100},
},
}
snapshot, err := meta.GetCollectionByID(context.Background(), "", 100, typeutil.MaxTimestamp, true)
require.NoError(t, err)
require.Same(t, originalPart, snapshot.Partitions[0])
err = meta.DropPartition(context.Background(), 100, 100, 9999)
require.NoError(t, err)
require.Same(t, originalPart, snapshot.Partitions[0])
assert.Equal(t, pb.PartitionState_PartitionCreated, snapshot.Partitions[0].State)
require.Len(t, meta.collID2Meta[100].Partitions, 1)
assert.NotSame(t, originalPart, meta.collID2Meta[100].Partitions[0])
assert.Equal(t, pb.PartitionState_PartitionDropping, meta.collID2Meta[100].Partitions[0].State)
assert.Equal(t, pb.PartitionState_PartitionCreated, originalPart.State)
assert.NotContains(t, meta.partitionName2ID[100], "p1")
}
func TestMetaTable_RemovePartition(t *testing.T) {
t.Run("catalog error", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropPartition",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.AnythingOfType("uint64"),
).Return(errors.New("error mock AlterPartition"))
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
CollectionID: 100,
DBID: int64(100),
Partitions: []*model.Partition{
{PartitionID: 100, State: pb.PartitionState_PartitionCreated},
},
},
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
ctx := context.Background()
err := meta.RemovePartition(ctx, 100, 100, 9999)
assert.Error(t, err)
meta.collID2Meta[100].Partitions[0].State = pb.PartitionState_PartitionDropping
err = meta.RemovePartition(ctx, 100, 100, 9999)
assert.Error(t, err)
})
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropPartition",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.AnythingOfType("uint64"),
).Return(nil)
meta := &MetaTable{
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
Name: "collection",
Partitions: []*model.Partition{
{PartitionID: 100, State: pb.PartitionState_PartitionDropping},
},
},
},
}
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
meta.names.insert("", "collection", 100)
meta.names.insert("", "alias1", 100)
meta.names.insert("", "alias2", 100)
ctx := context.Background()
err := meta.RemovePartition(ctx, 100, 100, 9999)
assert.NoError(t, err)
})
}
func TestMetaTable_reload(t *testing.T) {
createMetaTableFn := func(catalogOpts ...func(*mocks.RootCoordCatalog)) *MetaTable {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("ListDatabases",
mock.Anything,
mock.Anything,
).Return(make([]*model.Database, 0), nil)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("ListRLSPolicies",
mock.Anything,
mock.Anything,
).Return([]*model.RLSPolicy{}, nil).Maybe()
catalog.On("ListRLSPrincipals",
mock.Anything,
mock.Anything,
).Return([]*model.RLSPrincipal{}, nil).Maybe()
tso := mocktso.NewAllocator(t)
tso.On("GenerateTSO",
mock.Anything,
).Return(uint64(1), nil)
for _, opt := range catalogOpts {
opt(catalog)
}
return &MetaTable{
ctx: context.Background(),
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
tsoAllocator: tso,
}
}
t.Run("list db fail", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("ListDatabases",
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock ListDatabases"))
meta := &MetaTable{catalog: catalog}
err := meta.reload()
assert.Error(t, err)
assert.Empty(t, meta.collID2Meta)
assert.Equal(t, 0, len(meta.names.listDB()))
assert.Equal(t, 0, len(meta.aliases.listDB()))
})
t.Run("failed to list collections", func(t *testing.T) {
meta := createMetaTableFn(func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListCollections",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock ListCollections"))
})
err := meta.reload()
assert.Error(t, err)
assert.Empty(t, meta.collID2Meta)
})
t.Run("failed to list aliases", func(t *testing.T) {
meta := createMetaTableFn(
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListCollections",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Collection{{CollectionID: 100, Name: "test"}},
nil)
},
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListAliases",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, errors.New("error mock ListAliases"))
},
)
err := meta.reload()
assert.Error(t, err)
})
t.Run("no collections and default db doesn't exists", func(t *testing.T) {
meta := createMetaTableFn(
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListCollections",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(make([]*model.Collection, 0), nil)
},
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListAliases",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Alias{},
nil)
},
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListFileResource",
mock.Anything,
).Return(nil, uint64(0), nil)
},
)
channel.ResetStaticPChannelStatsManager()
err := meta.reload()
assert.NoError(t, err)
assert.NoError(t, err)
assert.Equal(t, 0, len(meta.collID2Meta))
assert.Equal(t, 1, len(meta.names.listDB()))
assert.True(t, meta.names.exist("default"))
assert.Equal(t, 1, len(meta.aliases.listDB()))
assert.True(t, meta.aliases.exist("default"))
})
t.Run("no collections and default db already exists", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("ListDatabases",
mock.Anything,
mock.Anything,
).Return([]*model.Database{model.NewDefaultDatabase(nil)}, nil)
catalog.On("ListCollections",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Collection{{CollectionID: 100, Name: "test", State: pb.CollectionState_CollectionCreated}},
nil)
catalog.On("ListAliases",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Alias{{Name: "alias", CollectionID: 100}},
nil)
catalog.On("ListRLSPolicies",
mock.Anything,
mock.Anything,
).Return([]*model.RLSPolicy{}, nil).Maybe()
catalog.On("ListRLSPrincipals",
mock.Anything,
mock.Anything,
).Return([]*model.RLSPrincipal{}, nil).Maybe()
catalog.On("ListFileResource",
mock.Anything,
).Return(nil, uint64(0), nil)
meta := &MetaTable{ctx: context.Background(), catalog: catalog}
channel.ResetStaticPChannelStatsManager()
err := meta.reload()
assert.NoError(t, err)
assert.NoError(t, err)
assert.Equal(t, 1, len(meta.collID2Meta))
assert.Equal(t, 1, len(meta.names.listDB()))
assert.True(t, meta.names.exist(util.DefaultDBName))
assert.Equal(t, 1, len(meta.aliases.listDB()))
assert.True(t, meta.aliases.exist(util.DefaultDBName))
})
t.Run("ok", func(t *testing.T) {
meta := createMetaTableFn(
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListCollections",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Collection{{CollectionID: 100, Name: "test", State: pb.CollectionState_CollectionCreated}},
nil)
},
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListAliases",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(
[]*model.Alias{{Name: "alias", CollectionID: 100}},
nil)
},
func(catalog *mocks.RootCoordCatalog) {
catalog.On("ListFileResource",
mock.Anything,
).Return(nil, uint64(0), nil)
},
)
channel.ResetStaticPChannelStatsManager()
err := meta.reload()
assert.NoError(t, err)
assert.Equal(t, 1, len(meta.collID2Meta))
assert.Equal(t, 1, len(meta.names.listDB()))
assert.True(t, meta.names.exist(util.DefaultDBName))
assert.Equal(t, 1, len(meta.aliases.listDB()))
assert.True(t, meta.aliases.exist(util.DefaultDBName))
colls, err := meta.names.listCollectionID(util.DefaultDBName)
assert.NoError(t, err)
assert.Equal(t, 1, len(colls))
assert.Equal(t, int64(100), colls[0])
colls, err = meta.aliases.listCollectionID(util.DefaultDBName)
assert.NoError(t, err)
assert.Equal(t, 1, len(colls))
assert.Equal(t, int64(100), colls[0])
})
}
func TestMetaTable_ListAllAvailCollections(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: {
ID: util.DefaultDBID,
},
"db2": {
ID: 11,
},
"db3": {
ID: 2,
},
"db4": {
ID: 1111,
},
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
111: {
CollectionID: 111,
DBID: 1111,
State: pb.CollectionState_CollectionDropped,
},
2: {
CollectionID: 2,
DBID: 11,
State: pb.CollectionState_CollectionCreated,
},
3: {
CollectionID: 3,
DBID: 11,
State: pb.CollectionState_CollectionCreated,
},
4: {
CollectionID: 4,
DBID: 2,
State: pb.CollectionState_CollectionCreated,
},
5: {
CollectionID: 5,
DBID: util.NonDBID,
State: pb.CollectionState_CollectionCreated,
},
},
}
ret := meta.ListAllAvailCollections(context.TODO())
assert.Equal(t, 4, len(ret))
db0, ok := ret[util.DefaultDBID]
assert.True(t, ok)
assert.Equal(t, 1, len(db0))
db1, ok := ret[11]
assert.True(t, ok)
assert.Equal(t, 2, len(db1))
db2, ok2 := ret[2]
assert.True(t, ok2)
assert.Equal(t, 1, len(db2))
db3, ok := ret[1111]
assert.True(t, ok)
assert.Equal(t, 0, len(db3))
}
func TestMetaTable_GetAvailableCollectionCount(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: {ID: util.DefaultDBID},
"db2": {ID: 11},
"db3": {ID: 2},
"db4": {ID: 1111},
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
111: {
CollectionID: 111,
DBID: 1111,
State: pb.CollectionState_CollectionDropped,
},
2: {
CollectionID: 2,
DBID: 11,
State: pb.CollectionState_CollectionCreated,
},
3: {
CollectionID: 3,
DBID: 11,
State: pb.CollectionState_CollectionCreated,
},
4: {
CollectionID: 4,
DBID: 2,
State: pb.CollectionState_CollectionCreated,
},
5: {
CollectionID: 5,
DBID: util.NonDBID,
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.rebuildAvailableCollectionCountLocked()
dbCount, total, ok := meta.GetAvailableCollectionCount(context.TODO(), util.DefaultDBID)
assert.True(t, ok)
assert.Equal(t, 1, dbCount)
assert.Equal(t, 4, total)
dbCount, total, ok = meta.GetAvailableCollectionCount(context.TODO(), int64(11))
assert.True(t, ok)
assert.Equal(t, 2, dbCount)
assert.Equal(t, 4, total)
dbCount, total, ok = meta.GetAvailableCollectionCount(context.TODO(), int64(1111))
assert.True(t, ok)
assert.Equal(t, 0, dbCount)
assert.Equal(t, 4, total)
dbCount, total, ok = meta.GetAvailableCollectionCount(context.TODO(), int64(9999))
assert.False(t, ok)
assert.Equal(t, 0, dbCount)
assert.Equal(t, 4, total)
}
func TestMetaTable_AvailableCollectionCountTransitions(t *testing.T) {
ctx := context.Background()
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateCollection", mock.Anything, mock.Anything, mock.Anything).Return(nil)
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
catalog.On("AlterCollectionDB", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
catalog.On("DeleteGrantByCollectionName", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
catalog.On("MigrateGrantCollectionName", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
meta := &MetaTable{
catalog: catalog,
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: {ID: util.DefaultDBID},
"db2": {ID: 11},
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{},
partitionName2ID: map[int64]map[string]int64{},
names: newNameDb(),
aliases: newNameDb(),
fileResourceRefCnt: map[int64]int{},
}
meta.names.createDbIfNotExist(util.DefaultDBName)
meta.names.createDbIfNotExist("db2")
meta.rebuildAvailableCollectionCountLocked()
err := meta.AddCollection(ctx, &model.Collection{
CollectionID: 100,
DBID: util.DefaultDBID,
DBName: util.DefaultDBName,
Name: "c1",
State: pb.CollectionState_CollectionCreated,
ShardsNum: 1,
Partitions: []*model.Partition{
{PartitionID: 10, PartitionName: "_default", State: pb.PartitionState_PartitionCreated},
},
})
require.NoError(t, err)
dbCount, total, ok := meta.GetAvailableCollectionCount(ctx, util.DefaultDBID)
require.True(t, ok)
assert.Equal(t, 1, dbCount)
assert.Equal(t, 1, total)
result := message.BroadcastResultAlterCollectionMessageV2{
Message: message.MustAsBroadcastAlterCollectionMessageV2(
message.NewAlterCollectionMessageBuilderV2().
WithHeader(&message.AlterCollectionMessageHeader{
CollectionId: 100,
UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{message.FieldMaskDB}},
}).
WithBody(&message.AlterCollectionMessageBody{
Updates: &message.AlterCollectionMessageUpdates{
DbId: 11,
DbName: "db2",
},
}).
WithBroadcast([]string{funcutil.GetControlChannel("by-dev-rootcoord-dml_1")}).
MustBuildBroadcast(),
),
Results: map[string]*message.AppendResult{
funcutil.GetControlChannel("by-dev-rootcoord-dml_1"): {TimeTick: 200},
},
}
err = meta.AlterCollection(ctx, result)
require.NoError(t, err)
dbCount, total, ok = meta.GetAvailableCollectionCount(ctx, util.DefaultDBID)
require.True(t, ok)
assert.Equal(t, 0, dbCount)
assert.Equal(t, 1, total)
dbCount, total, ok = meta.GetAvailableCollectionCount(ctx, int64(11))
require.True(t, ok)
assert.Equal(t, 1, dbCount)
assert.Equal(t, 1, total)
err = meta.DropCollection(ctx, 100, 300)
require.NoError(t, err)
dbCount, total, ok = meta.GetAvailableCollectionCount(ctx, int64(11))
require.True(t, ok)
assert.Equal(t, 0, dbCount)
assert.Equal(t, 0, total)
}
func TestMetaTable_DropCollectionDBLookupFailureIsRetryable(t *testing.T) {
ctx := context.Background()
channel.ResetStaticPChannelStatsManager()
channel.RecoverPChannelStatsManager([]string{})
t.Cleanup(channel.ResetStaticPChannelStatsManager)
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil).
Once()
catalog.On("DeleteGrantByCollectionName", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil).
Once()
const dbID = int64(11)
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
CollectionID: 100,
DBID: dbID,
DBName: "db2",
Name: "c1",
State: pb.CollectionState_CollectionCreated,
ShardsNum: 1,
Partitions: []*model.Partition{
{PartitionID: 10, PartitionName: "_default", State: pb.PartitionState_PartitionCreated},
},
},
},
dbName2Meta: map[string]*model.Database{},
fileResourceRefCnt: map[int64]int{},
generalCnt: 1,
}
meta.rebuildAvailableCollectionCountLocked()
err := meta.DropCollection(ctx, 100, 300)
require.ErrorIs(t, err, merr.ErrDatabaseNotFound)
assert.Equal(t, pb.CollectionState_CollectionCreated, meta.collID2Meta[100].State)
assert.Equal(t, 1, meta.generalCnt)
dbCount, total, ok := meta.GetAvailableCollectionCount(ctx, dbID)
require.True(t, ok)
assert.Equal(t, 1, dbCount)
assert.Equal(t, 1, total)
catalog.AssertNotCalled(t, "AlterCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
meta.dbName2Meta["db2"] = &model.Database{ID: dbID, Name: "db2"}
require.NoError(t, meta.DropCollection(ctx, 100, 300))
assert.Equal(t, pb.CollectionState_CollectionDropping, meta.collID2Meta[100].State)
assert.Equal(t, 0, meta.generalCnt)
dbCount, total, ok = meta.GetAvailableCollectionCount(ctx, dbID)
require.True(t, ok)
assert.Equal(t, 0, dbCount)
assert.Equal(t, 0, total)
}
func TestMetaTable_AddPartition(t *testing.T) {
t.Run("collection not available", func(t *testing.T) {
meta := &MetaTable{}
err := meta.AddPartition(context.TODO(), &model.Partition{CollectionID: 100})
assert.Error(t, err)
})
t.Run("add not-created partition", func(t *testing.T) {
meta := &MetaTable{
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {
Name: "test",
CollectionID: 100,
},
},
}
err := meta.AddPartition(context.TODO(), &model.Partition{CollectionID: 100, State: pb.PartitionState_PartitionDropping})
assert.Error(t, err)
})
t.Run("failed to create partition", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreatePartition",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("error mock CreatePartition"))
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {Name: "test", CollectionID: 100},
},
}
err := meta.AddPartition(context.TODO(), &model.Partition{CollectionID: 100, State: pb.PartitionState_PartitionCreated})
assert.Error(t, err)
})
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreatePartition",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
meta := &MetaTable{
catalog: catalog,
collID2Meta: map[typeutil.UniqueID]*model.Collection{
100: {Name: "test", CollectionID: 100},
},
partitionName2ID: make(map[int64]map[string]int64),
}
err := meta.AddPartition(context.TODO(), &model.Partition{CollectionID: 100, State: pb.PartitionState_PartitionCreated})
assert.NoError(t, err)
})
}
/*
func TestMetaTable_RenameCollection(t *testing.T) {
t.Run("unsupported use a alias to rename collection", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", "alias", 1)
err := meta.RenameCollection(context.TODO(), "", "alias", "", "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("target db doesn't exist", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
err := meta.RenameCollection(context.TODO(), "", "non-exists", "non-exists", "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("collection name not exist", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
err := meta.RenameCollection(context.TODO(), "", "non-exists", "", "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("collection id not exist", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
meta.names.insert("", "old", 1)
err := meta.RenameCollection(context.TODO(), "", "old", "", "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("new collection name already exist-1", func(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
util.DefaultDBID: {
CollectionID: 1,
Name: "old",
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "old", 1000)
assert.Error(t, err)
})
t.Run("new collection name already exist-2", func(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
Name: "old",
State: pb.CollectionState_CollectionCreated,
},
2: {
CollectionID: 2,
Name: "new",
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
meta.names.insert(util.DefaultDBName, "new", 2)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", 1000)
assert.Error(t, err)
})
t.Run("alter collection fail", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
Name: "old",
DBID: util.DefaultDBID,
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("new name is already an alias", func(t *testing.T) {
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
Name: "old",
DBID: util.DefaultDBID,
State: pb.CollectionState_CollectionCreated,
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
meta.aliases.insert(util.DefaultDBName, "new", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", typeutil.MaxTimestamp)
assert.Error(t, err)
})
t.Run("alter collection ok", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollectionDB",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, merr.WrapErrCollectionNotFound("error"))
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
Name: "old",
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", 1000)
assert.NoError(t, err)
id, ok := meta.names.get(util.DefaultDBName, "new")
assert.True(t, ok)
assert.Equal(t, int64(1), id)
coll, ok := meta.collID2Meta[1]
assert.True(t, ok)
assert.Equal(t, "new", coll.Name)
})
t.Run("rename collection ok", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, merr.WrapErrCollectionNotFound("error"))
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
DBID: 1,
Name: "old",
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", 1000)
assert.NoError(t, err)
id, ok := meta.names.get(util.DefaultDBName, "new")
assert.True(t, ok)
assert.Equal(t, int64(1), id)
coll, ok := meta.collID2Meta[1]
assert.True(t, ok)
assert.Equal(t, "new", coll.Name)
})
t.Run("rename collection rename collName with db encryption on", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterCollection",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("GetCollectionByName",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil, merr.WrapErrCollectionNotFound("error"))
// Create database with encryption enabled
encryptedDBProperties := []*commonpb.KeyValuePair{
{Key: "cipher.enabled", Value: "true"},
}
encryptedDB := model.NewDatabase(1, util.DefaultDBName, pb.DatabaseState_DatabaseCreated, encryptedDBProperties)
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: encryptedDB,
},
catalog: catalog,
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
DBID: 1,
Name: "old",
},
},
}
meta.names.insert(util.DefaultDBName, "old", 1)
// Should succeed when renaming within the same encrypted database
err := meta.RenameCollection(context.TODO(), util.DefaultDBName, "old", util.DefaultDBName, "new", 1000)
assert.NoError(t, err)
// Verify the collection name was updated
id, ok := meta.names.get(util.DefaultDBName, "new")
assert.True(t, ok)
assert.Equal(t, int64(1), id)
coll, ok := meta.collID2Meta[1]
assert.True(t, ok)
assert.Equal(t, "new", coll.Name)
})
}
*/
func TestMetaTable_CreateDatabase(t *testing.T) {
db := model.NewDatabase(1, "exist", pb.DatabaseState_DatabaseCreated, nil)
t.Run("database already exist", func(t *testing.T) {
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
dbName2Meta: make(map[string]*model.Database),
}
meta.names.insert("exist", "collection", 100)
err := meta.CheckIfDatabaseCreatable(context.TODO(), &milvuspb.CreateDatabaseRequest{
DbName: "exist",
})
assert.Error(t, err)
})
t.Run("database not persistent", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("error mock CreateDatabase"))
meta := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
err := meta.CreateDatabase(context.TODO(), db, 10000)
assert.Error(t, err)
})
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
meta := &MetaTable{
dbName2Meta: make(map[string]*model.Database),
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
err := meta.CheckIfDatabaseCreatable(context.TODO(), &milvuspb.CreateDatabaseRequest{
DbName: "exist",
})
assert.NoError(t, err)
err = meta.CreateDatabase(context.TODO(), db, 10000)
assert.NoError(t, err)
assert.True(t, meta.names.exist("exist"))
assert.True(t, meta.aliases.exist("exist"))
assert.True(t, meta.names.empty("exist"))
assert.True(t, meta.aliases.empty("exist"))
})
}
func TestCreateDefaultDb(t *testing.T) {
hookutil.InitTestCipher()
// Save original config and restore after test
originalDefaultKey := paramtable.GetCipherParams().DefaultRootKey.GetValue()
defer func() {
paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", originalDefaultKey)
}()
t.Run("default db without encryption when defaultKey is empty", func(t *testing.T) {
paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "")
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
tsoAllocator := mocktso.NewAllocator(t)
tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(100), nil)
meta := &MetaTable{
ctx: context.Background(),
dbName2Meta: make(map[string]*model.Database),
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
tsoAllocator: tsoAllocator,
}
err := meta.createDefaultDb()
assert.NoError(t, err)
// Verify default database was created
db, ok := meta.dbName2Meta[util.DefaultDBName]
assert.True(t, ok)
assert.Equal(t, util.DefaultDBName, db.Name)
// Verify no encryption properties
hasEncryption := hookutil.IsDBEncrypted(db.Properties)
assert.False(t, hasEncryption, "default DB should not be encrypted when defaultKey is empty")
})
t.Run("default db with encryption when defaultKey is set", func(t *testing.T) {
paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "default-test-key")
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
tsoAllocator := mocktso.NewAllocator(t)
tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(200), nil)
meta := &MetaTable{
ctx: context.Background(),
dbName2Meta: make(map[string]*model.Database),
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
tsoAllocator: tsoAllocator,
}
err := meta.createDefaultDb()
assert.NoError(t, err)
// Verify default database was created
db, ok := meta.dbName2Meta[util.DefaultDBName]
assert.True(t, ok)
assert.Equal(t, util.DefaultDBName, db.Name)
// Verify encryption properties are present
hasEzID := false
hasRootKey := false
for _, prop := range db.Properties {
if prop.Key == common.EncryptionEzIDKey {
hasEzID = true
assert.Equal(t, prop.GetValue(), "199")
}
if prop.Key == common.EncryptionRootKeyKey && prop.Value == "default-test-key" {
hasRootKey = true
}
}
assert.True(t, hasRootKey, "default DB should have root key when encrypted")
assert.True(t, hasEzID, "default DB should have ezID when encrypted")
})
t.Run("TSO allocation failure", func(t *testing.T) {
tsoAllocator := mocktso.NewAllocator(t)
tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(0), errors.New("TSO allocation failed"))
meta := &MetaTable{
ctx: context.Background(),
dbName2Meta: make(map[string]*model.Database),
tsoAllocator: tsoAllocator,
}
err := meta.createDefaultDb()
assert.Error(t, err)
assert.Contains(t, err.Error(), "TSO allocation failed")
})
t.Run("catalog CreateDatabase failure", func(t *testing.T) {
paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "")
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("CreateDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("catalog error"))
tsoAllocator := mocktso.NewAllocator(t)
tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(300), nil)
meta := &MetaTable{
ctx: context.Background(),
dbName2Meta: make(map[string]*model.Database),
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
tsoAllocator: tsoAllocator,
}
err := meta.createDefaultDb()
assert.Error(t, err)
assert.Contains(t, err.Error(), "catalog error")
})
}
func TestAlterDatabase(t *testing.T) {
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("AlterDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
db := model.NewDatabase(1, "db1", pb.DatabaseState_DatabaseCreated, nil)
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
"db1": db,
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
newDB := db.Clone()
db.Properties = []*commonpb.KeyValuePair{
{
Key: "key1",
Value: "value1",
},
}
err := meta.AlterDatabase(context.TODO(), newDB, typeutil.ZeroTimestamp)
assert.NoError(t, err)
})
t.Run("access catalog failed", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
mockErr := errors.New("access catalog failed")
catalog.On("AlterDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(mockErr)
db := model.NewDatabase(1, "db1", pb.DatabaseState_DatabaseCreated, nil)
meta := &MetaTable{
dbName2Meta: map[string]*model.Database{
"db1": db,
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
newDB := db.Clone()
db.Properties = []*commonpb.KeyValuePair{
{
Key: "key1",
Value: "value1",
},
}
err := meta.AlterDatabase(context.TODO(), newDB, typeutil.ZeroTimestamp)
assert.ErrorIs(t, err, mockErr)
})
}
func TestMetaTable_EmtpyDatabaseName(t *testing.T) {
t.Run("getDatabaseByNameInternal with empty db", func(t *testing.T) {
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: {ID: 1},
},
}
ret, err := mt.getDatabaseByNameInternal(context.TODO(), "", typeutil.MaxTimestamp)
assert.NoError(t, err)
assert.Equal(t, int64(1), ret.ID)
})
t.Run("getCollectionByNameInternal with empty db", func(t *testing.T) {
mt := &MetaTable{
aliases: newNameDb(),
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {CollectionID: 1},
},
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
},
}
mt.aliases.insert(util.DefaultDBName, "aliases", 1)
ret, err := mt.getCollectionByNameInternal(context.TODO(), "", "aliases", typeutil.MaxTimestamp, false)
assert.NoError(t, err)
assert.Equal(t, int64(1), ret.CollectionID)
})
t.Run("listCollectionFromCache with empty db", func(t *testing.T) {
mt := &MetaTable{
names: newNameDb(),
dbName2Meta: map[string]*model.Database{
util.DefaultDBName: model.NewDefaultDatabase(nil),
"db2": model.NewDatabase(2, "db2", pb.DatabaseState_DatabaseCreated, nil),
},
collID2Meta: map[typeutil.UniqueID]*model.Collection{
1: {
CollectionID: 1,
State: pb.CollectionState_CollectionCreated,
},
2: {
CollectionID: 2,
State: pb.CollectionState_CollectionDropping,
DBID: util.DefaultDBID,
},
3: {
CollectionID: 3,
State: pb.CollectionState_CollectionCreated,
DBID: 2,
},
},
}
ret, err := mt.listCollectionFromCache(context.TODO(), "none", false)
assert.Error(t, err)
assert.Nil(t, ret)
ret, err = mt.listCollectionFromCache(context.TODO(), "", false)
assert.NoError(t, err)
assert.Equal(t, 2, len(ret))
assert.Contains(t, []int64{1, 2}, ret[0].CollectionID)
assert.Contains(t, []int64{1, 2}, ret[1].CollectionID)
ret, err = mt.listCollectionFromCache(context.TODO(), "db2", false)
assert.NoError(t, err)
assert.Equal(t, 1, len(ret))
assert.Equal(t, int64(3), ret[0].CollectionID)
})
t.Run("CreateAlias with empty db", func(t *testing.T) {
mt := &MetaTable{
names: newNameDb(),
}
mt.names.insert(util.DefaultDBName, "name", 1)
err := mt.CheckIfAliasCreatable(context.TODO(), "", "name", "name")
assert.Error(t, err)
})
}
func TestMetaTable_DropDatabase(t *testing.T) {
t.Run("can't drop default database", func(t *testing.T) {
mt := &MetaTable{}
err := mt.CheckIfDatabaseDroppable(context.TODO(), &milvuspb.DropDatabaseRequest{
DbName: util.DefaultDBName,
})
assert.Error(t, err)
})
t.Run("database not exist", func(t *testing.T) {
mt := &MetaTable{
names: newNameDb(),
aliases: newNameDb(),
}
err := mt.CheckIfDatabaseDroppable(context.TODO(), &milvuspb.DropDatabaseRequest{
DbName: "not_exist",
})
assert.True(t, errors.Is(err, merr.ErrDatabaseNotFound))
})
t.Run("database not empty", func(t *testing.T) {
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_empty": model.NewDatabase(1, "not_empty", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
collID2Meta: map[int64]*model.Collection{
10000000: {
DBID: 1,
CollectionID: 10000000,
Name: "collection",
},
},
}
mt.names.insert("not_empty", "collection", 10000000)
err := mt.CheckIfDatabaseDroppable(context.TODO(), &milvuspb.DropDatabaseRequest{
DbName: "not_empty",
})
assert.Error(t, err)
})
t.Run("not commit", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("error mock DropDatabase"))
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
mt.names.createDbIfNotExist("not_commit")
mt.aliases.createDbIfNotExist("not_commit")
err := mt.CheckIfDatabaseDroppable(context.TODO(), &milvuspb.DropDatabaseRequest{
DbName: "not_commit",
})
assert.NoError(t, err)
err = mt.DropDatabase(context.TODO(), "not_commit", 10000)
assert.Error(t, err)
})
t.Run("normal case", func(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.On("DropDatabase",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
mt.names.createDbIfNotExist("not_commit")
mt.aliases.createDbIfNotExist("not_commit")
err := mt.CheckIfDatabaseDroppable(context.TODO(), &milvuspb.DropDatabaseRequest{
DbName: "not_commit",
})
assert.NoError(t, err)
err = mt.DropDatabase(context.TODO(), "not_commit", 10000)
assert.NoError(t, err)
assert.False(t, mt.names.exist("not_commit"))
assert.False(t, mt.aliases.exist("not_commit"))
})
}
func TestMetaTable_BackupRBAC(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(&milvuspb.RBACMeta{}, nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
_, err := mt.BackupRBAC(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
catalog.ExpectedCalls = nil
catalog.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(nil, errors.New("error mock BackupRBAC"))
_, err = mt.BackupRBAC(context.TODO(), util.DefaultTenant)
assert.Error(t, err)
}
func TestMetaTable_RestoreRBAC(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
err := mt.RestoreRBAC(context.TODO(), util.DefaultTenant, &milvuspb.RBACMeta{})
assert.NoError(t, err)
catalog.ExpectedCalls = nil
catalog.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error mock RestoreRBAC"))
err = mt.RestoreRBAC(context.TODO(), util.DefaultTenant, &milvuspb.RBACMeta{})
assert.Error(t, err)
}
func TestMetaTable_CheckIfRBACRestorable_Wildcard(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().ListRole(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil, nil)
catalog.EXPECT().ListPrivilegeGroups(mock.Anything).
Return(nil, nil)
catalog.EXPECT().ListUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil, nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
req := &milvuspb.RestoreRBACMetaRequest{
RBACMeta: &milvuspb.RBACMeta{
Roles: []*milvuspb.RoleEntity{{Name: "wildcard_role"}},
Grants: []*milvuspb.GrantEntity{
{
Role: &milvuspb.RoleEntity{Name: "wildcard_role"},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
ObjectName: util.AnyWord,
DbName: util.AnyWord,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: util.AnyWord},
},
},
},
},
}
assert.NoError(t, mt.CheckIfRBACRestorable(context.TODO(), req))
}
func TestMetaTable_PrivilegeGroup(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().ListPrivilegeGroups(mock.Anything).Return([]*milvuspb.PrivilegeGroupInfo{
{
GroupName: "pg1",
Privileges: []*milvuspb.PrivilegeEntity{{Name: "CreateCollection"}, {Name: "DescribeCollection"}},
},
}, nil)
catalog.EXPECT().ListRole(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil)
catalog.EXPECT().SavePrivilegeGroup(mock.Anything, mock.Anything).Return(nil)
catalog.EXPECT().DropPrivilegeGroup(mock.Anything, mock.Anything).Return(nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
err := mt.CheckIfPrivilegeGroupCreatable(context.TODO(), &milvuspb.CreatePrivilegeGroupRequest{
GroupName: "pg1",
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupCreatable(context.TODO(), &milvuspb.CreatePrivilegeGroupRequest{
GroupName: "",
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupCreatable(context.TODO(), &milvuspb.CreatePrivilegeGroupRequest{
GroupName: "Insert",
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupCreatable(context.TODO(), &milvuspb.CreatePrivilegeGroupRequest{
GroupName: "pg2",
})
assert.NoError(t, err)
err = mt.CreatePrivilegeGroup(context.TODO(), "pg1")
assert.NoError(t, err)
// idempotency
err = mt.CreatePrivilegeGroup(context.TODO(), "pg1")
assert.NoError(t, err)
err = mt.CheckIfPrivilegeGroupDropable(context.TODO(), &milvuspb.DropPrivilegeGroupRequest{
GroupName: "",
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupDropable(context.TODO(), &milvuspb.DropPrivilegeGroupRequest{
GroupName: "pg1",
})
assert.NoError(t, err)
err = mt.DropPrivilegeGroup(context.TODO(), "pg1")
assert.NoError(t, err)
err = mt.CheckIfPrivilegeGroupAlterable(context.TODO(), &milvuspb.OperatePrivilegeGroupRequest{
GroupName: "",
Privileges: []*milvuspb.PrivilegeEntity{},
Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup,
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupAlterable(context.TODO(), &milvuspb.OperatePrivilegeGroupRequest{
GroupName: "ClusterReadOnly",
Privileges: []*milvuspb.PrivilegeEntity{},
Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup,
})
assert.Error(t, err)
err = mt.CheckIfPrivilegeGroupAlterable(context.TODO(), &milvuspb.OperatePrivilegeGroupRequest{
GroupName: "pg3",
Privileges: []*milvuspb.PrivilegeEntity{},
Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup,
})
assert.Error(t, err)
_, err = mt.GetPrivilegeGroupRoles(context.TODO(), "")
assert.Error(t, err)
_, err = mt.ListPrivilegeGroups(context.TODO())
assert.NoError(t, err)
}
func TestMetaTable_TruncateCollection(t *testing.T) {
channel.ResetStaticPChannelStatsManager()
kv, _ := kvfactory.GetEtcdAndPath()
path := funcutil.RandomString(10) + "/meta"
catalogKV := etcdkv.NewEtcdKV(kv, path)
catalog := rootcoord.NewCatalog(catalogKV)
allocator := mocktso.NewAllocator(t)
allocator.EXPECT().GenerateTSO(mock.Anything).Return(1000, nil)
meta, err := NewMetaTable(context.Background(), catalog, allocator)
require.NoError(t, err)
err = meta.AddCollection(context.Background(), &model.Collection{
CollectionID: 1,
PhysicalChannelNames: []string{"pchannel1"},
VirtualChannelNames: []string{"vchannel1"},
State: pb.CollectionState_CollectionCreated,
DBID: util.DefaultDBID,
Properties: common.NewKeyValuePairs(map[string]string{}),
ShardInfos: map[string]*model.ShardInfo{
"vchannel1": {
VChannelName: "vchannel1",
PChannelName: "pchannel1",
LastTruncateTimeTick: 0,
},
},
})
require.NoError(t, err)
// begin truncate collection
err = meta.BeginTruncateCollection(context.Background(), 1)
require.NoError(t, err)
coll, err := meta.GetCollectionByID(context.Background(), util.DefaultDBName, 1, typeutil.MaxTimestamp, false)
require.NoError(t, err)
m := common.CloneKeyValuePairs(coll.Properties).ToMap()
require.Equal(t, "1", m[common.CollectionOnTruncatingKey])
require.Equal(t, uint64(0), coll.ShardInfos["vchannel1"].LastTruncateTimeTick)
// reload the meta
channel.ResetStaticPChannelStatsManager()
meta, err = NewMetaTable(context.Background(), catalog, allocator)
require.NoError(t, err)
coll, err = meta.GetCollectionByID(context.Background(), util.DefaultDBName, 1, typeutil.MaxTimestamp, false)
require.NoError(t, err)
m = common.CloneKeyValuePairs(coll.Properties).ToMap()
require.Equal(t, "1", m[common.CollectionOnTruncatingKey])
require.Equal(t, uint64(0), coll.ShardInfos["vchannel1"].LastTruncateTimeTick)
// remove the temp property
b := message.NewTruncateCollectionMessageBuilderV2().
WithHeader(&message.TruncateCollectionMessageHeader{
CollectionId: 1,
}).
WithBody(&message.TruncateCollectionMessageBody{}).
WithBroadcast(coll.VirtualChannelNames, message.OptBuildBroadcastAckSyncUp()).
MustBuildBroadcast()
meta.TruncateCollection(context.Background(), message.BroadcastResultTruncateCollectionMessageV2{
Message: message.MustAsBroadcastTruncateCollectionMessageV2(b),
Results: map[string]*message.AppendResult{
"vchannel1": {
TimeTick: 1000,
},
},
})
require.NoError(t, err)
coll, err = meta.GetCollectionByID(context.Background(), util.DefaultDBName, 1, typeutil.MaxTimestamp, false)
require.NoError(t, err)
m = common.CloneKeyValuePairs(coll.Properties).ToMap()
_, ok := m[common.CollectionOnTruncatingKey]
require.False(t, ok)
require.Equal(t, uint64(1000), coll.ShardInfos["vchannel1"].LastTruncateTimeTick)
// reload the meta again
channel.ResetStaticPChannelStatsManager()
meta, err = NewMetaTable(context.Background(), catalog, allocator)
require.NoError(t, err)
coll, err = meta.GetCollectionByID(context.Background(), util.DefaultDBName, 1, typeutil.MaxTimestamp, false)
require.NoError(t, err)
m = common.CloneKeyValuePairs(coll.Properties).ToMap()
_, ok = m[common.CollectionOnTruncatingKey]
require.False(t, ok)
require.Equal(t, 1, len(coll.ShardInfos))
require.Equal(t, uint64(1000), coll.ShardInfos["vchannel1"].LastTruncateTimeTick)
}
func TestMetaTableReloadNormalizesMaxFieldIDProperty(t *testing.T) {
channel.ResetStaticPChannelStatsManager()
kv, _ := kvfactory.GetEtcdAndPath()
path := funcutil.RandomString(10) + "/meta"
catalogKV := etcdkv.NewEtcdKV(kv, path)
catalog := rootcoord.NewCatalog(catalogKV)
allocator := mocktso.NewAllocator(t)
allocator.EXPECT().GenerateTSO(mock.Anything).Return(1000, nil)
meta, err := NewMetaTable(context.Background(), catalog, allocator)
require.NoError(t, err)
err = meta.AddCollection(context.Background(), &model.Collection{
CollectionID: 1,
DBID: util.DefaultDBID,
DBName: util.DefaultDBName,
Name: "test_reload_max_field_id",
PhysicalChannelNames: []string{"pchannel1"},
VirtualChannelNames: []string{"vchannel1"},
State: pb.CollectionState_CollectionCreated,
Fields: []*model.Field{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64},
{FieldID: 105, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
Properties: common.NewKeyValuePairs(map[string]string{
common.CollectionReplicaNumber: "1",
}),
ShardInfos: map[string]*model.ShardInfo{
"vchannel1": {
VChannelName: "vchannel1",
PChannelName: "pchannel1",
LastTruncateTimeTick: 0,
},
},
})
require.NoError(t, err)
channel.ResetStaticPChannelStatsManager()
meta, err = NewMetaTable(context.Background(), catalog, allocator)
require.NoError(t, err)
coll, err := meta.GetCollectionByID(context.Background(), util.DefaultDBName, 1, typeutil.MaxTimestamp, false)
require.NoError(t, err)
props := common.CloneKeyValuePairs(coll.Properties).ToMap()
require.Equal(t, "105", props[common.MaxFieldIDKey])
}
Reference in New Issue View Git Blame Copy Permalink