wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
2521 lines
92 KiB
Go
2521 lines
92 KiB
Go
// 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"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/samber/lo"
|
|
"golang.org/x/exp/maps"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
|
"github.com/milvus-io/milvus/internal/metastore"
|
|
"github.com/milvus-io/milvus/internal/metastore/model"
|
|
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
|
|
"github.com/milvus-io/milvus/internal/tso"
|
|
"github.com/milvus-io/milvus/internal/util/hookutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/metrics"
|
|
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
|
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/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/contextutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/crypto"
|
|
"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/rbacutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
var (
|
|
errIgnoredAlterAlias = errors.New("ignored alter alias") // alias already created on current collection, so it can be ignored.
|
|
errIgnoredAlterCollection = errors.New("ignored alter collection") // collection already created, so it can be ignored.
|
|
errIgnoredAlterDatabase = errors.New("ignored alter database") // database already created, so it can be ignored.
|
|
errIgnoredCreateCollection = errors.New("ignored create collection") // create collection with same schema, so it can be ignored.
|
|
errIgnoerdCreatePartition = errors.New("ignored create partition") // partition is already exist, so it can be ignored.
|
|
errIgnoredDropCollection = errors.New("ignored drop collection") // drop collection or database not found, so it can be ignored.
|
|
errIgnoredDropPartition = errors.New("ignored drop partition") // drop partition not found, so it can be ignored.
|
|
|
|
errAlterCollectionNotFound = errors.New("alter collection not found") // alter collection not found, so it can be ignored.
|
|
)
|
|
|
|
type MetaTableChecker interface {
|
|
RBACChecker
|
|
|
|
CheckIfDatabaseCreatable(ctx context.Context, req *milvuspb.CreateDatabaseRequest) error
|
|
CheckIfDatabaseDroppable(ctx context.Context, req *milvuspb.DropDatabaseRequest) error
|
|
|
|
CheckIfAliasCreatable(ctx context.Context, dbName string, alias string, collectionName string) error
|
|
CheckIfAliasAlterable(ctx context.Context, dbName string, alias string, collectionName string) error
|
|
CheckIfAliasDroppable(ctx context.Context, dbName string, alias string) error
|
|
}
|
|
|
|
//go:generate mockery --name=IMetaTable --structname=MockIMetaTable --output=./ --filename=mock_meta_table.go --with-expecter --inpackage
|
|
type IMetaTable interface {
|
|
MetaTableChecker
|
|
|
|
GetDatabaseByID(ctx context.Context, dbID int64, ts Timestamp) (*model.Database, error)
|
|
GetDatabaseByName(ctx context.Context, dbName string, ts Timestamp) (*model.Database, error)
|
|
CreateDatabase(ctx context.Context, db *model.Database, ts typeutil.Timestamp) error
|
|
DropDatabase(ctx context.Context, dbName string, ts typeutil.Timestamp) error
|
|
ListDatabases(ctx context.Context, ts typeutil.Timestamp) ([]*model.Database, error)
|
|
AlterDatabase(ctx context.Context, newDB *model.Database, ts typeutil.Timestamp) error
|
|
|
|
AddCollection(ctx context.Context, coll *model.Collection) error
|
|
DropCollection(ctx context.Context, collectionID UniqueID, ts Timestamp) error
|
|
RemoveCollection(ctx context.Context, collectionID UniqueID, ts Timestamp) error
|
|
// GetCollectionID retrieves the corresponding collectionID based on the collectionName.
|
|
// If the collection does not exist, it will return InvalidCollectionID.
|
|
// Please use the function with caution.
|
|
GetCollectionID(ctx context.Context, dbName string, collectionName string) UniqueID
|
|
GetCollectionByName(ctx context.Context, dbName string, collectionName string, ts Timestamp, allowUnavailable bool) (*model.Collection, error)
|
|
GetCollectionByID(ctx context.Context, dbName string, collectionID UniqueID, ts Timestamp, allowUnavailable bool) (*model.Collection, error)
|
|
GetCollectionByIDWithMaxTs(ctx context.Context, collectionID UniqueID) (*model.Collection, error)
|
|
ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error)
|
|
ListAllAvailCollections(ctx context.Context) map[int64][]int64
|
|
// ListAllAvailPartitions returns the partition ids of all available collections.
|
|
// The key of the map is the database id, and the value is a map of collection id to partition ids.
|
|
ListAllAvailPartitions(ctx context.Context) map[int64]map[int64][]int64
|
|
ListCollectionPhysicalChannels(ctx context.Context) map[typeutil.UniqueID][]string
|
|
GetCollectionVirtualChannels(ctx context.Context, colID int64) []string
|
|
GetPChannelInfo(ctx context.Context, pchannel string) *rootcoordpb.GetPChannelInfoResponse
|
|
AddPartition(ctx context.Context, partition *model.Partition) error
|
|
GetPartitionIDByName(collectionID int64, partitionName string) (int64, bool)
|
|
DropPartition(ctx context.Context, collectionID UniqueID, partitionID UniqueID, ts Timestamp) error
|
|
RemovePartition(ctx context.Context, collectionID UniqueID, partitionID UniqueID, ts Timestamp) error
|
|
|
|
// Alias
|
|
AlterAlias(ctx context.Context, result message.BroadcastResultAlterAliasMessageV2) error
|
|
DropAlias(ctx context.Context, result message.BroadcastResultDropAliasMessageV2) error
|
|
DescribeAlias(ctx context.Context, dbName string, alias string, ts Timestamp) (string, error)
|
|
ListAliases(ctx context.Context, dbName string, collectionName string, ts Timestamp) ([]string, error)
|
|
|
|
AlterCollection(ctx context.Context, result message.BroadcastResultAlterCollectionMessageV2) error
|
|
// Deprecated: will be removed in the 3.0 after implementing ack sync up semantic.
|
|
// It will be used to forbid the compaction of current collection when truncate collection operation is in progress.
|
|
BeginTruncateCollection(ctx context.Context, collectionID UniqueID) error
|
|
// TruncateCollection is called when the truncate collection message is acknowledged.
|
|
TruncateCollection(ctx context.Context, result message.BroadcastResultTruncateCollectionMessageV2) error
|
|
CheckIfCollectionRenamable(ctx context.Context, dbName string, oldName string, newDBName string, newName string) error
|
|
GetGeneralCount(ctx context.Context) int
|
|
|
|
// TODO: it'll be a big cost if we handle the time travel logic, since we should always list all aliases in catalog.
|
|
IsAlias(ctx context.Context, db, name string) bool
|
|
ListAliasesByID(ctx context.Context, collID UniqueID) []string
|
|
|
|
GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error)
|
|
InitCredential(ctx context.Context) error
|
|
DeleteCredential(ctx context.Context, result message.BroadcastResultDropUserMessageV2) error
|
|
AlterCredential(ctx context.Context, result message.BroadcastResultAlterUserMessageV2) error
|
|
ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error)
|
|
|
|
CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
|
|
AlterRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
|
|
DropRole(ctx context.Context, tenant string, roleName string) error
|
|
OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
|
|
SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
|
|
SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
|
|
OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
|
|
SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
|
|
DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error
|
|
ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error)
|
|
ListUserRole(ctx context.Context, tenant string) ([]string, error)
|
|
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
|
|
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
|
|
IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error)
|
|
CreatePrivilegeGroup(ctx context.Context, groupName string) error
|
|
DropPrivilegeGroup(ctx context.Context, groupName string) error
|
|
ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error)
|
|
OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
|
|
GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error)
|
|
|
|
AddFileResource(ctx context.Context, resource *internalpb.FileResourceInfo) error
|
|
RemoveFileResource(ctx context.Context, name string) (error, bool)
|
|
ListFileResource(ctx context.Context) ([]*internalpb.FileResourceInfo, uint64)
|
|
IncFileResourceRefCnt(ids []int64) error
|
|
DecFileResourceRefCnt(ids []int64)
|
|
RecoverFileResourceRefCnt(pendingCollections map[int64][]int64)
|
|
}
|
|
|
|
// MetaTable is a persistent meta set of all databases, collections and partitions.
|
|
type MetaTable struct {
|
|
ctx context.Context
|
|
catalog metastore.RootCoordCatalog
|
|
|
|
tsoAllocator tso.Allocator
|
|
|
|
dbName2Meta map[string]*model.Database // database name -> db meta
|
|
collID2Meta map[typeutil.UniqueID]*model.Collection // collection id -> collection meta
|
|
|
|
// partition name index: collectionID -> partitionName -> partitionID
|
|
partitionName2ID map[int64]map[string]int64
|
|
|
|
fileResourceName2Meta map[string]*internalpb.FileResourceInfo // file resource name -> file resource meta
|
|
fileResourceID2Meta map[int64]*internalpb.FileResourceInfo // file resource id -> file resource meta
|
|
fileResourceRefCnt map[int64]int // file resource id -> reference count
|
|
fileResourceVersion uint64
|
|
|
|
generalCnt int // sum of product of partition number and shard number
|
|
|
|
// collections *collectionDb
|
|
names *nameDb
|
|
aliases *nameDb
|
|
|
|
ddLock sync.RWMutex
|
|
permissionLock sync.RWMutex
|
|
}
|
|
|
|
// NewMetaTable creates a new MetaTable with specified catalog and allocator.
|
|
func NewMetaTable(ctx context.Context, catalog metastore.RootCoordCatalog, tsoAllocator tso.Allocator) (*MetaTable, error) {
|
|
mt := &MetaTable{
|
|
ctx: contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue()),
|
|
catalog: catalog,
|
|
tsoAllocator: tsoAllocator,
|
|
}
|
|
if err := mt.reload(); err != nil {
|
|
return nil, err
|
|
}
|
|
return mt, nil
|
|
}
|
|
|
|
func (mt *MetaTable) reload() error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
record := timerecord.NewTimeRecorder("rootcoord")
|
|
mt.dbName2Meta = make(map[string]*model.Database)
|
|
mt.collID2Meta = make(map[UniqueID]*model.Collection)
|
|
mt.partitionName2ID = make(map[int64]map[string]int64)
|
|
mt.fileResourceRefCnt = make(map[int64]int)
|
|
mt.names = newNameDb()
|
|
mt.aliases = newNameDb()
|
|
|
|
metrics.RootCoordNumOfCollections.Reset()
|
|
metrics.RootCoordNumOfPartitions.Reset()
|
|
metrics.RootCoordNumOfDatabases.Set(0)
|
|
|
|
// recover databases.
|
|
dbs, err := mt.catalog.ListDatabases(mt.ctx, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mlog.Info(mt.ctx, "recover databases", mlog.Int("num of dbs", len(dbs)))
|
|
for _, db := range dbs {
|
|
mt.dbName2Meta[db.Name] = db
|
|
}
|
|
dbNames := maps.Keys(mt.dbName2Meta)
|
|
// create default database.
|
|
if !funcutil.SliceContain(dbNames, util.DefaultDBName) {
|
|
if err := mt.createDefaultDb(); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
mt.names.createDbIfNotExist(util.DefaultDBName)
|
|
mt.aliases.createDbIfNotExist(util.DefaultDBName)
|
|
}
|
|
|
|
// in order to support backward compatibility with meta of the old version, it also
|
|
// needs to reload collections that have no database
|
|
if err := mt.reloadWithNonDatabase(); err != nil {
|
|
return err
|
|
}
|
|
|
|
// recover collections from db namespace
|
|
for dbName, db := range mt.dbName2Meta {
|
|
partitionNum := int64(0)
|
|
collectionNum := int64(0)
|
|
|
|
mt.names.createDbIfNotExist(dbName)
|
|
|
|
start := time.Now()
|
|
// TODO: async list collections to accelerate cases with multiple databases.
|
|
collections, err := mt.catalog.ListCollections(mt.ctx, db.ID, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, collection := range collections {
|
|
if collection.DBName != "" && collection.DBName != dbName {
|
|
mlog.Warn(mt.ctx,
|
|
"collection dbname is not correct, it will be fixed",
|
|
mlog.Int64("collection_id", collection.CollectionID),
|
|
mlog.String("db_name", dbName),
|
|
mlog.String("collection_name", collection.Name),
|
|
mlog.String("collection_dbname", collection.DBName),
|
|
)
|
|
}
|
|
collection.DBName = dbName // some collections may not have db name or its dbname is not correct, we should fix it here.
|
|
ensureCollectionMaxFieldIDProperty(collection)
|
|
mt.collID2Meta[collection.CollectionID] = collection
|
|
// Build partition name index
|
|
mt.partitionName2ID[collection.CollectionID] = make(map[string]int64)
|
|
for _, partition := range collection.Partitions {
|
|
if partition.Available() {
|
|
mt.partitionName2ID[collection.CollectionID][partition.PartitionName] = partition.PartitionID
|
|
}
|
|
}
|
|
if collection.Available() {
|
|
mt.names.insert(dbName, collection.Name, collection.CollectionID)
|
|
for _, fileResourceID := range collection.FileResourceIds {
|
|
mt.fileResourceRefCnt[fileResourceID]++
|
|
}
|
|
pn := collection.GetPartitionNum(true)
|
|
mt.generalCnt += pn * int(collection.ShardsNum)
|
|
collectionNum++
|
|
partitionNum += int64(pn)
|
|
}
|
|
}
|
|
|
|
metrics.RootCoordNumOfDatabases.Inc()
|
|
metrics.RootCoordNumOfCollections.WithLabelValues(dbName).Add(float64(collectionNum))
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Add(float64(partitionNum))
|
|
mlog.Info(mt.ctx, "collections recovered from db", mlog.String("db_name", dbName),
|
|
mlog.Int64("collection_num", collectionNum),
|
|
mlog.Int64("partition_num", partitionNum),
|
|
mlog.Duration("dur", time.Since(start)))
|
|
}
|
|
|
|
// recover aliases from db namespace
|
|
for dbName, db := range mt.dbName2Meta {
|
|
mt.aliases.createDbIfNotExist(dbName)
|
|
aliases, err := mt.catalog.ListAliases(mt.ctx, db.ID, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, alias := range aliases {
|
|
mt.aliases.insert(dbName, alias.Name, alias.CollectionID)
|
|
}
|
|
}
|
|
|
|
mlog.Info(mt.ctx, "rootcoord start to recover the channel stats for streaming coord balancer")
|
|
vchannels := make([]string, 0, len(mt.collID2Meta)*2)
|
|
for _, coll := range mt.collID2Meta {
|
|
if coll.Available() {
|
|
vchannels = append(vchannels, coll.VirtualChannelNames...)
|
|
}
|
|
}
|
|
channel.RecoverPChannelStatsManager(vchannels)
|
|
|
|
// reload file resources
|
|
resources, version, err := mt.catalog.ListFileResource(mt.ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
mt.fileResourceName2Meta = make(map[string]*internalpb.FileResourceInfo)
|
|
mt.fileResourceID2Meta = make(map[int64]*internalpb.FileResourceInfo)
|
|
for _, resource := range resources {
|
|
mt.fileResourceName2Meta[resource.Name] = resource
|
|
mt.fileResourceID2Meta[resource.Id] = resource
|
|
}
|
|
mt.fileResourceVersion = version
|
|
|
|
mlog.Info(mt.ctx, "RootCoord meta table reload done", mlog.Duration("duration", record.ElapseSpan()))
|
|
return nil
|
|
}
|
|
|
|
// insert into default database if the collections doesn't inside some database
|
|
func (mt *MetaTable) reloadWithNonDatabase() error {
|
|
collectionNum := int64(0)
|
|
partitionNum := int64(0)
|
|
oldCollections, err := mt.catalog.ListCollections(mt.ctx, util.NonDBID, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, collection := range oldCollections {
|
|
ensureCollectionMaxFieldIDProperty(collection)
|
|
mt.collID2Meta[collection.CollectionID] = collection
|
|
if collection.Available() {
|
|
mt.names.insert(util.DefaultDBName, collection.Name, collection.CollectionID)
|
|
for _, fileResourceID := range collection.FileResourceIds {
|
|
mt.fileResourceRefCnt[fileResourceID]++
|
|
}
|
|
pn := collection.GetPartitionNum(true)
|
|
mt.generalCnt += pn * int(collection.ShardsNum)
|
|
collectionNum++
|
|
partitionNum += int64(pn)
|
|
}
|
|
}
|
|
|
|
if collectionNum > 0 {
|
|
mlog.Info(mt.ctx, "recover collections without db", mlog.Int64("collection_num", collectionNum), mlog.Int64("partition_num", partitionNum))
|
|
}
|
|
|
|
aliases, err := mt.catalog.ListAliases(mt.ctx, util.NonDBID, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, alias := range aliases {
|
|
mt.aliases.insert(util.DefaultDBName, alias.Name, alias.CollectionID)
|
|
}
|
|
|
|
metrics.RootCoordNumOfCollections.WithLabelValues(util.DefaultDBName).Add(float64(collectionNum))
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Add(float64(partitionNum))
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) createDefaultDb() error {
|
|
// Generate ezID and db ts for default database
|
|
// Use unique ID as ezID because the default dbID(1) for each cluster is the same
|
|
ts, err := mt.tsoAllocator.GenerateTSO(2)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
s := Params.RootCoordCfg.DefaultDBProperties.GetValue()
|
|
defaultProperties, err := funcutil.String2KeyValuePair(s)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Apply same encryption logic as regular database creation
|
|
// This respects the defaultKey setting
|
|
defaultProperties, err = hookutil.TidyDBCipherProperties(int64(ts-1), defaultProperties)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Create EZ if encryption is enabled
|
|
if err := hookutil.CreateEZByDBProperties(defaultProperties); err != nil {
|
|
return err
|
|
}
|
|
|
|
return mt.createDatabasePrivate(mt.ctx, model.NewDefaultDatabase(defaultProperties), ts)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfDatabaseCreatable(ctx context.Context, req *milvuspb.CreateDatabaseRequest) error {
|
|
dbName := req.GetDbName()
|
|
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
if _, ok := mt.dbName2Meta[dbName]; ok || mt.aliases.exist(dbName) || mt.names.exist(dbName) {
|
|
// TODO: idempotency check here.
|
|
return merr.WrapErrParameterInvalidMsg("database already exist: %s", dbName)
|
|
}
|
|
|
|
cfgMaxDatabaseNum := Params.RootCoordCfg.MaxDatabaseNum.GetAsInt()
|
|
if len(mt.dbName2Meta) > cfgMaxDatabaseNum { // not include default database so use > instead of >= here.
|
|
return merr.WrapErrDatabaseNumLimitExceeded(cfgMaxDatabaseNum)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CreateDatabase(ctx context.Context, db *model.Database, ts typeutil.Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
if err := mt.createDatabasePrivate(ctx, db, ts); err != nil {
|
|
return err
|
|
}
|
|
metrics.RootCoordNumOfDatabases.Inc()
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) createDatabasePrivate(ctx context.Context, db *model.Database, ts typeutil.Timestamp) error {
|
|
dbName := db.Name
|
|
if err := mt.catalog.CreateDatabase(ctx, db, ts); err != nil {
|
|
return err
|
|
}
|
|
|
|
mt.names.createDbIfNotExist(dbName)
|
|
mt.aliases.createDbIfNotExist(dbName)
|
|
mt.dbName2Meta[dbName] = db
|
|
|
|
mlog.Info(ctx, "create database", mlog.String("db", dbName), mlog.Uint64("ts", ts))
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) AlterDatabase(ctx context.Context, newDB *model.Database, ts typeutil.Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.AlterDatabase(ctx1, newDB, ts); err != nil {
|
|
return err
|
|
}
|
|
mt.dbName2Meta[newDB.Name] = newDB
|
|
mlog.Info(ctx, "alter database finished", mlog.String("dbName", newDB.Name), mlog.Uint64("ts", ts))
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfDatabaseDroppable(ctx context.Context, req *milvuspb.DropDatabaseRequest) error {
|
|
dbName := req.GetDbName()
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
if dbName == util.DefaultDBName {
|
|
return merr.WrapErrParameterInvalidMsg("can not drop default database")
|
|
}
|
|
|
|
if _, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp); err != nil {
|
|
mlog.Warn(ctx, "not found database", mlog.String("db", dbName))
|
|
return err
|
|
}
|
|
|
|
colls, err := mt.listCollectionFromCache(ctx, dbName, true)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(colls) > 0 {
|
|
return merr.WrapErrParameterInvalidMsg("database:%s not empty, must drop all collections before drop database", dbName)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) DropDatabase(ctx context.Context, dbName string, ts typeutil.Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
db, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
mlog.Warn(ctx, "not found database", mlog.String("db", dbName))
|
|
return nil
|
|
}
|
|
if err := mt.catalog.DropDatabase(ctx, db.ID, ts); err != nil {
|
|
return err
|
|
}
|
|
|
|
mt.names.dropDb(dbName)
|
|
mt.aliases.dropDb(dbName)
|
|
delete(mt.dbName2Meta, dbName)
|
|
|
|
metrics.RootCoordNumOfDatabases.Dec()
|
|
mlog.Info(ctx, "drop database", mlog.String("db", dbName), mlog.Uint64("ts", ts))
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) ListDatabases(ctx context.Context, ts typeutil.Timestamp) ([]*model.Database, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
return maps.Values(mt.dbName2Meta), nil
|
|
}
|
|
|
|
func (mt *MetaTable) GetDatabaseByID(ctx context.Context, dbID int64, ts Timestamp) (*model.Database, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
return mt.getDatabaseByIDInternal(ctx, dbID, ts)
|
|
}
|
|
|
|
func (mt *MetaTable) getDatabaseByIDInternal(ctx context.Context, dbID int64, ts Timestamp) (*model.Database, error) {
|
|
for _, db := range maps.Values(mt.dbName2Meta) {
|
|
if db.ID == dbID {
|
|
return db, nil
|
|
}
|
|
}
|
|
return nil, merr.WrapErrDatabaseNotFound(dbID)
|
|
}
|
|
|
|
func (mt *MetaTable) GetDatabaseByName(ctx context.Context, dbName string, ts Timestamp) (*model.Database, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
return mt.getDatabaseByNameInternal(ctx, dbName, ts)
|
|
}
|
|
|
|
func (mt *MetaTable) getDatabaseByNameInternal(ctx context.Context, dbName string, _ Timestamp) (*model.Database, error) {
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty")
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
db, ok := mt.dbName2Meta[dbName]
|
|
if !ok {
|
|
return nil, merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
|
|
return db, nil
|
|
}
|
|
|
|
func (mt *MetaTable) AddCollection(ctx context.Context, coll *model.Collection) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
// Note:
|
|
// 1, idempotency check was already done outside;
|
|
// 2, no need to check time travel logic, since ts should always be the latest;
|
|
if coll.State != pb.CollectionState_CollectionCreated {
|
|
return merr.WrapErrServiceInternalMsg("collection state should be created, collection name: %s, collection id: %d, state: %s", coll.Name, coll.CollectionID, coll.State)
|
|
}
|
|
|
|
// check if there's a collection meta with the same collection id.
|
|
// merge the collection meta together.
|
|
if _, ok := mt.collID2Meta[coll.CollectionID]; ok {
|
|
mlog.Info(ctx, "collection already created, skip add collection to meta table", mlog.Int64("collectionID", coll.CollectionID))
|
|
return nil
|
|
}
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.CreateCollection(ctx1, coll, coll.CreateTime); err != nil {
|
|
return err
|
|
}
|
|
|
|
mt.collID2Meta[coll.CollectionID] = coll.Clone()
|
|
mt.names.insert(coll.DBName, coll.Name, coll.CollectionID)
|
|
// Build partition name index for the new collection
|
|
mt.partitionName2ID[coll.CollectionID] = make(map[string]int64)
|
|
for _, partition := range coll.Partitions {
|
|
if partition.Available() {
|
|
mt.partitionName2ID[coll.CollectionID][partition.PartitionName] = partition.PartitionID
|
|
}
|
|
}
|
|
|
|
pn := coll.GetPartitionNum(true)
|
|
mt.generalCnt += pn * int(coll.ShardsNum)
|
|
metrics.RootCoordNumOfCollections.WithLabelValues(coll.DBName).Inc()
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Add(float64(pn))
|
|
|
|
channel.StaticPChannelStatsManager.MustGet().AddVChannel(coll.VirtualChannelNames...)
|
|
mlog.Info(ctx, "add collection to meta table",
|
|
mlog.Int64("dbID", coll.DBID),
|
|
mlog.String("collection", coll.Name),
|
|
mlog.Int64("id", coll.CollectionID),
|
|
mlog.Uint64("ts", coll.CreateTime),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) DropCollection(ctx context.Context, collectionID UniqueID, ts Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
if coll.State == pb.CollectionState_CollectionDropping {
|
|
return nil
|
|
}
|
|
|
|
clone := coll.Clone()
|
|
clone.State = pb.CollectionState_CollectionDropping
|
|
clone.UpdateTimestamp = ts
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.AlterCollection(ctx1, coll, clone, metastore.MODIFY, ts, false); err != nil {
|
|
return err
|
|
}
|
|
mt.collID2Meta[collectionID] = clone
|
|
for _, fileResourceID := range coll.FileResourceIds {
|
|
if mt.fileResourceRefCnt[fileResourceID] > 0 {
|
|
mt.fileResourceRefCnt[fileResourceID]--
|
|
} else {
|
|
mlog.Warn(context.TODO(), "DropCollection: file resource refCnt underflow",
|
|
mlog.Int64("collectionID", collectionID), mlog.Int64("fileResourceID", fileResourceID))
|
|
}
|
|
}
|
|
|
|
mlog.Info(ctx, "update coll state to dropping",
|
|
mlog.Int64("collectionID", collectionID),
|
|
mlog.String("state", clone.State.String()),
|
|
)
|
|
|
|
db, err := mt.getDatabaseByIDInternal(ctx, coll.DBID, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return merr.Wrapf(err, "dbID not found for collection:%d", collectionID)
|
|
}
|
|
|
|
pn := coll.GetPartitionNum(true)
|
|
|
|
mt.generalCnt -= pn * int(coll.ShardsNum)
|
|
channel.StaticPChannelStatsManager.MustGet().RemoveVChannel(coll.VirtualChannelNames...)
|
|
metrics.RootCoordNumOfCollections.WithLabelValues(db.Name).Dec()
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Sub(float64(pn))
|
|
|
|
mlog.Info(ctx, "drop collection from meta table", mlog.Int64("collection", collectionID),
|
|
mlog.String("state", coll.State.String()), mlog.Uint64("ts", ts))
|
|
|
|
// Delete all grants referencing this collection immediately so they don't
|
|
// linger until the tombstone sweeper runs (which can take minutes).
|
|
if err := mt.catalog.DeleteGrantByCollectionName(ctx1, util.DefaultTenant, db.Name, coll.Name); err != nil {
|
|
mlog.Warn(ctx, "failed to delete grants for dropped collection, skipping",
|
|
mlog.String("dbName", db.Name), mlog.String("collectionName", coll.Name), mlog.Err(err))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) removeIfNameMatchedInternal(ctx context.Context, collectionID UniqueID, name string) {
|
|
mt.names.removeIf(func(db string, collection string, id UniqueID) bool {
|
|
if collectionID == id {
|
|
mlog.Info(ctx, "remove from names",
|
|
mlog.String("dbName", db),
|
|
mlog.String("collectionName", collection),
|
|
mlog.Int64("collectionID", id),
|
|
)
|
|
return true
|
|
}
|
|
return false
|
|
})
|
|
}
|
|
|
|
func (mt *MetaTable) removeIfAliasMatchedInternal(ctx context.Context, collectionID UniqueID, alias string) {
|
|
mt.aliases.removeIf(func(db string, collection string, id UniqueID) bool {
|
|
if collectionID == id {
|
|
mlog.Info(ctx, "remove from aliases",
|
|
mlog.String("dbName", db),
|
|
mlog.String("alias", collection),
|
|
mlog.Int64("collectionID", id),
|
|
)
|
|
return true
|
|
}
|
|
return false
|
|
})
|
|
}
|
|
|
|
func (mt *MetaTable) removeIfMatchedInternal(ctx context.Context, collectionID UniqueID, name string) {
|
|
mt.removeIfNameMatchedInternal(ctx, collectionID, name)
|
|
mt.removeIfAliasMatchedInternal(ctx, collectionID, name)
|
|
}
|
|
|
|
func (mt *MetaTable) removeAllNamesIfMatchedInternal(ctx context.Context, collectionID UniqueID, names []string) {
|
|
for _, name := range names {
|
|
mt.removeIfMatchedInternal(ctx, collectionID, name)
|
|
}
|
|
}
|
|
|
|
func (mt *MetaTable) removeCollectionByIDInternal(ctx context.Context, collectionID UniqueID) {
|
|
delete(mt.collID2Meta, collectionID)
|
|
delete(mt.partitionName2ID, collectionID)
|
|
mlog.Info(ctx, "delete from collID2Meta",
|
|
mlog.Int64("collectionID", collectionID),
|
|
)
|
|
}
|
|
|
|
func (mt *MetaTable) RemoveCollection(ctx context.Context, collectionID UniqueID, ts Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
// Note: we cannot handle case that dropping collection with `ts1` but a collection exists in catalog with newer ts
|
|
// which is bigger than `ts1`. So we assume that ts should always be the latest.
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
mlog.Warn(ctx, "not found collection, skip remove", mlog.Int64("collectionID", collectionID))
|
|
return nil
|
|
}
|
|
if coll.State != pb.CollectionState_CollectionDropping {
|
|
return merr.WrapErrServiceInternalMsg("remove collection which state is not dropping, collectionID: %d, state: %s", collectionID, coll.State.String())
|
|
}
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
aliases := mt.listAliasesByID(collectionID)
|
|
newColl := &model.Collection{
|
|
CollectionID: collectionID,
|
|
Partitions: model.ClonePartitions(coll.Partitions),
|
|
Fields: model.CloneFields(coll.Fields),
|
|
StructArrayFields: model.CloneStructArrayFields(coll.StructArrayFields),
|
|
Aliases: aliases,
|
|
DBID: coll.DBID,
|
|
}
|
|
if err := mt.catalog.DropCollection(ctx1, newColl, ts); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := mt.catalog.DeleteGrantByCollectionName(ctx1, util.DefaultTenant, coll.DBName, coll.Name); err != nil {
|
|
mlog.Warn(ctx, "failed to delete grants for dropped collection, skipping",
|
|
mlog.String("dbName", coll.DBName), mlog.String("collectionName", coll.Name), mlog.Err(err))
|
|
}
|
|
|
|
allNames := common.CloneStringList(aliases)
|
|
allNames = append(allNames, coll.Name)
|
|
|
|
// We cannot delete the name directly, since newly collection with same name may be created.
|
|
mt.removeAllNamesIfMatchedInternal(ctx, collectionID, allNames)
|
|
mt.removeCollectionByIDInternal(ctx, collectionID)
|
|
|
|
mlog.Info(ctx, "remove collection",
|
|
mlog.Int64("dbID", coll.DBID),
|
|
mlog.String("name", coll.Name),
|
|
mlog.Int64("id", collectionID),
|
|
mlog.Strings("aliases", aliases),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
// Note: The returned model.Collection is read-only. Do NOT modify it directly,
|
|
// as it may cause unexpected behavior or inconsistencies.
|
|
func filterUnavailablePartition(coll *model.Collection) *model.Collection {
|
|
clone := coll.ShallowClone()
|
|
// pick available partitions.
|
|
clone.Partitions = make([]*model.Partition, 0, len(coll.Partitions))
|
|
for _, partition := range coll.Partitions {
|
|
if partition.Available() {
|
|
clone.Partitions = append(clone.Partitions, partition)
|
|
}
|
|
}
|
|
return clone
|
|
}
|
|
|
|
// getLatestCollectionByIDInternal should be called with ts = typeutil.MaxTimestamp
|
|
// Note: The returned model.Collection is read-only. Do NOT modify it directly,
|
|
// as it may cause unexpected behavior or inconsistencies.
|
|
func (mt *MetaTable) getLatestCollectionByIDInternal(ctx context.Context, collectionID UniqueID, allowUnavailable bool) (*model.Collection, error) {
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok || coll == nil {
|
|
mlog.Warn(context.TODO(), "not found collection", mlog.Int64("collectionID", collectionID))
|
|
return nil, merr.WrapErrCollectionNotFound(collectionID)
|
|
}
|
|
if allowUnavailable {
|
|
return coll.ShallowClone(), nil
|
|
}
|
|
if !coll.Available() {
|
|
return nil, merr.WrapErrCollectionNotFound(collectionID)
|
|
}
|
|
return filterUnavailablePartition(coll), nil
|
|
}
|
|
|
|
// getCollectionByIDInternal get collection by collection id without lock.
|
|
// Note: The returned model.Collection is read-only. Do NOT modify it directly,
|
|
// as it may cause unexpected behavior or inconsistencies.
|
|
func (mt *MetaTable) getCollectionByIDInternal(ctx context.Context, dbName string, collectionID UniqueID, ts Timestamp, allowUnavailable bool) (*model.Collection, error) {
|
|
if isMaxTs(ts) {
|
|
return mt.getLatestCollectionByIDInternal(ctx, collectionID, allowUnavailable)
|
|
}
|
|
|
|
var coll *model.Collection
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
// Use UpdateTimestamp (not CreateTime) for cache invalidation.
|
|
// This is a bug fix for time-travel correctness, not merely an enhancement for backfill:
|
|
// collID2Meta always holds the *latest* collection state. After a schema alteration (e.g.
|
|
// AlterCollectionSchema), the in-memory entry has a schema that didn't exist at timestamps
|
|
// before the alteration. The old code compared against CreateTime, so a time-travel query
|
|
// at T where CreateTime < T < AlterTime would incorrectly serve the altered schema from the
|
|
// in-memory cache instead of falling back to the catalog for the historical version.
|
|
// Comparing against UpdateTimestamp ensures the cache is bypassed for any ts that predates
|
|
// the last schema modification, fixing time-travel semantics for all schema-altering DDLs.
|
|
if !ok || coll == nil || !coll.Available() || coll.UpdateTimestamp > ts {
|
|
// travel meta information from catalog.
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
db, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
coll, err = mt.catalog.GetCollectionByID(ctx1, db.ID, ts, collectionID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
if coll == nil {
|
|
return nil, merr.WrapErrCollectionNotFound(collectionID)
|
|
}
|
|
|
|
if allowUnavailable {
|
|
return coll.ShallowClone(), nil
|
|
}
|
|
|
|
if !coll.Available() {
|
|
return nil, merr.WrapErrCollectionNotFound(collectionID)
|
|
}
|
|
|
|
return filterUnavailablePartition(coll), nil
|
|
}
|
|
|
|
func (mt *MetaTable) GetCollectionByName(ctx context.Context, dbName string, collectionName string, ts Timestamp, allowUnavailable bool) (*model.Collection, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
return mt.getCollectionByNameInternal(ctx, dbName, collectionName, ts, allowUnavailable)
|
|
}
|
|
|
|
// GetCollectionID retrieves the corresponding collectionID based on the collectionName.
|
|
// If the collection does not exist, it will return InvalidCollectionID.
|
|
// Please use the function with caution.
|
|
func (mt *MetaTable) GetCollectionID(ctx context.Context, dbName string, collectionName string) UniqueID {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(context.TODO(), "db name is empty", mlog.String("collectionName", collectionName))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
_, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return InvalidCollectionID
|
|
}
|
|
|
|
collectionID, ok := mt.aliases.get(dbName, collectionName)
|
|
if ok {
|
|
return collectionID
|
|
}
|
|
|
|
collectionID, ok = mt.names.get(dbName, collectionName)
|
|
if ok {
|
|
return collectionID
|
|
}
|
|
return InvalidCollectionID
|
|
}
|
|
|
|
// Note: The returned model.Collection is read-only. Do NOT modify it directly,
|
|
// as it may cause unexpected behavior or inconsistencies.
|
|
func (mt *MetaTable) getCollectionByNameInternal(ctx context.Context, dbName string, collectionName string, ts Timestamp, allowUnavailable bool) (*model.Collection, error) {
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty", mlog.String("collectionName", collectionName), mlog.Uint64("ts", ts))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
db, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
collectionID, ok := mt.aliases.get(dbName, collectionName)
|
|
if ok {
|
|
return mt.getCollectionByIDInternal(ctx, dbName, collectionID, ts, allowUnavailable)
|
|
}
|
|
|
|
collectionID, ok = mt.names.get(dbName, collectionName)
|
|
if ok {
|
|
return mt.getCollectionByIDInternal(ctx, dbName, collectionID, ts, allowUnavailable)
|
|
}
|
|
|
|
if isMaxTs(ts) {
|
|
return nil, merr.WrapErrCollectionNotFoundWithDB(dbName, collectionName)
|
|
}
|
|
|
|
// travel meta information from catalog. No need to check time travel logic again, since catalog already did.
|
|
ctx = contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
coll, err := mt.catalog.GetCollectionByName(ctx, db.ID, db.Name, collectionName, ts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if coll == nil || !coll.Available() {
|
|
return nil, merr.WrapErrCollectionNotFoundWithDB(dbName, collectionName)
|
|
}
|
|
if allowUnavailable {
|
|
return coll.ShallowClone(), nil
|
|
}
|
|
return filterUnavailablePartition(coll), nil
|
|
}
|
|
|
|
func (mt *MetaTable) GetCollectionByID(ctx context.Context, dbName string, collectionID UniqueID, ts Timestamp, allowUnavailable bool) (*model.Collection, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
return mt.getCollectionByIDInternal(ctx, dbName, collectionID, ts, allowUnavailable)
|
|
}
|
|
|
|
// GetCollectionByIDWithMaxTs get collection, dbName can be ignored if ts is max timestamps
|
|
func (mt *MetaTable) GetCollectionByIDWithMaxTs(ctx context.Context, collectionID UniqueID) (*model.Collection, error) {
|
|
return mt.GetCollectionByID(ctx, "", collectionID, typeutil.MaxTimestamp, false)
|
|
}
|
|
|
|
func (mt *MetaTable) ListAllAvailCollections(ctx context.Context) map[int64][]int64 {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
ret := make(map[int64][]int64, len(mt.dbName2Meta))
|
|
for _, dbMeta := range mt.dbName2Meta {
|
|
ret[dbMeta.ID] = make([]int64, 0)
|
|
}
|
|
|
|
for collID, collMeta := range mt.collID2Meta {
|
|
if !collMeta.Available() {
|
|
continue
|
|
}
|
|
dbID := collMeta.DBID
|
|
if dbID == util.NonDBID {
|
|
ret[util.DefaultDBID] = append(ret[util.DefaultDBID], collID)
|
|
continue
|
|
}
|
|
ret[dbID] = append(ret[dbID], collID)
|
|
}
|
|
|
|
return ret
|
|
}
|
|
|
|
func (mt *MetaTable) ListAllAvailPartitions(ctx context.Context) map[int64]map[int64][]int64 {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
ret := make(map[int64]map[int64][]int64, len(mt.dbName2Meta))
|
|
for _, dbMeta := range mt.dbName2Meta {
|
|
// Database may not have available collections.
|
|
ret[dbMeta.ID] = make(map[int64][]int64, 64)
|
|
}
|
|
for _, collMeta := range mt.collID2Meta {
|
|
if !collMeta.Available() {
|
|
continue
|
|
}
|
|
dbID := collMeta.DBID
|
|
if dbID == util.NonDBID {
|
|
dbID = util.DefaultDBID
|
|
}
|
|
if _, ok := ret[dbID]; !ok {
|
|
ret[dbID] = make(map[int64][]int64, 64)
|
|
}
|
|
ret[dbID][collMeta.CollectionID] = lo.Map(collMeta.Partitions, func(part *model.Partition, _ int) int64 { return part.PartitionID })
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func (mt *MetaTable) ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
if isMaxTs(ts) {
|
|
return mt.listCollectionFromCache(ctx, dbName, onlyAvail)
|
|
}
|
|
|
|
db, err := mt.getDatabaseByNameInternal(ctx, dbName, typeutil.MaxTimestamp)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// list collections should always be loaded from catalog.
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
colls, err := mt.catalog.ListCollections(ctx1, db.ID, ts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
onlineCollections := make([]*model.Collection, 0, len(colls))
|
|
for _, coll := range colls {
|
|
if onlyAvail && !coll.Available() {
|
|
continue
|
|
}
|
|
onlineCollections = append(onlineCollections, coll)
|
|
}
|
|
return onlineCollections, nil
|
|
}
|
|
|
|
func (mt *MetaTable) listCollectionFromCache(ctx context.Context, dbName string, onlyAvail bool) ([]*model.Collection, error) {
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty")
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
db, ok := mt.dbName2Meta[dbName]
|
|
if !ok {
|
|
return nil, merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
|
|
collectionFromCache := make([]*model.Collection, 0, len(mt.collID2Meta))
|
|
for _, collMeta := range mt.collID2Meta {
|
|
if (collMeta.DBID != util.NonDBID && db.ID == collMeta.DBID) ||
|
|
(collMeta.DBID == util.NonDBID && dbName == util.DefaultDBName) {
|
|
if onlyAvail && !collMeta.Available() {
|
|
continue
|
|
}
|
|
|
|
collectionFromCache = append(collectionFromCache, collMeta)
|
|
}
|
|
}
|
|
return collectionFromCache, nil
|
|
}
|
|
|
|
// ListCollectionPhysicalChannels list physical channels of all collections.
|
|
func (mt *MetaTable) ListCollectionPhysicalChannels(ctx context.Context) map[typeutil.UniqueID][]string {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
chanMap := make(map[UniqueID][]string)
|
|
|
|
for id, collInfo := range mt.collID2Meta {
|
|
chanMap[id] = common.CloneStringList(collInfo.PhysicalChannelNames)
|
|
}
|
|
|
|
return chanMap
|
|
}
|
|
|
|
// AlterCollection is used to alter a collection in the meta table.
|
|
func (mt *MetaTable) AlterCollection(ctx context.Context, result message.BroadcastResultAlterCollectionMessageV2) error {
|
|
header := result.Message.Header()
|
|
body := result.Message.MustBody()
|
|
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[header.CollectionId]
|
|
if !ok {
|
|
// collection not exists, return directly.
|
|
return errAlterCollectionNotFound
|
|
}
|
|
oldColl := coll.Clone()
|
|
newColl := coll.Clone()
|
|
newColl.ApplyUpdates(header, body)
|
|
fieldModify := false
|
|
dbChanged := false
|
|
for _, path := range header.UpdateMask.GetPaths() {
|
|
switch path {
|
|
case message.FieldMaskCollectionSchema:
|
|
fieldModify = true
|
|
case message.FieldMaskDB:
|
|
dbChanged = true
|
|
}
|
|
}
|
|
newColl.UpdateTimestamp = result.GetMaxTimeTick()
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if !dbChanged {
|
|
if err := mt.catalog.AlterCollection(ctx1, oldColl, newColl, metastore.MODIFY, newColl.UpdateTimestamp, fieldModify); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
if err := mt.catalog.AlterCollectionDB(ctx1, oldColl, newColl, newColl.UpdateTimestamp); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if oldColl.Name != newColl.Name || oldColl.DBName != newColl.DBName {
|
|
if err := mt.catalog.MigrateGrantCollectionName(ctx1, util.DefaultTenant, oldColl.DBName, oldColl.Name, newColl.DBName, newColl.Name); err != nil {
|
|
mlog.Warn(ctx, "failed to migrate grants for renamed collection, skipping",
|
|
mlog.String("oldDBName", oldColl.DBName), mlog.String("oldName", oldColl.Name),
|
|
mlog.String("newDBName", newColl.DBName), mlog.String("newName", newColl.Name), mlog.Err(err))
|
|
}
|
|
}
|
|
|
|
mt.names.remove(oldColl.DBName, oldColl.Name)
|
|
mt.names.insert(newColl.DBName, newColl.Name, newColl.CollectionID)
|
|
mt.collID2Meta[header.CollectionId] = newColl
|
|
mlog.Info(ctx, "alter collection finished",
|
|
mlog.String("oldDBName", oldColl.DBName),
|
|
mlog.String("newDBName", newColl.DBName),
|
|
mlog.String("oldCollectionName", oldColl.Name),
|
|
mlog.String("newCollectionName", newColl.Name),
|
|
mlog.Int64("headerCollectionID", header.CollectionId),
|
|
mlog.Int64("newCollectionID", newColl.CollectionID),
|
|
mlog.Int64("oldCollectionID", oldColl.CollectionID),
|
|
mlog.Bool("dbChanged", dbChanged),
|
|
mlog.Uint64("ts", newColl.UpdateTimestamp),
|
|
mlog.Int32("schemaVersion", newColl.SchemaVersion),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) BeginTruncateCollection(ctx context.Context, collectionID UniqueID) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return errAlterCollectionNotFound
|
|
}
|
|
|
|
// Apply the properties to override the existing properties.
|
|
newProperties := common.CloneKeyValuePairs(coll.Properties).ToMap()
|
|
key := common.CollectionOnTruncatingKey
|
|
if _, ok := newProperties[key]; ok && newProperties[key] == "1" {
|
|
return nil
|
|
}
|
|
newProperties[key] = "1"
|
|
oldColl := coll.Clone()
|
|
newColl := coll.Clone()
|
|
newColl.Properties = common.NewKeyValuePairs(newProperties)
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.AlterCollection(ctx1, oldColl, newColl, metastore.MODIFY, newColl.UpdateTimestamp, false); err != nil {
|
|
return err
|
|
}
|
|
mt.collID2Meta[coll.CollectionID] = newColl
|
|
mlog.Info(ctx, "update collID2Meta for begin truncate collection",
|
|
mlog.Int64("collectionID", coll.CollectionID),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) TruncateCollection(ctx context.Context, result message.BroadcastResultTruncateCollectionMessageV2) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
collectionID := result.Message.Header().CollectionId
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return errAlterCollectionNotFound
|
|
}
|
|
|
|
oldColl := coll.Clone()
|
|
|
|
// remmove the truncating key from the properties and update the last truncate time tick of the shard infos
|
|
newColl := coll.Clone()
|
|
newProperties := common.CloneKeyValuePairs(coll.Properties).ToMap()
|
|
delete(newProperties, common.CollectionOnTruncatingKey)
|
|
newColl.Properties = common.NewKeyValuePairs(newProperties)
|
|
for vchannel := range newColl.ShardInfos {
|
|
newColl.ShardInfos[vchannel].LastTruncateTimeTick = result.Results[vchannel].TimeTick
|
|
}
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.AlterCollection(ctx1, oldColl, newColl, metastore.MODIFY, newColl.UpdateTimestamp, false); err != nil {
|
|
return err
|
|
}
|
|
mt.collID2Meta[coll.CollectionID] = newColl
|
|
mlog.Info(ctx, "update collID2Meta for truncate collection",
|
|
mlog.Int64("collectionID", coll.CollectionID),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfCollectionRenamable(ctx context.Context, dbName string, oldName string, newDBName string, newName string) error {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
ctx = contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
|
|
// DB name already filled in rename collection task prepare
|
|
// get target db
|
|
targetDB, ok := mt.dbName2Meta[newDBName]
|
|
if !ok {
|
|
return merr.WrapErrDatabaseNotFound(newDBName)
|
|
}
|
|
|
|
// old collection should not be an alias
|
|
_, ok = mt.aliases.get(dbName, oldName)
|
|
if ok {
|
|
mlog.Warn(context.TODO(), "unsupported use a alias to rename collection")
|
|
return merr.WrapErrParameterInvalidMsg("unsupported use an alias to rename collection, alias:%s", oldName)
|
|
}
|
|
|
|
_, ok = mt.aliases.get(newDBName, newName)
|
|
if ok {
|
|
mlog.Warn(context.TODO(), "cannot rename collection to an existing alias")
|
|
return merr.WrapErrAsInputError(merr.WrapErrAliasCollectionNameConflict(newDBName, newName))
|
|
}
|
|
|
|
// check new collection already exists
|
|
coll, err := mt.getCollectionByNameInternal(ctx, newDBName, newName, typeutil.MaxTimestamp, false)
|
|
if coll != nil {
|
|
mlog.Warn(context.TODO(), "duplicated new collection name, already taken by another collection or alias.")
|
|
return merr.WrapErrParameterInvalidMsg("duplicated new collection name %s:%s with other collection name or alias", newDBName, newName)
|
|
}
|
|
if err != nil && !errors.Is(err, merr.ErrCollectionNotFound) {
|
|
mlog.Warn(context.TODO(), "fail to check if new collection name is already taken", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
// get old collection meta
|
|
oldColl, err := mt.getCollectionByNameInternal(ctx, dbName, oldName, typeutil.MaxTimestamp, false)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "fail to find collection with old name", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
// unsupported rename collection while the collection has aliases
|
|
aliases := mt.listAliasesByID(oldColl.CollectionID)
|
|
if len(aliases) > 0 && oldColl.DBID != targetDB.ID {
|
|
return merr.WrapErrParameterInvalidMsg("fail to rename db name, must drop all aliases of this collection before rename")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetCollectionVirtualChannels returns virtual channels of a given collection.
|
|
func (mt *MetaTable) GetCollectionVirtualChannels(ctx context.Context, colID int64) []string {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
for id, collInfo := range mt.collID2Meta {
|
|
if id == colID {
|
|
return common.CloneStringList(collInfo.VirtualChannelNames)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetPChannelInfo returns infos on pchannel.
|
|
func (mt *MetaTable) GetPChannelInfo(ctx context.Context, pchannel string) *rootcoordpb.GetPChannelInfoResponse {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
resp := &rootcoordpb.GetPChannelInfoResponse{
|
|
Status: merr.Success(),
|
|
Collections: make([]*rootcoordpb.CollectionInfoOnPChannel, 0),
|
|
}
|
|
for _, collInfo := range mt.collID2Meta {
|
|
if collInfo.State != pb.CollectionState_CollectionCreated && collInfo.State != pb.CollectionState_CollectionDropping {
|
|
// streamingnode will receive the createCollectionMessage to recover if the collection is creating.
|
|
// streamingnode use it to recover the collection state at first time streaming arch enabled.
|
|
// streamingnode will get the dropping collection and drop it before streaming arch enabled.
|
|
continue
|
|
}
|
|
if idx := lo.IndexOf(collInfo.PhysicalChannelNames, pchannel); idx >= 0 {
|
|
partitions := make([]*rootcoordpb.PartitionInfoOnPChannel, 0, len(collInfo.Partitions))
|
|
for _, part := range collInfo.Partitions {
|
|
partitions = append(partitions, &rootcoordpb.PartitionInfoOnPChannel{
|
|
PartitionId: part.PartitionID,
|
|
})
|
|
}
|
|
resp.Collections = append(resp.Collections, &rootcoordpb.CollectionInfoOnPChannel{
|
|
CollectionId: collInfo.CollectionID,
|
|
Partitions: partitions,
|
|
Vchannel: collInfo.VirtualChannelNames[idx],
|
|
State: collInfo.State,
|
|
})
|
|
}
|
|
}
|
|
return resp
|
|
}
|
|
|
|
func (mt *MetaTable) AddPartition(ctx context.Context, partition *model.Partition) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[partition.CollectionID]
|
|
if !ok || !coll.Available() {
|
|
return merr.WrapErrServiceInternalMsg("collection not exists: %d", partition.CollectionID)
|
|
}
|
|
|
|
if partition.State != pb.PartitionState_PartitionCreated {
|
|
return merr.WrapErrServiceInternalMsg("partition state is not created, collection: %d, partition: %d, state: %s", partition.CollectionID, partition.PartitionID, partition.State)
|
|
}
|
|
|
|
// idempotency check here.
|
|
for _, part := range coll.Partitions {
|
|
if part.PartitionID == partition.PartitionID {
|
|
mlog.Info(ctx, "partition already exists, ignore the operation", mlog.Int64("collection", partition.CollectionID), mlog.Int64("partition", partition.PartitionID))
|
|
return nil
|
|
}
|
|
}
|
|
if err := mt.catalog.CreatePartition(ctx, coll.DBID, partition, partition.PartitionCreatedTimestamp); err != nil {
|
|
return err
|
|
}
|
|
// COW: create new slice to avoid data race with concurrent readers
|
|
oldPartitions := coll.Partitions
|
|
newPartitions := make([]*model.Partition, len(oldPartitions)+1)
|
|
copy(newPartitions, oldPartitions)
|
|
newPartitions[len(oldPartitions)] = partition.Clone()
|
|
coll.Partitions = newPartitions
|
|
|
|
// Update partition name index
|
|
if mt.partitionName2ID[partition.CollectionID] == nil {
|
|
mt.partitionName2ID[partition.CollectionID] = make(map[string]int64)
|
|
}
|
|
mt.partitionName2ID[partition.CollectionID][partition.PartitionName] = partition.PartitionID
|
|
|
|
mlog.Info(ctx, "add partition to meta table",
|
|
mlog.Int64("collection", partition.CollectionID), mlog.String("partition", partition.PartitionName),
|
|
mlog.Int64("partitionid", partition.PartitionID), mlog.Uint64("ts", partition.PartitionCreatedTimestamp))
|
|
mt.generalCnt += int(coll.ShardsNum) // 1 partition * shardNum
|
|
// support Dynamic load/release partitions
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Inc()
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetPartitionIDByName returns partition ID by collection ID and partition name.
|
|
// Returns (partitionID, true) if found, (0, false) if not found.
|
|
func (mt *MetaTable) GetPartitionIDByName(collectionID int64, partitionName string) (int64, bool) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
if partitions, ok := mt.partitionName2ID[collectionID]; ok {
|
|
if partitionID, exists := partitions[partitionName]; exists {
|
|
return partitionID, true
|
|
}
|
|
}
|
|
return 0, false
|
|
}
|
|
|
|
func (mt *MetaTable) DropPartition(ctx context.Context, collectionID UniqueID, partitionID UniqueID, ts Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
for idx, part := range coll.Partitions {
|
|
if part.PartitionID == partitionID {
|
|
if part.State == pb.PartitionState_PartitionDropping {
|
|
// promise idempotency here.
|
|
return nil
|
|
}
|
|
clone := part.Clone()
|
|
clone.State = pb.PartitionState_PartitionDropping
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.AlterPartition(ctx1, coll.DBID, part, clone, metastore.MODIFY, ts); err != nil {
|
|
return err
|
|
}
|
|
// Copy-on-write so snapshots returned by ShallowClone never observe
|
|
// in-place updates through a shared partitions backing array.
|
|
newPartitions := make([]*model.Partition, len(coll.Partitions))
|
|
copy(newPartitions, coll.Partitions)
|
|
newPartitions[idx] = clone
|
|
mt.collID2Meta[collectionID].Partitions = newPartitions
|
|
|
|
// Remove from partition name index when dropping
|
|
if mt.partitionName2ID[collectionID] != nil {
|
|
delete(mt.partitionName2ID[collectionID], part.PartitionName)
|
|
}
|
|
|
|
mlog.Info(ctx, "drop partition", mlog.Int64("collection", collectionID),
|
|
mlog.Int64("partition", partitionID),
|
|
mlog.Uint64("ts", ts))
|
|
|
|
mt.generalCnt -= int(coll.ShardsNum) // 1 partition * shardNum
|
|
metrics.RootCoordNumOfPartitions.WithLabelValues().Dec()
|
|
return nil
|
|
}
|
|
}
|
|
// partition not found, so promise idempotency here.
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) RemovePartition(ctx context.Context, collectionID UniqueID, partitionID UniqueID, ts Timestamp) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
loc := -1
|
|
for idx, part := range coll.Partitions {
|
|
if part.PartitionID == partitionID {
|
|
loc = idx
|
|
break
|
|
}
|
|
}
|
|
if loc == -1 {
|
|
mlog.Warn(ctx, "not found partition, skip remove", mlog.Int64("collection", collectionID), mlog.Int64("partition", partitionID))
|
|
return nil
|
|
}
|
|
partition := coll.Partitions[loc]
|
|
if partition.State != pb.PartitionState_PartitionDropping {
|
|
return merr.WrapErrServiceInternalMsg("remove partition which state is not dropping, collection: %d, partition: %d, state: %s", collectionID, partitionID, partition.State.String())
|
|
}
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.DropPartition(ctx1, coll.DBID, collectionID, partitionID, ts); err != nil {
|
|
return err
|
|
}
|
|
// COW: create new slice to avoid data race with concurrent readers
|
|
oldPartitions := coll.Partitions
|
|
newPartitions := make([]*model.Partition, len(oldPartitions)-1)
|
|
copy(newPartitions[:loc], oldPartitions[:loc])
|
|
copy(newPartitions[loc:], oldPartitions[loc+1:])
|
|
coll.Partitions = newPartitions
|
|
|
|
// Remove from partition name index
|
|
if mt.partitionName2ID[collectionID] != nil {
|
|
delete(mt.partitionName2ID[collectionID], partition.PartitionName)
|
|
}
|
|
|
|
mlog.Info(ctx, "remove partition", mlog.Int64("collection", collectionID), mlog.Int64("partition", partitionID), mlog.Uint64("ts", ts))
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfAliasCreatable(ctx context.Context, dbName string, alias string, collectionName string) error {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty", mlog.String("alias", alias), mlog.String("collection", collectionName))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
// It's ok that we don't read from catalog when cache missed.
|
|
// Since cache always keep the latest version, and the ts should always be the latest.
|
|
|
|
if !mt.names.exist(dbName) {
|
|
return merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
|
|
if collID, ok := mt.names.get(dbName, alias); ok {
|
|
coll, ok := mt.collID2Meta[collID]
|
|
if !ok {
|
|
return merr.WrapErrServiceInternalMsg("meta error, name mapped non-exist collection id")
|
|
}
|
|
// allow alias with dropping&dropped
|
|
if coll.State != pb.CollectionState_CollectionDropping && coll.State != pb.CollectionState_CollectionDropped {
|
|
return merr.WrapErrAliasCollectionNameConflict(dbName, alias)
|
|
}
|
|
}
|
|
|
|
collectionID, ok := mt.names.get(dbName, collectionName)
|
|
if !ok {
|
|
// you cannot alias to a non-existent collection.
|
|
return merr.WrapErrCollectionNotFoundWithDB(dbName, collectionName)
|
|
}
|
|
|
|
// check if alias exists.
|
|
aliasedCollectionID, ok := mt.aliases.get(dbName, alias)
|
|
if ok && aliasedCollectionID == collectionID {
|
|
mlog.Warn(ctx, "add duplicate alias", mlog.String("alias", alias), mlog.String("collection", collectionName))
|
|
return errIgnoredAlterAlias
|
|
} else if ok {
|
|
// TODO: better to check if aliasedCollectionID exist or is available, though not very possible.
|
|
aliasedColl := mt.collID2Meta[aliasedCollectionID]
|
|
msg := fmt.Sprintf("%s is alias to another collection: %s", alias, aliasedColl.Name)
|
|
return merr.WrapErrAliasAlreadyExist(dbName, alias, msg)
|
|
}
|
|
// alias didn't exist.
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok || !coll.Available() {
|
|
// you cannot alias to a non-existent collection.
|
|
return merr.WrapErrCollectionNotFoundWithDB(dbName, collectionName)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfAliasDroppable(ctx context.Context, dbName string, alias string) error {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
if _, ok := mt.aliases.get(dbName, alias); !ok {
|
|
return merr.WrapErrAliasNotFound(dbName, alias)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) DropAlias(ctx context.Context, result message.BroadcastResultDropAliasMessageV2) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
header := result.Message.Header()
|
|
|
|
ctx1 := contextutil.WithTenantID(ctx, Params.CommonCfg.ClusterName.GetValue())
|
|
if err := mt.catalog.DropAlias(ctx1, header.DbId, header.Alias, result.GetControlChannelResult().TimeTick); err != nil {
|
|
return err
|
|
}
|
|
mt.aliases.remove(header.DbName, header.Alias)
|
|
|
|
mlog.Info(ctx, "drop alias",
|
|
mlog.String("db", header.DbName),
|
|
mlog.String("alias", header.Alias),
|
|
mlog.Uint64("ts", result.GetControlChannelResult().TimeTick),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) AlterAlias(ctx context.Context, result message.BroadcastResultAlterAliasMessageV2) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
header := result.Message.Header()
|
|
if err := mt.catalog.AlterAlias(ctx, &model.Alias{
|
|
Name: header.Alias,
|
|
CollectionID: header.CollectionId,
|
|
CreatedTime: result.GetControlChannelResult().TimeTick,
|
|
State: pb.AliasState_AliasCreated,
|
|
DbID: header.DbId,
|
|
}, result.GetControlChannelResult().TimeTick); err != nil {
|
|
return err
|
|
}
|
|
|
|
// alias switch to another collection anyway.
|
|
mt.aliases.insert(header.DbName, header.Alias, header.CollectionId)
|
|
|
|
mlog.Info(ctx, "alter alias",
|
|
mlog.String("db", header.DbName),
|
|
mlog.String("alias", header.Alias),
|
|
mlog.String("collectionName", header.CollectionName),
|
|
mlog.Int64("collectionID", header.CollectionId),
|
|
mlog.Uint64("ts", result.GetControlChannelResult().TimeTick),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfAliasAlterable(ctx context.Context, dbName string, alias string, collectionName string) error {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
// backward compatibility for rolling upgrade
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty", mlog.String("alias", alias), mlog.String("collection", collectionName))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
// It's ok that we don't read from catalog when cache missed.
|
|
// Since cache always keep the latest version, and the ts should always be the latest.
|
|
|
|
if !mt.names.exist(dbName) {
|
|
return merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
|
|
if collID, ok := mt.names.get(dbName, alias); ok {
|
|
coll := mt.collID2Meta[collID]
|
|
// allow alias with dropping&dropped
|
|
if coll.State != pb.CollectionState_CollectionDropping && coll.State != pb.CollectionState_CollectionDropped {
|
|
return merr.WrapErrAliasCollectionNameConflict(dbName, alias)
|
|
}
|
|
}
|
|
|
|
collectionID, ok := mt.names.get(dbName, collectionName)
|
|
if !ok {
|
|
// you cannot alias to a non-existent collection.
|
|
return merr.WrapErrCollectionNotFound(collectionName)
|
|
}
|
|
|
|
coll, ok := mt.collID2Meta[collectionID]
|
|
if !ok || !coll.Available() {
|
|
// you cannot alias to a non-existent collection.
|
|
return merr.WrapErrCollectionNotFound(collectionName)
|
|
}
|
|
|
|
// check if alias exists.
|
|
existAliasCollectionID, ok := mt.aliases.get(dbName, alias)
|
|
if !ok {
|
|
return merr.WrapErrAliasNotFound(dbName, alias)
|
|
}
|
|
if existAliasCollectionID == collectionID {
|
|
return errIgnoredAlterAlias
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) DescribeAlias(ctx context.Context, dbName string, alias string, ts Timestamp) (string, error) {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty", mlog.String("alias", alias))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
// check if database exists.
|
|
dbExist := mt.aliases.exist(dbName)
|
|
if !dbExist {
|
|
return "", merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
// check if alias exists.
|
|
collectionID, ok := mt.aliases.get(dbName, alias)
|
|
if !ok {
|
|
return "", merr.WrapErrAliasNotFound(dbName, alias)
|
|
}
|
|
|
|
collectionMeta, ok := mt.collID2Meta[collectionID]
|
|
if !ok {
|
|
return "", merr.WrapErrCollectionIDOfAliasNotFound(collectionID)
|
|
}
|
|
if collectionMeta.State == pb.CollectionState_CollectionCreated {
|
|
return collectionMeta.Name, nil
|
|
}
|
|
return "", merr.WrapErrAliasNotFound(dbName, alias)
|
|
}
|
|
|
|
func (mt *MetaTable) ListAliases(ctx context.Context, dbName string, collectionName string, ts Timestamp) ([]string, error) {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
if dbName == "" {
|
|
mlog.Warn(ctx, "db name is empty", mlog.String("collection", collectionName))
|
|
dbName = util.DefaultDBName
|
|
}
|
|
|
|
// check if database exists.
|
|
dbExist := mt.aliases.exist(dbName)
|
|
if !dbExist {
|
|
return nil, merr.WrapErrDatabaseNotFound(dbName)
|
|
}
|
|
var aliases []string
|
|
if collectionName == "" {
|
|
collections := mt.aliases.listCollections(dbName)
|
|
for name, collectionID := range collections {
|
|
if collectionMeta, ok := mt.collID2Meta[collectionID]; ok &&
|
|
collectionMeta.State == pb.CollectionState_CollectionCreated {
|
|
aliases = append(aliases, name)
|
|
}
|
|
}
|
|
} else {
|
|
collectionID, exist := mt.names.get(dbName, collectionName)
|
|
collectionMeta, exist2 := mt.collID2Meta[collectionID]
|
|
if exist && exist2 && collectionMeta.State == pb.CollectionState_CollectionCreated {
|
|
aliases = mt.listAliasesByID(collectionID)
|
|
} else {
|
|
return nil, merr.WrapErrCollectionNotFound(collectionName)
|
|
}
|
|
}
|
|
return aliases, nil
|
|
}
|
|
|
|
func (mt *MetaTable) IsAlias(ctx context.Context, db, name string) bool {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
_, ok := mt.aliases.get(db, name)
|
|
return ok
|
|
}
|
|
|
|
func (mt *MetaTable) listAliasesByID(collID UniqueID) []string {
|
|
ret := make([]string, 0)
|
|
mt.aliases.iterate(func(db string, collection string, id UniqueID) bool {
|
|
if collID == id {
|
|
ret = append(ret, collection)
|
|
}
|
|
return true
|
|
})
|
|
return ret
|
|
}
|
|
|
|
func (mt *MetaTable) ListAliasesByID(ctx context.Context, collID UniqueID) []string {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
return mt.listAliasesByID(collID)
|
|
}
|
|
|
|
// GetGeneralCount gets the general count(sum of product of partition number and shard number).
|
|
func (mt *MetaTable) GetGeneralCount(ctx context.Context) int {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
return mt.generalCnt
|
|
}
|
|
|
|
func (mt *MetaTable) InitCredential(ctx context.Context) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
credInfo, err := mt.catalog.GetCredential(ctx, util.UserRoot)
|
|
if err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return err
|
|
}
|
|
if credInfo != nil {
|
|
return nil
|
|
}
|
|
encryptedRootPassword, err := crypto.PasswordEncrypt(Params.CommonCfg.DefaultRootPassword.GetValue())
|
|
if err != nil {
|
|
mlog.Warn(ctx, "RootCoord init user root failed", mlog.Err(err))
|
|
return err
|
|
}
|
|
mlog.Info(ctx, "RootCoord init user root")
|
|
err = mt.catalog.AlterCredential(ctx, &model.Credential{
|
|
Username: util.UserRoot,
|
|
EncryptedPassword: encryptedRootPassword,
|
|
})
|
|
if err != nil {
|
|
mlog.Warn(ctx, "RootCoord init user root failed", mlog.Err(err))
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfAddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
|
|
if funcutil.IsEmptyString(credInfo.GetUsername()) {
|
|
return merr.WrapErrParameterInvalidMsg("username is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
usernames, err := mt.catalog.ListCredentials(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// check if the username already exists.
|
|
for _, username := range usernames {
|
|
if username == credInfo.GetUsername() {
|
|
return errUserAlreadyExists
|
|
}
|
|
}
|
|
|
|
// check if the number of users has reached the limit.
|
|
maxUserNum := Params.ProxyCfg.MaxUserNum.GetAsInt()
|
|
if len(usernames) >= maxUserNum {
|
|
errMsg := "unable to add user because the number of users has reached the limit"
|
|
mlog.Error(ctx, errMsg, mlog.Int("maxUserNum", maxUserNum))
|
|
return merr.WrapErrServiceQuotaExceeded(errMsg)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfUpdateCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
|
|
if funcutil.IsEmptyString(credInfo.GetUsername()) {
|
|
return merr.WrapErrParameterInvalidMsg("username is empty")
|
|
}
|
|
hasEncryptedPassword := credInfo.GetEncryptedPassword() != ""
|
|
hasSha256Password := credInfo.GetSha256Password() != ""
|
|
if hasEncryptedPassword != hasSha256Password {
|
|
return merr.WrapErrParameterInvalidMsg("credential password update must include both encrypted and sha256 password")
|
|
}
|
|
if !hasEncryptedPassword && !hasSha256Password && credInfo.Description == nil {
|
|
return merr.WrapErrParameterInvalidMsg("credential update must change password or description")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
// check if the number of credential exists.
|
|
if _, err := mt.catalog.GetCredential(ctx, credInfo.GetUsername()); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errUserNotFound
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// AlterCredential update credential
|
|
func (mt *MetaTable) AlterCredential(ctx context.Context, result message.BroadcastResultAlterUserMessageV2) error {
|
|
body := result.Message.MustBody()
|
|
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
existsCredential, err := mt.catalog.GetCredential(ctx, body.CredentialInfo.Username)
|
|
if err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return err
|
|
}
|
|
// if the credential already exists and the version is not greater than the current timetick.
|
|
if existsCredential != nil && existsCredential.TimeTick >= result.GetControlChannelResult().TimeTick {
|
|
mlog.Info(ctx, "credential already exists and the version is not greater than the current timetick",
|
|
mlog.String("username", body.CredentialInfo.Username),
|
|
mlog.Uint64("incoming", result.GetControlChannelResult().TimeTick),
|
|
mlog.Uint64("current", existsCredential.TimeTick),
|
|
)
|
|
return nil
|
|
}
|
|
encryptedPassword := body.CredentialInfo.EncryptedPassword
|
|
description := body.CredentialInfo.GetDescription()
|
|
if existsCredential != nil {
|
|
if encryptedPassword == "" {
|
|
encryptedPassword = existsCredential.EncryptedPassword
|
|
}
|
|
if body.CredentialInfo.Description == nil {
|
|
description = existsCredential.Description
|
|
}
|
|
}
|
|
credential := &model.Credential{
|
|
Username: body.CredentialInfo.Username,
|
|
EncryptedPassword: encryptedPassword,
|
|
Description: description,
|
|
TimeTick: result.GetControlChannelResult().TimeTick,
|
|
}
|
|
return mt.catalog.AlterCredential(ctx, credential)
|
|
}
|
|
|
|
// GetCredential get credential by username
|
|
func (mt *MetaTable) GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
credential, err := mt.catalog.GetCredential(ctx, username)
|
|
return model.MarshalCredentialModel(credential), err
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfDeleteCredential(ctx context.Context, req *milvuspb.DeleteCredentialRequest) error {
|
|
if funcutil.IsEmptyString(req.GetUsername()) {
|
|
return merr.WrapErrParameterInvalidMsg("username is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
// check if the number of credential exists.
|
|
if _, err := mt.catalog.GetCredential(ctx, req.GetUsername()); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errUserNotFound
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DeleteCredential delete credential
|
|
func (mt *MetaTable) DeleteCredential(ctx context.Context, result message.BroadcastResultDropUserMessageV2) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
existsCredential, err := mt.catalog.GetCredential(ctx, result.Message.Header().UserName)
|
|
if err != nil && !errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return err
|
|
}
|
|
// if the credential already exists and the version is not greater than the current timetick.
|
|
if existsCredential != nil && existsCredential.TimeTick >= result.GetControlChannelResult().TimeTick {
|
|
mlog.Info(ctx, "credential already exists and the version is not greater than the current timetick",
|
|
mlog.String("username", result.Message.Header().UserName),
|
|
mlog.Uint64("incoming", result.GetControlChannelResult().TimeTick),
|
|
mlog.Uint64("current", existsCredential.TimeTick),
|
|
)
|
|
return nil
|
|
}
|
|
return mt.catalog.DropCredential(ctx, result.Message.Header().UserName)
|
|
}
|
|
|
|
// ListCredentialUsernames list credential usernames
|
|
func (mt *MetaTable) ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
usernames, err := mt.catalog.ListCredentials(ctx)
|
|
if err != nil {
|
|
return nil, merr.Wrap(err, "failed to list credential usernames")
|
|
}
|
|
return &milvuspb.ListCredUsersResponse{Usernames: usernames}, nil
|
|
}
|
|
|
|
// CheckIfCreateRole checks if the role can be created.
|
|
func (mt *MetaTable) CheckIfCreateRole(ctx context.Context, in *milvuspb.CreateRoleRequest) error {
|
|
if funcutil.IsEmptyString(in.GetEntity().GetName()) {
|
|
return merr.WrapErrParameterInvalidMsg("role name is empty")
|
|
}
|
|
if err := validateRoleDescription(in.GetEntity().GetDescription()); err != nil {
|
|
return err
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
results, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
|
|
if err != nil {
|
|
mlog.Warn(ctx, "fail to list roles", mlog.Err(err))
|
|
return err
|
|
}
|
|
for _, result := range results {
|
|
if result.GetRole().GetName() == in.GetEntity().GetName() {
|
|
mlog.Info(ctx, "role already exists", mlog.String("role", in.GetEntity().GetName()))
|
|
return errRoleAlreadyExists
|
|
}
|
|
}
|
|
if len(results) >= Params.ProxyCfg.MaxRoleNum.GetAsInt() {
|
|
errMsg := "unable to create role because the number of roles has reached the limit"
|
|
mlog.Warn(ctx, errMsg, mlog.Int("max_role_num", Params.ProxyCfg.MaxRoleNum.GetAsInt()))
|
|
return merr.WrapErrServiceQuotaExceeded(errMsg)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CreateRole create role
|
|
func (mt *MetaTable) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.CreateRole(ctx, tenant, entity)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfAlterRole(ctx context.Context, in *milvuspb.AlterRoleRequest) error {
|
|
if funcutil.IsEmptyString(in.GetRoleName()) {
|
|
return merr.WrapErrParameterInvalidMsg("role name is empty")
|
|
}
|
|
if util.IsBuiltinRole(in.GetRoleName()) || lo.Contains(util.DefaultRoles, in.GetRoleName()) {
|
|
return merr.WrapErrPrivilegeNotPermitted("the role[%s] is a builtin role, which can't be altered", in.GetRoleName())
|
|
}
|
|
if err := validateRoleDescription(in.GetDescription()); err != nil {
|
|
return err
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
if _, err := mt.catalog.ListRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.GetRoleName()}, false); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errRoleNotExists
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) AlterRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
|
|
if funcutil.IsEmptyString(entity.GetName()) {
|
|
return merr.WrapErrParameterInvalidMsg("role name is empty")
|
|
}
|
|
if util.IsBuiltinRole(entity.GetName()) || lo.Contains(util.DefaultRoles, entity.GetName()) {
|
|
return merr.WrapErrPrivilegeNotPermitted("the role[%s] is a builtin role, which can't be altered", entity.GetName())
|
|
}
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
if _, err := mt.catalog.ListRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.GetName()}, false); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errRoleNotExists
|
|
}
|
|
return err
|
|
}
|
|
return mt.catalog.AlterRole(ctx, tenant, entity)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfDropRole(ctx context.Context, in *milvuspb.DropRoleRequest) error {
|
|
if funcutil.IsEmptyString(in.GetRoleName()) {
|
|
return merr.WrapErrParameterInvalidMsg("role name is empty")
|
|
}
|
|
if util.IsBuiltinRole(in.GetRoleName()) {
|
|
return merr.WrapErrPrivilegeNotPermitted("the role[%s] is a builtin role, which can't be dropped", in.GetRoleName())
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
if _, err := mt.catalog.ListRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.GetRoleName()}, false); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errRoleNotExists
|
|
}
|
|
return err
|
|
}
|
|
if in.GetForceDrop() {
|
|
return nil
|
|
}
|
|
|
|
grantEntities, err := mt.catalog.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
|
|
Role: &milvuspb.RoleEntity{Name: in.GetRoleName()},
|
|
DbName: "*",
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(grantEntities) != 0 {
|
|
errMsg := "fail to drop the role that it has privileges. Use REVOKE API to revoke privileges"
|
|
return merr.WrapErrParameterInvalidMsg(errMsg)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateRoleDescription(description string) error {
|
|
return rbacutil.ValidateRoleDescription(description, Params.ProxyCfg.MaxRoleDescriptionLength.GetAsInt())
|
|
}
|
|
|
|
// DropRole drop role info
|
|
func (mt *MetaTable) DropRole(ctx context.Context, tenant string, roleName string) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.DropRole(ctx, tenant, roleName)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfOperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) error {
|
|
if funcutil.IsEmptyString(req.GetUsername()) {
|
|
return merr.WrapErrParameterInvalidMsg("username in the user entity is empty")
|
|
}
|
|
if funcutil.IsEmptyString(req.GetRoleName()) {
|
|
return merr.WrapErrParameterInvalidMsg("role name in the role entity is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
if _, err := mt.catalog.ListRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: req.RoleName}, false); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return errRoleNotExists
|
|
}
|
|
return err
|
|
}
|
|
if req.Type != milvuspb.OperateUserRoleType_RemoveUserFromRole {
|
|
if _, err := mt.catalog.ListUser(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: req.Username}, false); err != nil {
|
|
if errors.Is(err, merr.ErrIoKeyNotFound) {
|
|
return merr.WrapErrParameterInvalidMsg("user %q not found", req.GetUsername())
|
|
}
|
|
return merr.Wrap(err, "failed to check user existence")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// OperateUserRole operate the relationship between a user and a role, including adding a user to a role and removing a user from a role
|
|
func (mt *MetaTable) OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.AlterUserRole(ctx, tenant, userEntity, roleEntity, operateType)
|
|
}
|
|
|
|
// SelectRole select role.
|
|
// Enter the role condition by the entity param. And this param is nil, which means selecting all roles.
|
|
// Get all users that are added to the role by setting the includeUserInfo param to true.
|
|
func (mt *MetaTable) SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListRole(ctx, tenant, entity, includeUserInfo)
|
|
}
|
|
|
|
// SelectUser select user.
|
|
// Enter the user condition by the entity param. And this param is nil, which means selecting all users.
|
|
// Get all roles that are added the user to by setting the includeRoleInfo param to true.
|
|
func (mt *MetaTable) SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListUser(ctx, tenant, entity, includeRoleInfo)
|
|
}
|
|
|
|
// OperatePrivilege grant or revoke privilege by setting the operateType param
|
|
func (mt *MetaTable) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
|
|
if funcutil.IsEmptyString(entity.ObjectName) {
|
|
return merr.WrapErrParameterInvalidMsg("the object name in the grant entity is empty")
|
|
}
|
|
if entity.Object == nil || funcutil.IsEmptyString(entity.Object.Name) {
|
|
return merr.WrapErrParameterInvalidMsg("the object entity in the grant entity is invalid")
|
|
}
|
|
if entity.Role == nil || funcutil.IsEmptyString(entity.Role.Name) {
|
|
return merr.WrapErrParameterInvalidMsg("the role entity in the grant entity is invalid")
|
|
}
|
|
if entity.Grantor == nil {
|
|
return merr.WrapErrParameterInvalidMsg("the grantor in the grant entity is empty")
|
|
}
|
|
if entity.Grantor.Privilege == nil || funcutil.IsEmptyString(entity.Grantor.Privilege.Name) {
|
|
return merr.WrapErrParameterInvalidMsg("the privilege name in the grant entity is empty")
|
|
}
|
|
if entity.Grantor.User == nil || funcutil.IsEmptyString(entity.Grantor.User.Name) {
|
|
return merr.WrapErrParameterInvalidMsg("the grantor name in the grant entity is empty")
|
|
}
|
|
if !funcutil.IsRevoke(operateType) && !funcutil.IsGrant(operateType) {
|
|
return merr.WrapErrParameterInvalidMsg("the operate type in the grant entity is invalid")
|
|
}
|
|
if entity.DbName == "" {
|
|
entity.DbName = util.DefaultDBName
|
|
}
|
|
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.AlterGrant(ctx, tenant, entity, operateType)
|
|
}
|
|
|
|
// SelectGrant select grant
|
|
// The principal entity MUST be not empty in the grant entity
|
|
// The resource entity and the resource name are optional, and the two params should be not empty together when you select some grants about the resource kind.
|
|
func (mt *MetaTable) SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
|
|
var entities []*milvuspb.GrantEntity
|
|
if entity == nil {
|
|
return entities, merr.WrapErrParameterInvalidMsg("the grant entity is nil")
|
|
}
|
|
|
|
if entity.Role == nil || funcutil.IsEmptyString(entity.Role.Name) {
|
|
return entities, merr.WrapErrParameterInvalidMsg("the role entity in the grant entity is invalid")
|
|
}
|
|
if entity.DbName == "" {
|
|
entity.DbName = util.DefaultDBName
|
|
}
|
|
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListGrant(ctx, tenant, entity)
|
|
}
|
|
|
|
func (mt *MetaTable) DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
|
|
if role == nil || funcutil.IsEmptyString(role.Name) {
|
|
return merr.WrapErrParameterInvalidMsg("the role entity is invalid when dropping the grant")
|
|
}
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.DeleteGrant(ctx, tenant, role)
|
|
}
|
|
|
|
func (mt *MetaTable) ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListPolicy(ctx, tenant)
|
|
}
|
|
|
|
func (mt *MetaTable) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListUserRole(ctx, tenant)
|
|
}
|
|
|
|
func (mt *MetaTable) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.BackupRBAC(ctx, tenant)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfRBACRestorable(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest) error {
|
|
meta := req.GetRBACMeta()
|
|
if len(meta.GetRoles()) == 0 && len(meta.GetPrivilegeGroups()) == 0 && len(meta.GetGrants()) == 0 && len(meta.GetUsers()) == 0 {
|
|
return errEmptyRBACMeta
|
|
}
|
|
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
// check if role already exists
|
|
existRoles, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
existRoleMap := lo.SliceToMap(existRoles, func(entity *milvuspb.RoleResult) (string, struct{}) { return entity.GetRole().GetName(), struct{}{} })
|
|
existRoleAfterRestoreMap := lo.SliceToMap(existRoles, func(entity *milvuspb.RoleResult) (string, struct{}) { return entity.GetRole().GetName(), struct{}{} })
|
|
for _, role := range meta.GetRoles() {
|
|
if err := validateRoleDescription(role.GetDescription()); err != nil {
|
|
return err
|
|
}
|
|
if _, ok := existRoleMap[role.GetName()]; ok {
|
|
return merr.WrapErrParameterInvalidMsg("role [%s] already exists", role.GetName())
|
|
}
|
|
existRoleAfterRestoreMap[role.GetName()] = struct{}{}
|
|
}
|
|
|
|
// check if privilege group already exists
|
|
existPrivGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
existPrivGroupMap := lo.SliceToMap(existPrivGroups, func(entity *milvuspb.PrivilegeGroupInfo) (string, struct{}) { return entity.GetGroupName(), struct{}{} })
|
|
existPrivGroupAfterRestoreMap := lo.SliceToMap(existPrivGroups, func(entity *milvuspb.PrivilegeGroupInfo) (string, struct{}) { return entity.GetGroupName(), struct{}{} })
|
|
for _, group := range meta.GetPrivilegeGroups() {
|
|
if _, ok := existPrivGroupMap[group.GetGroupName()]; ok {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group [%s] already exists", group.GetGroupName())
|
|
}
|
|
existPrivGroupAfterRestoreMap[group.GetGroupName()] = struct{}{}
|
|
}
|
|
|
|
// check if grant can be restored
|
|
for _, grant := range meta.GetGrants() {
|
|
privName := grant.GetGrantor().GetPrivilege().GetName()
|
|
if util.IsAnyWord(privName) {
|
|
continue
|
|
}
|
|
if _, ok := existPrivGroupAfterRestoreMap[privName]; !ok && !util.IsPrivilegeNameDefined(privName) {
|
|
return merr.WrapErrParameterInvalidMsg("privilege [%s] does not exist", privName)
|
|
}
|
|
}
|
|
|
|
// check if user can be restored
|
|
existUser, err := mt.catalog.ListUser(ctx, util.DefaultTenant, nil, false)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
existUserMap := lo.SliceToMap(existUser, func(entity *milvuspb.UserResult) (string, struct{}) { return entity.GetUser().GetName(), struct{}{} })
|
|
for _, user := range meta.GetUsers() {
|
|
if _, ok := existUserMap[user.GetUser()]; ok {
|
|
return merr.WrapErrParameterInvalidMsg("user [%s] already exists", user.GetUser())
|
|
}
|
|
|
|
// check if user-role can be restored
|
|
for _, role := range user.GetRoles() {
|
|
if _, ok := existRoleAfterRestoreMap[role.GetName()]; !ok {
|
|
return merr.WrapErrParameterInvalidMsg("role [%s] does not exist", role.GetName())
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.RestoreRBAC(ctx, tenant, meta)
|
|
}
|
|
|
|
// check if the privilege group name is defined by users
|
|
func (mt *MetaTable) IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error) {
|
|
privGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
for _, group := range privGroups {
|
|
if group.GroupName == groupName {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfPrivilegeGroupCreatable(ctx context.Context, req *milvuspb.CreatePrivilegeGroupRequest) error {
|
|
if funcutil.IsEmptyString(req.GetGroupName()) {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group name is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, req.GetGroupName())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if definedByUsers {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group name [%s] is defined by users", req.GetGroupName())
|
|
}
|
|
if util.IsPrivilegeNameDefined(req.GetGroupName()) {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group name [%s] is defined by built in privileges or privilege groups in system", req.GetGroupName())
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
data := &milvuspb.PrivilegeGroupInfo{
|
|
GroupName: groupName,
|
|
Privileges: make([]*milvuspb.PrivilegeEntity, 0),
|
|
}
|
|
return mt.catalog.SavePrivilegeGroup(ctx, data)
|
|
}
|
|
|
|
func (mt *MetaTable) CheckIfPrivilegeGroupDropable(ctx context.Context, req *milvuspb.DropPrivilegeGroupRequest) error {
|
|
if funcutil.IsEmptyString(req.GetGroupName()) {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group name is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, req.GetGroupName())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !definedByUsers {
|
|
return errNotCustomPrivilegeGroup
|
|
}
|
|
|
|
// check if the group is used by any role
|
|
roles, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
roleEntity := lo.Map(roles, func(entity *milvuspb.RoleResult, _ int) *milvuspb.RoleEntity {
|
|
return entity.GetRole()
|
|
})
|
|
for _, role := range roleEntity {
|
|
grants, err := mt.catalog.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
|
|
Role: role,
|
|
DbName: util.AnyWord,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, grant := range grants {
|
|
if grant.Grantor.Privilege.Name == req.GetGroupName() {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group [%s] is used by role [%s], Use REVOKE API to revoke it first", req.GetGroupName(), role.GetName())
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) DropPrivilegeGroup(ctx context.Context, groupName string) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
return mt.catalog.DropPrivilegeGroup(ctx, groupName)
|
|
}
|
|
|
|
func (mt *MetaTable) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
return mt.catalog.ListPrivilegeGroups(ctx)
|
|
}
|
|
|
|
// CheckIfPrivilegeGroupAlterable checks if the privilege group can be altered.
|
|
func (mt *MetaTable) CheckIfPrivilegeGroupAlterable(ctx context.Context, req *milvuspb.OperatePrivilegeGroupRequest) error {
|
|
if funcutil.IsEmptyString(req.GetGroupName()) {
|
|
return merr.WrapErrParameterInvalidMsg("privilege group name is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
groups, err := mt.catalog.ListPrivilegeGroups(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
currenctGroups := lo.SliceToMap(groups, func(group *milvuspb.PrivilegeGroupInfo) (string, []*milvuspb.PrivilegeEntity) {
|
|
return group.GroupName, group.Privileges
|
|
})
|
|
// check if the privilege group is defined by users
|
|
if _, ok := currenctGroups[req.GroupName]; !ok {
|
|
return merr.WrapErrParameterInvalidMsg("there is no privilege group name [%s] defined in system to operate", req.GroupName)
|
|
}
|
|
|
|
if len(req.Privileges) == 0 {
|
|
return merr.WrapErrParameterInvalidMsg("privileges is empty when alter the privilege group")
|
|
}
|
|
// check if the new incoming privileges are defined by users or built in
|
|
for _, p := range req.Privileges {
|
|
if util.IsPrivilegeNameDefined(p.Name) {
|
|
continue
|
|
}
|
|
if _, ok := currenctGroups[p.Name]; !ok {
|
|
return merr.WrapErrParameterInvalidMsg("there is no privilege name or privilege group name [%s] defined in system to operate", p.Name)
|
|
}
|
|
}
|
|
|
|
if req.Type == milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup {
|
|
// Check if all privileges are the same privilege level
|
|
privilegeLevels := lo.SliceToMap(lo.Union(req.Privileges, currenctGroups[req.GroupName]), func(p *milvuspb.PrivilegeEntity) (string, struct{}) {
|
|
return util.GetPrivilegeLevel(p.Name), struct{}{}
|
|
})
|
|
if len(privilegeLevels) > 1 {
|
|
return merr.WrapErrParameterInvalidMsg("privileges are not the same privilege level")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
|
|
mt.permissionLock.Lock()
|
|
defer mt.permissionLock.Unlock()
|
|
|
|
// merge with current privileges
|
|
group, err := mt.catalog.GetPrivilegeGroup(ctx, groupName)
|
|
if err != nil {
|
|
mlog.Warn(ctx, "fail to get privilege group", mlog.String("privilege_group", groupName), mlog.Err(err))
|
|
return err
|
|
}
|
|
privSet := lo.SliceToMap(group.Privileges, func(p *milvuspb.PrivilegeEntity) (string, struct{}) {
|
|
return p.Name, struct{}{}
|
|
})
|
|
switch operateType {
|
|
case milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup:
|
|
for _, p := range privileges {
|
|
privSet[p.Name] = struct{}{}
|
|
}
|
|
case milvuspb.OperatePrivilegeGroupType_RemovePrivilegesFromGroup:
|
|
for _, p := range privileges {
|
|
delete(privSet, p.Name)
|
|
}
|
|
default:
|
|
mlog.Warn(ctx, "unsupported operate type", mlog.Any("operate_type", operateType))
|
|
return merr.WrapErrParameterInvalidMsg("unsupported operate type: %v", operateType)
|
|
}
|
|
|
|
mergedPrivs := lo.Map(lo.Keys(privSet), func(priv string, _ int) *milvuspb.PrivilegeEntity {
|
|
return &milvuspb.PrivilegeEntity{Name: priv}
|
|
})
|
|
data := &milvuspb.PrivilegeGroupInfo{
|
|
GroupName: groupName,
|
|
Privileges: mergedPrivs,
|
|
}
|
|
return mt.catalog.SavePrivilegeGroup(ctx, data)
|
|
}
|
|
|
|
func (mt *MetaTable) GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
|
|
if funcutil.IsEmptyString(groupName) {
|
|
return nil, merr.WrapErrParameterInvalidMsg("the privilege group name is empty")
|
|
}
|
|
mt.permissionLock.RLock()
|
|
defer mt.permissionLock.RUnlock()
|
|
|
|
// get all roles
|
|
roles, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
roleEntity := lo.Map(roles, func(entity *milvuspb.RoleResult, _ int) *milvuspb.RoleEntity {
|
|
return entity.GetRole()
|
|
})
|
|
|
|
rolesMap := make(map[*milvuspb.RoleEntity]struct{})
|
|
for _, role := range roleEntity {
|
|
grants, err := mt.catalog.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
|
|
Role: role,
|
|
DbName: util.AnyWord,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, grant := range grants {
|
|
if grant.Grantor.Privilege.Name == groupName {
|
|
rolesMap[role] = struct{}{}
|
|
}
|
|
}
|
|
}
|
|
return lo.Keys(rolesMap), nil
|
|
}
|
|
|
|
func (mt *MetaTable) AddFileResource(ctx context.Context, resource *internalpb.FileResourceInfo) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
if old, ok := mt.fileResourceName2Meta[resource.Name]; ok {
|
|
if old.Path == resource.Path {
|
|
return nil
|
|
}
|
|
return merr.WrapErrParameterInvalidMsg("file resource %s already exists", resource.Name)
|
|
}
|
|
|
|
err := mt.catalog.SaveFileResource(ctx, resource, mt.fileResourceVersion+1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mt.fileResourceName2Meta[resource.Name] = resource
|
|
mt.fileResourceID2Meta[resource.Id] = resource
|
|
mt.fileResourceVersion++
|
|
return nil
|
|
}
|
|
|
|
func (mt *MetaTable) RemoveFileResource(ctx context.Context, name string) (error, bool) {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
|
|
if resource, ok := mt.fileResourceName2Meta[name]; ok {
|
|
if mt.fileResourceRefCnt[resource.Id] > 0 {
|
|
return merr.WrapErrParameterInvalidMsg("file resource %s is still in use: %d", resource.Name, mt.fileResourceRefCnt[resource.Id]), false
|
|
}
|
|
|
|
err := mt.catalog.RemoveFileResource(ctx, resource.Id, mt.fileResourceVersion+1)
|
|
if err != nil {
|
|
return err, false
|
|
}
|
|
|
|
delete(mt.fileResourceName2Meta, resource.Name)
|
|
delete(mt.fileResourceID2Meta, resource.Id)
|
|
delete(mt.fileResourceRefCnt, resource.Id)
|
|
mt.fileResourceVersion++
|
|
return nil, true
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
func (mt *MetaTable) ListFileResource(ctx context.Context) ([]*internalpb.FileResourceInfo, uint64) {
|
|
mt.ddLock.RLock()
|
|
defer mt.ddLock.RUnlock()
|
|
|
|
return lo.Values(mt.fileResourceID2Meta), mt.fileResourceVersion
|
|
}
|
|
|
|
// IncFileResourceRefCnt increments refCnt for file resources, binding them to a
|
|
// collection being created. Under ddLock, atomic with RemoveFileResource.
|
|
// Returns error if any resource ID does not exist.
|
|
func (mt *MetaTable) IncFileResourceRefCnt(ids []int64) error {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
for _, id := range ids {
|
|
if _, ok := mt.fileResourceID2Meta[id]; !ok {
|
|
return merr.WrapErrParameterInvalidMsg("file resource %d not found", id)
|
|
}
|
|
}
|
|
for _, id := range ids {
|
|
mt.fileResourceRefCnt[id]++
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DecFileResourceRefCnt decrements refCnt. Used for early release when
|
|
// CreateCollection fails after validation.
|
|
func (mt *MetaTable) DecFileResourceRefCnt(ids []int64) {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
for _, id := range ids {
|
|
if mt.fileResourceRefCnt[id] > 0 {
|
|
mt.fileResourceRefCnt[id]--
|
|
} else {
|
|
mlog.Warn(context.TODO(), "DecFileResourceRefCnt underflow", mlog.Int64("id", id))
|
|
}
|
|
}
|
|
}
|
|
|
|
// RecoverFileResourceRefCnt re-increments refCnt for file resources referenced by
|
|
// pending CreateCollection broadcast tasks whose collections have not yet been
|
|
// persisted. Called during startup before rootcoord becomes Healthy.
|
|
func (mt *MetaTable) RecoverFileResourceRefCnt(pendingCollections map[int64][]int64) {
|
|
mt.ddLock.Lock()
|
|
defer mt.ddLock.Unlock()
|
|
for collID, resourceIds := range pendingCollections {
|
|
if _, exists := mt.collID2Meta[collID]; exists {
|
|
continue // collection already persisted, reload already counted it
|
|
}
|
|
for _, id := range resourceIds {
|
|
if _, ok := mt.fileResourceID2Meta[id]; ok {
|
|
mt.fileResourceRefCnt[id]++
|
|
} else {
|
|
mlog.Warn(context.TODO(), "RecoverFileResourceRefCnt: pending task references missing file resource",
|
|
mlog.Int64("collectionID", collID), mlog.Int64("resourceID", id))
|
|
}
|
|
}
|
|
}
|
|
}
|