infoschema_v2: support create/alter/drop policy (#51594)

ref #50959
pingcap · Mar 8, 2024 · 13304a2 · 13304a2
1 parent 72182b2
commit 13304a2
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 36 deletions.
diff --git a/pkg/infoschema/builder_misc.go b/pkg/infoschema/builder_misc.go
@@ -23,7 +23,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/model"
 )
 
-func (b *Builder) applyCreatePolicy(m *meta.Meta, diff *model.SchemaDiff) error {
+func applyCreatePolicy(b *Builder, m *meta.Meta, diff *model.SchemaDiff) error {
 	po, err := m.GetPolicy(diff.SchemaID)
 	if err != nil {
 		return errors.Trace(err)
@@ -44,7 +44,7 @@ func (b *Builder) applyCreatePolicy(m *meta.Meta, diff *model.SchemaDiff) error
 	return nil
 }
 
-func (b *Builder) applyAlterPolicy(m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
+func applyAlterPolicy(b *Builder, m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
 	po, err := m.GetPolicy(diff.SchemaID)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -62,7 +62,7 @@ func (b *Builder) applyAlterPolicy(m *meta.Meta, diff *model.SchemaDiff) ([]int6
 	return []int64{}, nil
 }
 
-func (b *Builder) applyDropPolicy(PolicyID int64) []int64 {
+func applyDropPolicy(b *Builder, PolicyID int64) []int64 {
 	po, ok := b.infoSchema.PolicyByID(PolicyID)
 	if !ok {
 		return nil

diff --git a/pkg/infoschema/infoschema_v2.go b/pkg/infoschema/infoschema_v2.go
@@ -609,27 +609,6 @@ func applyModifySchemaDefaultPlacement(b *Builder, m *meta.Meta, diff *model.Sch
 	return b.applyModifySchemaDefaultPlacement(m, diff)
 }
 
-func applyCreatePolicy(b *Builder, m *meta.Meta, diff *model.SchemaDiff) error {
-	if b.enableV2 {
-		return b.applyCreatePolicyV2(m, diff)
-	}
-	return b.applyCreatePolicy(m, diff)
-}
-
-func applyDropPolicy(b *Builder, PolicyID int64) []int64 {
-	if b.enableV2 {
-		return b.applyDropPolicyV2(PolicyID)
-	}
-	return b.applyDropPolicy(PolicyID)
-}
-
-func applyAlterPolicy(b *Builder, m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
-	if b.enableV2 {
-		return b.applyAlterPolicyV2(m, diff)
-	}
-	return b.applyAlterPolicy(m, diff)
-}
-
 func applyDropTableOrPartition(b *Builder, m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
 	if b.enableV2 {
 		// return b.applyDropTableOrPartitionV2(m, diff)
@@ -725,18 +704,6 @@ func (b *Builder) applyModifySchemaDefaultPlacementV2(m *meta.Meta, diff *model.
 	panic("TODO")
 }
 
-func (b *Builder) applyCreatePolicyV2(m *meta.Meta, diff *model.SchemaDiff) error {
-	panic("TODO")
-}
-
-func (b *Builder) applyAlterPolicyV2(m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
-	panic("TODO")
-}
-
-func (b *Builder) applyDropPolicyV2(PolicyID int64) []int64 {
-	panic("TODO")
-}
-
 func (b *Builder) applyTruncateTableOrPartitionV2(m *meta.Meta, diff *model.SchemaDiff) ([]int64, error) {
 	panic("TODO")
 }

diff --git a/pkg/infoschema/infoschema_v2_test.go b/pkg/infoschema/infoschema_v2_test.go
@@ -135,4 +135,58 @@ func TestMisc(t *testing.T) {
 	require.True(t, ok)
 	require.Equal(t, resourceGroupInfo2, getResourceGroupInfo)
 	require.NoError(t, txn.Rollback())
+
+	// test create policy
+	policyInfo := internal.MockPolicyInfo(t, r.Store(), "test")
+	internal.CreatePolicy(t, r.Store(), policyInfo)
+	txn, err = r.Store().Begin()
+	require.NoError(t, err)
+	err = applyCreatePolicy(builder, meta.NewMeta(txn), &model.SchemaDiff{SchemaID: policyInfo.ID})
+	require.NoError(t, err)
+	is = builder.Build()
+	require.Len(t, is.AllPlacementPolicies(), 1)
+	getPolicyInfo, ok := is.PolicyByName(policyInfo.Name)
+	require.True(t, ok)
+	require.Equal(t, policyInfo, getPolicyInfo)
+	require.NoError(t, txn.Rollback())
+
+	// create another policy
+	policyInfo2 := internal.MockPolicyInfo(t, r.Store(), "test2")
+	internal.CreatePolicy(t, r.Store(), policyInfo2)
+	txn, err = r.Store().Begin()
+	require.NoError(t, err)
+	err = applyCreatePolicy(builder, meta.NewMeta(txn), &model.SchemaDiff{SchemaID: policyInfo2.ID})
+	require.NoError(t, err)
+	is = builder.Build()
+	require.Len(t, is.AllPlacementPolicies(), 2)
+	getPolicyInfo, ok = is.PolicyByName(policyInfo2.Name)
+	require.True(t, ok)
+	require.Equal(t, policyInfo2, getPolicyInfo)
+	require.NoError(t, txn.Rollback())
+
+	// test alter policy
+	policyInfo.State = model.StatePublic
+	internal.UpdatePolicy(t, r.Store(), policyInfo)
+	txn, err = r.Store().Begin()
+	require.NoError(t, err)
+	err = applyCreatePolicy(builder, meta.NewMeta(txn), &model.SchemaDiff{SchemaID: policyInfo.ID})
+	require.NoError(t, err)
+	is = builder.Build()
+	require.Len(t, is.AllPlacementPolicies(), 2)
+	getPolicyInfo, ok = is.PolicyByName(policyInfo.Name)
+	require.True(t, ok)
+	require.Equal(t, policyInfo, getPolicyInfo)
+	require.NoError(t, txn.Rollback())
+
+	// test drop policy
+	internal.DropPolicy(t, r.Store(), policyInfo)
+	txn, err = r.Store().Begin()
+	require.NoError(t, err)
+	_ = applyDropPolicy(builder, policyInfo.ID)
+	is = builder.Build()
+	require.Len(t, is.AllPlacementPolicies(), 1)
+	getPolicyInfo, ok = is.PolicyByName(policyInfo2.Name)
+	require.True(t, ok)
+	require.Equal(t, policyInfo2, getPolicyInfo)
+	require.NoError(t, txn.Rollback())
 }
diff --git a/pkg/infoschema/internal/testkit.go b/pkg/infoschema/internal/testkit.go
@@ -190,6 +190,16 @@ func MockResourceGroupInfo(t *testing.T, store kv.Storage, groupName string) *mo
 	}
 }
 
+// MockPolicyInfo mock policy for testing.
+func MockPolicyInfo(t *testing.T, store kv.Storage, policyName string) *model.PolicyInfo {
+	id, err := GenGlobalID(store)
+	require.NoError(t, err)
+	return &model.PolicyInfo{
+		ID:   id,
+		Name: model.NewCIStr(policyName),
+	}
+}
+
 // AddTable add mock table for testing.
 func AddTable(t *testing.T, store kv.Storage, dbInfo *model.DBInfo, tblInfo *model.TableInfo) {
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL)
@@ -267,3 +277,36 @@ func DropResourceGroup(t *testing.T, store kv.Storage, group *model.ResourceGrou
 	})
 	require.NoError(t, err)
 }
+
+// CreatePolicy create mock policy for testing.
+func CreatePolicy(t *testing.T, store kv.Storage, policy *model.PolicyInfo) {
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL)
+	err := kv.RunInNewTxn(ctx, store, true, func(ctx context.Context, txn kv.Transaction) error {
+		err := meta.NewMeta(txn).CreatePolicy(policy)
+		require.NoError(t, err)
+		return errors.Trace(err)
+	})
+	require.NoError(t, err)
+}
+
+// UpdatePolicy update mock policy for testing.
+func UpdatePolicy(t *testing.T, store kv.Storage, policy *model.PolicyInfo) {
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL)
+	err := kv.RunInNewTxn(ctx, store, true, func(ctx context.Context, txn kv.Transaction) error {
+		err := meta.NewMeta(txn).UpdatePolicy(policy)
+		require.NoError(t, err)
+		return errors.Trace(err)
+	})
+	require.NoError(t, err)
+}
+
+// DropPolicy drop mock policy for testing.
+func DropPolicy(t *testing.T, store kv.Storage, policy *model.PolicyInfo) {
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL)
+	err := kv.RunInNewTxn(ctx, store, true, func(ctx context.Context, txn kv.Transaction) error {
+		err := meta.NewMeta(txn).DropPolicy(policy.ID)
+		require.NoError(t, err)
+		return errors.Trace(err)
+	})
+	require.NoError(t, err)
+}