dumpling: supports foreign key limitations (#39914)

close #39913

dumpling/export/dump.go

@@ -1219,9 +1219,7 @@ func dumpTableMeta(tctx *tcontext.Context, conf *Config, conn *BaseConn, db stri
 		selectedField:    selectField,
 		selectedLen:      selectLen,
 		hasImplicitRowID: hasImplicitRowID,
-		specCmts: []string{
-			"/*!40101 SET NAMES binary*/;",
-		},
+		specCmts:         getSpecialComments(conf.ServerInfo.ServerType),
 	}
 
 	if conf.NoSchemas {

dumpling/export/ir.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/br/pkg/version"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
 )
 
@@ -85,7 +86,7 @@ type MetaIR interface {
 	MetaSQL() string
 }
 
-func setTableMetaFromRows(rows *sql.Rows) (TableMeta, error) {
+func setTableMetaFromRows(serverType version.ServerType, rows *sql.Rows) (TableMeta, error) {
 	tps, err := rows.ColumnTypes()
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -101,6 +102,6 @@ func setTableMetaFromRows(rows *sql.Rows) (TableMeta, error) {
 		colTypes:      tps,
 		selectedField: strings.Join(nms, ","),
 		selectedLen:   len(nms),
-		specCmts:      []string{"/*!40101 SET NAMES binary*/;"},
+		specCmts:      getSpecialComments(serverType),
 	}, nil
 }

dumpling/export/ir_impl.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/br/pkg/version"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
 	"go.uber.org/zap"
 )
@@ -370,3 +371,22 @@ func (td *multiQueriesChunk) Close() error {
 func (*multiQueriesChunk) RawRows() *sql.Rows {
 	return nil
 }
+
+var serverSpecialComments = map[version.ServerType][]string{
+	version.ServerTypeMySQL: {
+		"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;",
+		"/*!40101 SET NAMES binary*/;",
+	},
+	version.ServerTypeTiDB: {
+		"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;",
+		"/*!40101 SET NAMES binary*/;",
+	},
+	version.ServerTypeMariaDB: {
+		"/*!40101 SET NAMES binary*/;",
+		"SET FOREIGN_KEY_CHECKS=0;",
+	},
+}
+
+func getSpecialComments(serverType version.ServerType) []string {
+	return serverSpecialComments[serverType]
+}

dumpling/export/writer.go

@@ -133,7 +133,7 @@ func (w *Writer) WritePolicyMeta(policy, createSQL string) error {
 	if err != nil {
 		return err
 	}
-	return writeMetaToFile(tctx, "placement-policy", createSQL, w.extStorage, fileName+".sql", conf.CompressType)
+	return w.writeMetaToFile(tctx, "placement-policy", createSQL, fileName+".sql")
 }
 
 // WriteDatabaseMeta writes database meta to a file
@@ -143,7 +143,7 @@ func (w *Writer) WriteDatabaseMeta(db, createSQL string) error {
 	if err != nil {
 		return err
 	}
-	return writeMetaToFile(tctx, db, createSQL, w.extStorage, fileName+".sql", conf.CompressType)
+	return w.writeMetaToFile(tctx, db, createSQL, fileName+".sql")
 }
 
 // WriteTableMeta writes table meta to a file
@@ -153,7 +153,7 @@ func (w *Writer) WriteTableMeta(db, table, createSQL string) error {
 	if err != nil {
 		return err
 	}
-	return writeMetaToFile(tctx, db, createSQL, w.extStorage, fileName+".sql", conf.CompressType)
+	return w.writeMetaToFile(tctx, db, createSQL, fileName+".sql")
 }
 
 // WriteViewMeta writes view meta to a file
@@ -167,11 +167,11 @@ func (w *Writer) WriteViewMeta(db, view, createTableSQL, createViewSQL string) e
 	if err != nil {
 		return err
 	}
-	err = writeMetaToFile(tctx, db, createTableSQL, w.extStorage, fileNameTable+".sql", conf.CompressType)
+	err = w.writeMetaToFile(tctx, db, createTableSQL, fileNameTable+".sql")
 	if err != nil {
 		return err
 	}
-	return writeMetaToFile(tctx, db, createViewSQL, w.extStorage, fileNameView+".sql", conf.CompressType)
+	return w.writeMetaToFile(tctx, db, createViewSQL, fileNameView+".sql")
 }
 
 // WriteSequenceMeta writes sequence meta to a file
@@ -181,7 +181,7 @@ func (w *Writer) WriteSequenceMeta(db, sequence, createSQL string) error {
 	if err != nil {
 		return err
 	}
-	return writeMetaToFile(tctx, db, createSQL, w.extStorage, fileName+".sql", conf.CompressType)
+	return w.writeMetaToFile(tctx, db, createSQL, fileName+".sql")
 }
 
 // WriteTableData writes table data to a file with retry
@@ -213,10 +213,14 @@ func (w *Writer) WriteTableData(meta TableMeta, ir TableDataIR, currentChunk int
 			return
 		}
 		if conf.SQL != "" {
-			meta, err = setTableMetaFromRows(ir.RawRows())
+			rows := ir.RawRows()
+			meta, err = setTableMetaFromRows(w.conf.ServerInfo.ServerType, rows)
 			if err != nil {
 				return err
 			}
+			if err = rows.Err(); err != nil {
+				return errors.Trace(err)
+			}
 		}
 		defer func() {
 			_ = ir.Close()
@@ -270,19 +274,17 @@ func (w *Writer) tryToWriteTableData(tctx *tcontext.Context, meta TableMeta, ir
 	return nil
 }
 
-func writeMetaToFile(tctx *tcontext.Context, target, metaSQL string, s storage.ExternalStorage, path string, compressType storage.CompressType) error {
-	fileWriter, tearDown, err := buildFileWriter(tctx, s, path, compressType)
+func (w *Writer) writeMetaToFile(tctx *tcontext.Context, target, metaSQL string, path string) error {
+	fileWriter, tearDown, err := buildFileWriter(tctx, w.extStorage, path, w.conf.CompressType)
 	if err != nil {
 		return errors.Trace(err)
 	}
 	defer tearDown(tctx)
 
 	return WriteMeta(tctx, &metaData{
-		target:  target,
-		metaSQL: metaSQL,
-		specCmts: []string{
-			"/*!40101 SET NAMES binary*/;",
-		},
+		target:   target,
+		metaSQL:  metaSQL,
+		specCmts: getSpecialComments(w.conf.ServerInfo.ServerType),
 	}, fileWriter)
 }
 

dumpling/export/writer_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/pingcap/tidb/br/pkg/version"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
 	"github.com/pingcap/tidb/util/promutil"
 	"github.com/stretchr/testify/require"
@@ -32,7 +33,7 @@ func TestWriteDatabaseMeta(t *testing.T) {
 
 	bytes, err := os.ReadFile(p)
 	require.NoError(t, err)
-	require.Equal(t, "/*!40101 SET NAMES binary*/;\nCREATE DATABASE `test`;\n", string(bytes))
+	require.Equal(t, "/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n/*!40101 SET NAMES binary*/;\nCREATE DATABASE `test`;\n", string(bytes))
 }
 
 func TestWritePolicyMeta(t *testing.T) {
@@ -51,7 +52,7 @@ func TestWritePolicyMeta(t *testing.T) {
 
 	bytes, err := os.ReadFile(p)
 	require.NoError(t, err)
-	require.Equal(t, "/*!40101 SET NAMES binary*/;\ncreate placement policy `y` followers=2;\n", string(bytes))
+	require.Equal(t, "/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n/*!40101 SET NAMES binary*/;\ncreate placement policy `y` followers=2;\n", string(bytes))
 }
 
 func TestWriteTableMeta(t *testing.T) {
@@ -69,7 +70,7 @@ func TestWriteTableMeta(t *testing.T) {
 	require.NoError(t, err)
 	bytes, err := os.ReadFile(p)
 	require.NoError(t, err)
-	require.Equal(t, "/*!40101 SET NAMES binary*/;\nCREATE TABLE t (a INT);\n", string(bytes))
+	require.Equal(t, "/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n/*!40101 SET NAMES binary*/;\nCREATE TABLE t (a INT);\n", string(bytes))
 }
 
 func TestWriteViewMeta(t *testing.T) {
@@ -79,7 +80,7 @@ func TestWriteViewMeta(t *testing.T) {
 
 	writer := createTestWriter(config, t)
 
-	specCmt := "/*!40101 SET NAMES binary*/;\n"
+	specCmt := "/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n/*!40101 SET NAMES binary*/;\n"
 	createTableSQL := "CREATE TABLE `v`(\n`a` int\n)ENGINE=MyISAM;\n"
 	createViewSQL := "DROP TABLE IF EXISTS `v`;\nDROP VIEW IF EXISTS `v`;\nSET @PREV_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT;\nSET @PREV_CHARACTER_SET_RESULTS=@@CHARACTER_SET_RESULTS;\nSET @PREV_COLLATION_CONNECTION=@@COLLATION_CONNECTION;\nSET character_set_client = utf8;\nSET character_set_results = utf8;\nSET collation_connection = utf8_general_ci;\nCREATE ALGORITHM=UNDEFINED DEFINER=`root`@`localhost` SQL SECURITY DEFINER VIEW `v` (`a`) AS SELECT `t`.`a` AS `a` FROM `test`.`t`;\nSET character_set_client = @PREV_CHARACTER_SET_CLIENT;\nSET character_set_results = @PREV_CHARACTER_SET_RESULTS;\nSET collation_connection = @PREV_COLLATION_CONNECTION;\n"
 	err := writer.WriteViewMeta("test", "v", createTableSQL, createViewSQL)
@@ -331,6 +332,7 @@ var mu sync.Mutex
 
 func createTestWriter(conf *Config, t *testing.T) *Writer {
 	t.Helper()
+	conf.ServerInfo.ServerType = version.ServerTypeMySQL
 
 	mu.Lock()
 	extStore, err := conf.createExternalStorage(context.Background())

dumpling/tests/file_size/run.sh

@@ -14,8 +14,8 @@ chars_20="1111_0000_1111_0000_"
 # insert 100 records, each occupies 20 bytes
 run_sql "insert into t values $(seq -s, 100 | sed 's/,*$//g' | sed "s/[0-9]*/('$chars_20')/g");"
 
-# dumping with file size = 311 bytes, actually 10 rows
-run_dumpling -F 311B
+# dumping with file size = 348 bytes, actually 10 rows
+run_dumpling -F 348B
 
 # the dumping result is expected to be:
 # 10 files for insertion(each conatins 10 records / 200 bytes)

dumpling/tests/naughty_strings/data/naughty_strings.escape.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `escape` VALUES
 ('''', '"'),

dumpling/tests/naughty_strings/data/naughty_strings.t.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `t` VALUES
 (''),

dumpling/tests/naughty_strings/expect/naughty_strings.escape.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `escape` VALUES
 ('\'','\"'),

dumpling/tests/naughty_strings/expect/naughty_strings.t.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `t` VALUES
 (''),

dumpling/tests/placement_policy/result/x-placement-policy-create.sql

@@ -1,2 +1,3 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 /*T![placement] CREATE PLACEMENT POLICY `x` PRIMARY_REGION="cn-east-1" REGIONS="cn-east-1,cn-east" */;

dumpling/tests/placement_policy/result/x1-placement-policy-create.sql

@@ -1,2 +1,3 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 /*T![placement] CREATE PLACEMENT POLICY `x1` FOLLOWERS=4 */;

dumpling/tests/primary_key/result/pk_case_0.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `pk_case_0` VALUES
 (0,10),

dumpling/tests/primary_key/result/pk_case_1.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `pk_case_1` VALUES
 (0,10),

dumpling/tests/primary_key/result/pk_case_2.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `pk_case_2` VALUES
 (0,10),

dumpling/tests/primary_key/result/pk_case_3.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `pk_case_3` VALUES
 (6,4,x'000000000101000000000000000000f03f000000000000f03f'),

dumpling/tests/quote/data/quote-database-schema-create-mysql57.sql

@@ -1,2 +1,3 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE DATABASE `quo``te/database` /*!40100 DEFAULT CHARACTER SET latin1 */;

dumpling/tests/quote/data/quote-database-schema-create.sql

@@ -1,2 +1,3 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE DATABASE `quo``te/database` /*!40100 DEFAULT CHARACTER SET latin1 */ /*!80016 DEFAULT ENCRYPTION='N' */;

dumpling/tests/quote/data/quote-database.quote-table-schema-mysql57.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE TABLE `quo``te/table` (
   `quo``te/col` int(11) NOT NULL,

dumpling/tests/quote/data/quote-database.quote-table-schema.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE TABLE `quo``te/table` (
   `quo``te/col` int NOT NULL,

dumpling/tests/quote/data/quote-database.quote-table.000000000-mysql57.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `quo``te/table` (`quo``te/col`,`a`) VALUES
 (0,10),

dumpling/tests/quote/data/quote-database.quote-table.000000000.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 INSERT INTO `quo``te/table` (`quo``te/col`,`a`) VALUES
 (0,10),

dumpling/tests/rows/run.sh

@@ -46,7 +46,7 @@ check_sync_diff $cur/conf/diff_config.toml
 
 # test dumpling with both rows and filesize
 rm -rf "$DUMPLING_OUTPUT_DIR"
-run_dumpling --rows 10 --filesize 100B --loglevel debug
+run_dumpling --rows 10 --filesize 140B --loglevel debug
 # the dumping result is expected to be:
 # 50 files for insertion
 file_num=$(find "$DUMPLING_OUTPUT_DIR" -maxdepth 1 -iname "$DB_NAME.$TABLE_NAME.*.sql" | wc -l)

dumpling/tests/sequences/data/sequences.s-schema-sequence-expect.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE SEQUENCE `s` start with 1 minvalue 1 maxvalue 9223372036854775806 increment by 1 cache 1000 nocycle ENGINE=InnoDB;
 SELECT SETVAL(`s`,1001);

dumpling/tests/views/data/views-schema-create.sql

@@ -1,2 +1,3 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE DATABASE `views` /*!40100 DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin */ /*!80016 DEFAULT ENCRYPTION='N' */;

dumpling/tests/views/data/views.v-schema-view.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 DROP TABLE IF EXISTS `v`;
 DROP VIEW IF EXISTS `v`;

dumpling/tests/views/data/views.v-schema.sql

@@ -1,3 +1,4 @@
+/*!40014 SET FOREIGN_KEY_CHECKS=0*/;
 /*!40101 SET NAMES binary*/;
 CREATE TABLE `v`(
 `a` int,