From 4b826a0925b8a7c1618f0efad83b03620699cf63 Mon Sep 17 00:00:00 2001 From: sylzd Date: Mon, 29 Nov 2021 16:43:52 +0800 Subject: [PATCH] dumpling: Add support for `Create Placement` (#29724) --- dumpling/export/dump.go | 29 ++++++++++ dumpling/export/dump_test.go | 1 + dumpling/export/sql.go | 34 ++++++++++++ dumpling/export/sql_test.go | 54 +++++++++++++++++++ dumpling/export/task.go | 22 +++++++- dumpling/export/writer.go | 2 + dumpling/install.sh | 2 +- dumpling/tests/basic/run.sh | 2 +- .../result/x-placement-policy-create.sql | 3 ++ .../result/x1-placement-policy-create.sql | 3 ++ dumpling/tests/placement_policy/run.sh | 30 +++++++++++ 11 files changed, 179 insertions(+), 3 deletions(-) create mode 100644 dumpling/tests/placement_policy/result/x-placement-policy-create.sql create mode 100644 dumpling/tests/placement_policy/result/x1-placement-policy-create.sql create mode 100644 dumpling/tests/placement_policy/run.sh diff --git a/dumpling/export/dump.go b/dumpling/export/dump.go index 39637aa4f1b89..de7eddcfa3b44 100755 --- a/dumpling/export/dump.go +++ b/dumpling/export/dump.go @@ -8,6 +8,7 @@ import ( "database/sql" "encoding/hex" "fmt" + "github.com/go-sql-driver/mysql" "math/big" "sort" "strconv" @@ -298,6 +299,34 @@ func (d *Dumper) startWriters(tctx *tcontext.Context, wg *errgroup.Group, taskCh func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskChan chan<- Task) error { conf := d.conf allTables := conf.Tables + + // policy should be created before database + // placement policy in other server type can be different, so we only handle the tidb server + if conf.ServerInfo.ServerType == version.ServerTypeTiDB { + policyNames, err := ListAllPlacementPolicyNames(metaConn) + if err != nil { + errCause := errors.Cause(err) + if mysqlErr, ok := errCause.(*mysql.MySQLError); ok && mysqlErr.Number == ErrNoSuchTable { + // some old tidb version and other server type doesn't support placement rules, we can skip it. + tctx.L().Debug("cannot dump placement policy, maybe the server doesn't support it", log.ShortError(err)) + } else { + tctx.L().Warn("fail to dump placement policy: ", log.ShortError(err)) + } + } + for _, policy := range policyNames { + createPolicySQL, err := ShowCreatePlacementPolicy(metaConn, policy) + if err != nil { + return err + } + wrappedCreatePolicySQL := fmt.Sprintf("/*T![placement] %s */", createPolicySQL) + task := NewTaskPolicyMeta(policy, wrappedCreatePolicySQL) + ctxDone := d.sendTaskToChan(tctx, task, taskChan) + if ctxDone { + return tctx.Err() + } + } + } + for dbName, tables := range allTables { if !conf.NoSchemas { createDatabaseSQL, err := ShowCreateDatabase(metaConn, dbName) diff --git a/dumpling/export/dump_test.go b/dumpling/export/dump_test.go index 1626908f511e6..284cec9435017 100644 --- a/dumpling/export/dump_test.go +++ b/dumpling/export/dump_test.go @@ -55,6 +55,7 @@ func TestDumpBlock(t *testing.T) { taskChan := make(chan Task, 1) taskChan <- &TaskDatabaseMeta{} d.conf.Tables = DatabaseTables{}.AppendTable(database, nil) + d.conf.ServerInfo.ServerType = version.ServerTypeMySQL require.ErrorIs(t, d.dumpDatabases(writerCtx, conn, taskChan), context.Canceled) require.ErrorIs(t, wg.Wait(), writerErr) } diff --git a/dumpling/export/sql.go b/dumpling/export/sql.go index 4b0203a1665da..1fb5688338b2e 100644 --- a/dumpling/export/sql.go +++ b/dumpling/export/sql.go @@ -92,6 +92,21 @@ func ShowCreateTable(db *sql.Conn, database, table string) (string, error) { return oneRow[1], nil } +// ShowCreatePlacementPolicy constructs the create policy SQL for a specified table +// returns (createPoilicySQL, error) +func ShowCreatePlacementPolicy(db *sql.Conn, policy string) (string, error) { + var oneRow [2]string + handleOneRow := func(rows *sql.Rows) error { + return rows.Scan(&oneRow[0], &oneRow[1]) + } + query := fmt.Sprintf("SHOW CREATE PLACEMENT POLICY `%s`", escapeString(policy)) + err := simpleQuery(db, query, handleOneRow) + if err != nil { + return "", errors.Annotatef(err, "sql: %s", query) + } + return oneRow[1], nil +} + // ShowCreateView constructs the create view SQL for a specified view // returns (createFakeTableSQL, createViewSQL, error) func ShowCreateView(db *sql.Conn, database, view string) (createFakeTableSQL string, createRealViewSQL string, err error) { @@ -279,6 +294,25 @@ func ListAllDatabasesTables(tctx *tcontext.Context, db *sql.Conn, databaseNames return dbTables, nil } +func ListAllPlacementPolicyNames(db *sql.Conn) ([]string, error) { + var policyList []string + var policy string + const query = "select distinct policy_name from information_schema.placement_rules where policy_name is not null;" + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + return policyList, errors.Annotatef(err, "sql: %s", query) + } + defer rows.Close() + for rows.Next() { + err := rows.Scan(&policy) + if err != nil { + return policyList, errors.Trace(err) + } + policyList = append(policyList, policy) + } + return policyList, errors.Annotatef(rows.Err(), "sql: %s", query) +} + // SelectVersion gets the version information from the database server func SelectVersion(db *sql.DB) (string, error) { var versionInfo string diff --git a/dumpling/export/sql_test.go b/dumpling/export/sql_test.go index b6bd835d27e35..6a69c7dc5c164 100644 --- a/dumpling/export/sql_test.go +++ b/dumpling/export/sql_test.go @@ -9,6 +9,7 @@ import ( "encoding/csv" "encoding/json" "fmt" + "github.com/go-sql-driver/mysql" "io" "os" "path" @@ -340,6 +341,59 @@ func TestShowCreateView(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestShowCreatePolicy(t *testing.T) { + t.Parallel() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { + require.NoError(t, db.Close()) + }() + + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + + mock.ExpectQuery("SHOW CREATE PLACEMENT POLICY `policy_x`"). + WillReturnRows(sqlmock.NewRows([]string{"Policy", "Create Policy"}). + AddRow("policy_x", "CREATE PLACEMENT POLICY `policy_x` LEARNERS=1")) + + createPolicySQL, err := ShowCreatePlacementPolicy(conn, "policy_x") + require.NoError(t, err) + require.Equal(t, "CREATE PLACEMENT POLICY `policy_x` LEARNERS=1", createPolicySQL) + require.NoError(t, mock.ExpectationsWereMet()) + +} + +func TestListPolicyNames(t *testing.T) { + t.Parallel() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { + require.NoError(t, db.Close()) + }() + + conn, err := db.Conn(context.Background()) + require.NoError(t, err) + + mock.ExpectQuery("select distinct policy_name from information_schema.placement_rules where policy_name is not null;"). + WillReturnRows(sqlmock.NewRows([]string{"policy_name"}). + AddRow("policy_x")) + policies, err := ListAllPlacementPolicyNames(conn) + require.NoError(t, err) + require.Equal(t, []string{"policy_x"}, policies) + require.NoError(t, mock.ExpectationsWereMet()) + + // some old tidb version doesn't support placement rules returns error + expectedErr := &mysql.MySQLError{Number: ErrNoSuchTable, Message: "Table 'information_schema.placement_rules' doesn't exist"} + mock.ExpectExec("select distinct policy_name from information_schema.placement_rules where policy_name is not null;"). + WillReturnError(expectedErr) + policies, err = ListAllPlacementPolicyNames(conn) + if mysqlErr, ok := err.(*mysql.MySQLError); ok { + require.Equal(t, mysqlErr.Number, ErrNoSuchTable) + } +} + func TestGetSuitableRows(t *testing.T) { t.Parallel() diff --git a/dumpling/export/task.go b/dumpling/export/task.go index 8cfec59b24859..a9e5874fa350b 100644 --- a/dumpling/export/task.go +++ b/dumpling/export/task.go @@ -4,7 +4,7 @@ package export import "fmt" -// Task is a file dump task for dumpling, it could either be dumping database/table/view metadata, table data +// Task is a file dump task for dumpling, it could either be dumping database/table/view/policy metadata, table data type Task interface { // Brief is the brief for a dumping task Brief() string @@ -34,6 +34,13 @@ type TaskViewMeta struct { CreateViewSQL string } +// TaskPolicyMeta is a dumping view metadata task +type TaskPolicyMeta struct { + Task + PolicyName string + CreatePolicySQL string +} + // TaskTableData is a dumping table data task type TaskTableData struct { Task @@ -70,6 +77,14 @@ func NewTaskViewMeta(dbName, tblName, createTableSQL, createViewSQL string) *Tas } } +// NewTaskPolicyMeta returns a new dumping placement policy metadata task +func NewTaskPolicyMeta(policyName, createPolicySQL string) *TaskPolicyMeta { + return &TaskPolicyMeta{ + PolicyName: policyName, + CreatePolicySQL: createPolicySQL, + } +} + // NewTaskTableData returns a new dumping table data task func NewTaskTableData(meta TableMeta, data TableDataIR, currentChunk, totalChunks int) *TaskTableData { return &TaskTableData{ @@ -95,6 +110,11 @@ func (t *TaskViewMeta) Brief() string { return fmt.Sprintf("meta of view '%s'.'%s'", t.DatabaseName, t.ViewName) } +// Brief implements task.Brief +func (t *TaskPolicyMeta) Brief() string { + return fmt.Sprintf("meta of placement policy '%s'", t.PolicyName) +} + // Brief implements task.Brief func (t *TaskTableData) Brief() string { db, tbl := t.Meta.DatabaseName(), t.Meta.TableName() diff --git a/dumpling/export/writer.go b/dumpling/export/writer.go index 4d8d5301db929..7520ff670462f 100644 --- a/dumpling/export/writer.go +++ b/dumpling/export/writer.go @@ -99,6 +99,8 @@ func (w *Writer) handleTask(task Task) error { return w.WriteTableMeta(t.DatabaseName, t.TableName, t.CreateTableSQL) case *TaskViewMeta: return w.WriteViewMeta(t.DatabaseName, t.ViewName, t.CreateTableSQL, t.CreateViewSQL) + case *TaskPolicyMeta: + return w.WritePolicyMeta(t.PolicyName, t.CreatePolicySQL) case *TaskTableData: err := w.WriteTableData(t.Meta, t.Data, t.ChunkIndex) if err != nil { diff --git a/dumpling/install.sh b/dumpling/install.sh index e7eb9a8e2f687..e51ee861b099c 100644 --- a/dumpling/install.sh +++ b/dumpling/install.sh @@ -17,7 +17,7 @@ cd $pwd/tidb-lightning && make cd $pwd mv tidb-lightning/bin/tidb-lightning bin/ -TIDB_TAG="v4.0.4" +TIDB_TAG="master" # download tidb-server git clone -b $TIDB_TAG https://github.com/pingcap/tidb cd $pwd/tidb && make diff --git a/dumpling/tests/basic/run.sh b/dumpling/tests/basic/run.sh index 6caccce221433..d801970643b75 100644 --- a/dumpling/tests/basic/run.sh +++ b/dumpling/tests/basic/run.sh @@ -90,7 +90,7 @@ echo "expected 2, actual ${actual}" # Test for tidb_mem_quota_query configuration export GO_FAILPOINTS="github.com/pingcap/tidb/dumpling/export/PrintTiDBMemQuotaQuery=1*return" -run_dumpling > ${DUMPLING_OUTPUT_DIR}/dumpling.log +run_dumpling | tee ${DUMPLING_OUTPUT_DIR}/dumpling.log actual=$(grep -w "tidb_mem_quota_query == 1073741824" ${DUMPLING_OUTPUT_DIR}/dumpling.log|wc -l) echo "expected 1, actual ${actual}" [ "$actual" = 1 ] diff --git a/dumpling/tests/placement_policy/result/x-placement-policy-create.sql b/dumpling/tests/placement_policy/result/x-placement-policy-create.sql new file mode 100644 index 0000000000000..0b68c742ee1cf --- /dev/null +++ b/dumpling/tests/placement_policy/result/x-placement-policy-create.sql @@ -0,0 +1,3 @@ +/*!40101 SET NAMES binary*/; +/*T![placement] SET PLACEMENT_CHECKS = 0*/; +/*T![placement] CREATE PLACEMENT POLICY `x` PRIMARY_REGION="cn-east-1" REGIONS="cn-east-1,cn-east" */; diff --git a/dumpling/tests/placement_policy/result/x1-placement-policy-create.sql b/dumpling/tests/placement_policy/result/x1-placement-policy-create.sql new file mode 100644 index 0000000000000..d2dafe29b2a85 --- /dev/null +++ b/dumpling/tests/placement_policy/result/x1-placement-policy-create.sql @@ -0,0 +1,3 @@ +/*!40101 SET NAMES binary*/; +/*T![placement] SET PLACEMENT_CHECKS = 0*/; +/*T![placement] CREATE PLACEMENT POLICY `x1` FOLLOWERS=4 */; diff --git a/dumpling/tests/placement_policy/run.sh b/dumpling/tests/placement_policy/run.sh new file mode 100644 index 0000000000000..b7653fc921378 --- /dev/null +++ b/dumpling/tests/placement_policy/run.sh @@ -0,0 +1,30 @@ +#!/bin/sh +# +# Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. + +set -eu + +export DUMPLING_TEST_PORT=4000 + +run_sql "drop database if exists policy" +run_sql "drop placement policy if exists x" +run_sql "drop placement policy if exists x1" +run_sql "create database policy" + +export DUMPLING_TEST_DATABASE="policy" + +run_sql 'CREATE PLACEMENT POLICY x PRIMARY_REGION="cn-east-1" REGIONS="cn-east-1,cn-east";' +run_sql 'CREATE PLACEMENT POLICY x1 FOLLOWERS=4;' + +run_dumpling + +file_should_exist "$DUMPLING_OUTPUT_DIR/policy-schema-create.sql" +file_should_exist "$DUMPLING_OUTPUT_DIR/x-placement-policy-create.sql" +file_should_exist "$DUMPLING_OUTPUT_DIR/x1-placement-policy-create.sql" + +diff "$DUMPLING_BASE_NAME/result/x-placement-policy-create.sql" "$DUMPLING_OUTPUT_DIR/x-placement-policy-create.sql" +diff "$DUMPLING_BASE_NAME/result/x1-placement-policy-create.sql" "$DUMPLING_OUTPUT_DIR/x1-placement-policy-create.sql" + +run_sql "drop database if exists policy" +run_sql "drop placement policy if exists x" +run_sql "drop placement policy if exists x1"