ddl: support create and drop placement policy

ddl/ddl.go

@@ -115,6 +115,8 @@ type DDL interface {
 	CreateSequence(ctx sessionctx.Context, stmt *ast.CreateSequenceStmt) error
 	DropSequence(ctx sessionctx.Context, tableIdent ast.Ident, ifExists bool) (err error)
 	AlterSequence(ctx sessionctx.Context, stmt *ast.AlterSequenceStmt) error
+	CreatePlacementPolicy(ctx sessionctx.Context, stmt *ast.CreatePlacementPolicyStmt) error
+	DropPlacementPolicy(ctx sessionctx.Context, stmt *ast.DropPlacementPolicyStmt) error
 
 	// CreateSchemaWithInfo creates a database (schema) given its database info.
 	//

ddl/ddl_api.go

@@ -57,6 +57,7 @@ import (
 	"github.com/pingcap/tidb/util/domainutil"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/mock"
+	"github.com/pingcap/tidb/util/placementpolicy"
 	"github.com/pingcap/tidb/util/set"
 	"go.uber.org/zap"
 )
@@ -6144,3 +6145,93 @@ func (d *ddl) AlterTablePartitionAttributes(ctx sessionctx.Context, ident ast.Id
 	err = d.callHookOnChanged(err)
 	return errors.Trace(err)
 }
+
+func buildPolicyInfo(stmt *ast.CreatePlacementPolicyStmt) (*placementpolicy.PolicyInfo, error) {
+	policyInfo := &placementpolicy.PolicyInfo{}
+	policyInfo.Name = stmt.PolicyName
+	for _, opt := range stmt.PlacementOptions {
+		switch opt.Tp {
+		case ast.PlacementOptionPrimaryRegion:
+			policyInfo.PrimaryRegion = opt.StrValue
+		case ast.PlacementOptionRegions:
+			policyInfo.Regions = opt.StrValue
+		case ast.PlacementOptionFollowerCount:
+			policyInfo.Followers = opt.UintValue
+		case ast.PlacementOptionVoterCount:
+			policyInfo.Voters = opt.UintValue
+		case ast.PlacementOptionLearnerCount:
+			policyInfo.Learners = opt.UintValue
+		case ast.PlacementOptionSchedule:
+			policyInfo.Schedule = opt.StrValue
+		case ast.PlacementOptionConstraints:
+			policyInfo.Constraints = opt.StrValue
+		case ast.PlacementOptionLearnerConstraints:
+			policyInfo.LearnerConstraints = opt.StrValue
+		case ast.PlacementOptionFollowerConstraints:
+			policyInfo.FollowerConstraints = opt.StrValue
+		case ast.PlacementOptionVoterConstraints:
+			policyInfo.VoterConstraints = opt.StrValue
+		case ast.PlacementOptionLeaderConstraints:
+			policyInfo.LeaderConstraints = opt.StrValue
+		default:
+			return nil, errors.Trace(errors.New("unknown placement policy option"))
+		}
+	}
+	return policyInfo, nil
+}
+
+func (d *ddl) CreatePlacementPolicy(ctx sessionctx.Context, stmt *ast.CreatePlacementPolicyStmt) (err error) {
+	policyName := stmt.PolicyName
+	is := d.GetInfoSchemaWithInterceptor(ctx)
+	// Check policy existence.
+	_, ok := is.PolicyByName(policyName)
+	if ok {
+		err = infoschema.ErrPlacementPolicyExists.GenWithStackByArgs(policyName)
+		if stmt.IfNotExists {
+			ctx.GetSessionVars().StmtCtx.AppendNote(err)
+			return nil
+		}
+		return err
+	}
+	// Auto fill the policyID when it is inserted.
+	policyInfo, err := buildPolicyInfo(stmt)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	job := &model.Job{
+		SchemaName: policyInfo.Name.L,
+		Type:       model.ActionCreatePlacementPolicy,
+		BinlogInfo: &model.HistoryInfo{},
+		Args:       []interface{}{policyInfo},
+	}
+	err = d.doDDLJob(ctx, job)
+	err = d.callHookOnChanged(err)
+	return errors.Trace(err)
+}
+
+func (d *ddl) DropPlacementPolicy(ctx sessionctx.Context, stmt *ast.DropPlacementPolicyStmt) (err error) {
+	policyName := stmt.PolicyName
+	is := d.GetInfoSchemaWithInterceptor(ctx)
+	// Check policy existence.
+	policy, ok := is.PolicyByName(policyName)
+	if !ok {
+		err = infoschema.ErrPlacementPolicyNotExists.GenWithStackByArgs(policyName)
+		if stmt.IfExists {
+			ctx.GetSessionVars().StmtCtx.AppendNote(err)
+			return nil
+		}
+		return err
+	}
+
+	job := &model.Job{
+		SchemaID:   policy.ID,
+		SchemaName: policy.Name.L,
+		Type:       model.ActionDropPlacementPolicy,
+		BinlogInfo: &model.HistoryInfo{},
+		Args:       []interface{}{policyName},
+	}
+	err = d.doDDLJob(ctx, job)
+	err = d.callHookOnChanged(err)
+	return errors.Trace(err)
+}

ddl/ddl_worker.go

@@ -834,6 +834,10 @@ func (w *worker) runDDLJob(d *ddlCtx, t *meta.Meta, job *model.Job) (ver int64,
 		ver, err = onAlterTableAttributes(t, job)
 	case model.ActionAlterTablePartitionAttributes:
 		ver, err = onAlterTablePartitionAttributes(t, job)
+	case model.ActionCreatePlacementPolicy:
+		ver, err = onCreatePlacementPolicy(d, t, job)
+	case model.ActionDropPlacementPolicy:
+		ver, err = onDropPlacementPolicy(t, job)
 	default:
 		// Invalid job, cancel it.
 		job.State = model.JobStateCancelled

ddl/placement_policy.go

@@ -0,0 +1,166 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed 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 ddl
+
+import (
+	"fmt"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/parser/model"
+	"github.com/pingcap/tidb/infoschema"
+	"github.com/pingcap/tidb/meta"
+	"github.com/pingcap/tidb/util/placementpolicy"
+)
+
+func onCreatePlacementPolicy(d *ddlCtx, t *meta.Meta, job *model.Job) (ver int64, _ error) {
+	policyInfo := &placementpolicy.PolicyInfo{}
+	if err := job.DecodeArgs(policyInfo); err != nil {
+		job.State = model.JobStateCancelled
+		return ver, errors.Trace(err)
+	}
+	policyInfo.State = model.StateNone
+
+	err := checkPlacementPolicyNotExistAndCancelExistJob(d, t, job, policyInfo)
+	if err != nil {
+		return ver, errors.Trace(err)
+	}
+	switch policyInfo.State {
+	case model.StateNone:
+		// none -> public
+		policyInfo.State = model.StatePublic
+		err = t.CreatePolicy(policyInfo)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		job.SchemaID = policyInfo.ID
+
+		ver, err = updateSchemaVersion(t, job)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		// Finish this job.
+		job.FinishDBJob(model.JobStateDone, model.StatePublic, ver, nil)
+		return ver, nil
+	default:
+		// We can't enter here.
+		return ver, ErrInvalidDDLState.GenWithStackByArgs("policy", policyInfo.State)
+	}
+}
+
+func getPolicyInfo(t *meta.Meta, policyID int64) (*placementpolicy.PolicyInfo, error) {
+	policy, err := t.GetPolicy(policyID)
+	if err != nil {
+		if meta.ErrPolicyNotExists.Equal(err) {
+			return nil, infoschema.ErrPlacementPolicyNotExists.GenWithStackByArgs(
+				fmt.Sprintf("(Policy ID %d)", policyID),
+			)
+		}
+		return nil, err
+	}
+	return policy, nil
+}
+
+func checkPlacementPolicyNotExistAndCancelExistJob(d *ddlCtx, t *meta.Meta, job *model.Job, info *placementpolicy.PolicyInfo) error {
+	currVer, err := t.GetSchemaVersion()
+	if err != nil {
+		return err
+	}
+	is := d.infoCache.GetLatest()
+	if is.SchemaMetaVersion() == currVer {
+		// Use cached policy.
+		_, ok := is.PolicyByName(info.Name)
+		if ok {
+			job.State = model.JobStateCancelled
+			return infoschema.ErrPlacementPolicyExists.GenWithStackByArgs(info.Name)
+		}
+		return nil
+	}
+	// Check in meta directly.
+	policies, err := t.ListPolicies()
+	if err != nil {
+		return errors.Trace(err)
+	}
+	for _, policy := range policies {
+		if policy.Name.L == info.Name.L {
+			job.State = model.JobStateCancelled
+			return infoschema.ErrPlacementPolicyExists.GenWithStackByArgs(info.Name)
+		}
+	}
+	return nil
+}
+
+func checkPlacementPolicyExistAndCancelNonExistJob(t *meta.Meta, job *model.Job, policyID int64) (*placementpolicy.PolicyInfo, error) {
+	policy, err := getPolicyInfo(t, policyID)
+	if err == nil {
+		return policy, nil
+	}
+	if infoschema.ErrPlacementPolicyNotExists.Equal(err) {
+		job.State = model.JobStateCancelled
+	}
+	return nil, err
+}
+
+func onDropPlacementPolicy(t *meta.Meta, job *model.Job) (ver int64, _ error) {
+	policyInfo, err := checkPlacementPolicyExistAndCancelNonExistJob(t, job, job.SchemaID)
+	if err != nil {
+		return ver, errors.Trace(err)
+	}
+	switch policyInfo.State {
+	case model.StatePublic:
+		// public -> write only
+		policyInfo.State = model.StateWriteOnly
+		err = t.UpdatePolicy(policyInfo)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		ver, err = updateSchemaVersion(t, job)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		// Update the job state when all affairs done.
+		job.SchemaState = model.StateWriteOnly
+	case model.StateWriteOnly:
+		// write only -> delete only
+		policyInfo.State = model.StateDeleteOnly
+		err = t.UpdatePolicy(policyInfo)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		ver, err = updateSchemaVersion(t, job)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		// Update the job state when all affairs done.
+		job.SchemaState = model.StateDeleteOnly
+	case model.StateDeleteOnly:
+		policyInfo.State = model.StateNone
+		if err = t.DropPolicy(policyInfo.ID); err != nil {
+			return ver, errors.Trace(err)
+		}
+		ver, err = updateSchemaVersion(t, job)
+		if err != nil {
+			return ver, errors.Trace(err)
+		}
+		// TODO: Reset all the policy reference, (modify meta & notify pd)
+		// If any partitions currently use this policy, they will be converted to the policy used by the table
+		// they belong to. If any databases use this policy, they will be converted to the default placement_policy policy.
+
+		// Finish this job. By now policy don't consider the binlog sync.
+		job.FinishDBJob(model.JobStateDone, model.StateNone, ver, nil)
+	default:
+		err = ErrInvalidDDLState.GenWithStackByArgs("policy", policyInfo.State)
+	}
+	return ver, errors.Trace(err)
+}