Remove db from global meta cache after drop database (#24474)

Signed-off-by: Njaime <yun.zhang@zilliz.com>

internal/proxy/impl_test.go

@@ -20,11 +20,10 @@ import (
 	"context"
 	"testing"
 
-	"google.golang.org/grpc/metadata"
-
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/milvus-io/milvus-proto/go-api/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
@@ -641,6 +640,7 @@ func TestProxyDropDatabase(t *testing.T) {
 	err = node.sched.Start()
 	assert.NoError(t, err)
 	defer node.sched.Close()
+	globalMetaCache, _ = NewMetaCache(nil, nil, nil)
 
 	t.Run("drop database fail", func(t *testing.T) {
 		rc := mocks.NewRootCoord(t)

internal/proxy/meta_cache.go

@@ -80,6 +80,8 @@ type Cache interface {
 	GetUserRole(username string) []string
 	RefreshPolicyInfo(op typeutil.CacheOp) error
 	InitPolicyInfo(info []string, userRoles []string)
+
+	RemoveDatabase(ctx context.Context, database string)
 }
 
 type collectionInfo struct {
@@ -940,3 +942,9 @@ func (m *MetaCache) RefreshPolicyInfo(op typeutil.CacheOp) error {
 	}
 	return nil
 }
+
+func (m *MetaCache) RemoveDatabase(ctx context.Context, database string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.collInfo, database)
+}

internal/proxy/mock_cache_test.go

@@ -731,6 +731,35 @@ func (_c *MockCache_RemoveCredential_Call) Return() *MockCache_RemoveCredential_
 	return _c
 }
 
+// RemoveDatabase provides a mock function with given fields: ctx, database
+func (_m *MockCache) RemoveDatabase(ctx context.Context, database string) {
+	_m.Called(ctx, database)
+}
+
+// MockCache_RemoveDatabase_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveDatabase'
+type MockCache_RemoveDatabase_Call struct {
+	*mock.Call
+}
+
+// RemoveDatabase is a helper method to define mock.On call
+//   - ctx context.Context
+//   - database string
+func (_e *MockCache_Expecter) RemoveDatabase(ctx interface{}, database interface{}) *MockCache_RemoveDatabase_Call {
+	return &MockCache_RemoveDatabase_Call{Call: _e.mock.On("RemoveDatabase", ctx, database)}
+}
+
+func (_c *MockCache_RemoveDatabase_Call) Run(run func(ctx context.Context, database string)) *MockCache_RemoveDatabase_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(string))
+	})
+	return _c
+}
+
+func (_c *MockCache_RemoveDatabase_Call) Return() *MockCache_RemoveDatabase_Call {
+	_c.Call.Return()
+	return _c
+}
+
 // RemovePartition provides a mock function with given fields: ctx, database, collectionName, partitionName
 func (_m *MockCache) RemovePartition(ctx context.Context, database string, collectionName string, partitionName string) {
 	_m.Called(ctx, database, collectionName, partitionName)

internal/proxy/rootcoord_mock_test.go

@@ -24,27 +24,21 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/milvus-io/milvus/internal/common"
-	"github.com/milvus-io/milvus/internal/types"
-
-	"github.com/milvus-io/milvus/internal/util/funcutil"
-
-	"github.com/milvus-io/milvus/internal/util/uniquegenerator"
-
 	"github.com/golang/protobuf/proto"
 
-	"github.com/milvus-io/milvus/internal/util/milvuserrors"
-
-	"github.com/milvus-io/milvus-proto/go-api/schemapb"
-
-	"github.com/milvus-io/milvus/internal/util/typeutil"
-
 	"github.com/milvus-io/milvus-proto/go-api/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
+	"github.com/milvus-io/milvus-proto/go-api/schemapb"
+	"github.com/milvus-io/milvus/internal/common"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/proxypb"
 	"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/internal/util/funcutil"
+	"github.com/milvus-io/milvus/internal/util/milvuserrors"
+	"github.com/milvus-io/milvus/internal/util/typeutil"
+	"github.com/milvus-io/milvus/internal/util/uniquegenerator"
 )
 
 type collectionMeta struct {
@@ -1166,7 +1160,7 @@ func (coord *RootCoordMock) CreateDatabase(ctx context.Context, in *milvuspb.Cre
 }
 
 func (coord *RootCoordMock) DropDatabase(ctx context.Context, in *milvuspb.DropDatabaseRequest) (*commonpb.Status, error) {
-	return &commonpb.Status{}, nil
+	return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
 }
 
 func (coord *RootCoordMock) ListDatabases(ctx context.Context, in *milvuspb.ListDatabasesRequest) (*milvuspb.ListDatabasesResponse, error) {
@@ -1185,10 +1179,15 @@ func (coord *RootCoordMock) RenameCollection(ctx context.Context, req *milvuspb.
 }
 
 type DescribeCollectionFunc func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error)
+
 type ShowPartitionsFunc func(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error)
+
 type ShowSegmentsFunc func(ctx context.Context, request *milvuspb.ShowSegmentsRequest) (*milvuspb.ShowSegmentsResponse, error)
+
 type DescribeSegmentsFunc func(ctx context.Context, request *rootcoordpb.DescribeSegmentsRequest) (*rootcoordpb.DescribeSegmentsResponse, error)
+
 type ImportFunc func(ctx context.Context, req *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error)
+
 type DropCollectionFunc func(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error)
 
 type GetGetCredentialFunc func(ctx context.Context, req *rootcoordpb.GetCredentialRequest) (*rootcoordpb.GetCredentialResponse, error)

internal/proxy/task_database.go

@@ -126,6 +126,10 @@ func (ddt *dropDatabaseTask) Execute(ctx context.Context) error {
 	var err error
 	ddt.result = &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}
 	ddt.result, err = ddt.rootCoord.DropDatabase(ctx, ddt.DropDatabaseRequest)
+
+	if ddt.result != nil && ddt.result.ErrorCode == commonpb.ErrorCode_Success {
+		globalMetaCache.RemoveDatabase(ctx, ddt.DbName)
+	}
 	return err
 }
 

internal/proxy/task_database_test.go

@@ -87,6 +87,8 @@ func TestDropDatabaseTask(t *testing.T) {
 		assert.Equal(t, Timestamp(100), task.BeginTs())
 		assert.Equal(t, Timestamp(100), task.EndTs())
 		assert.Equal(t, "db", task.GetDbName())
+
+		globalMetaCache, _ = NewMetaCache(nil, nil, nil)
 		err = task.Execute(ctx)
 		assert.NoError(t, err)