support complete insert (pingcap#135)

tisonkun · Aug 19, 2020 · c1da150 · c1da150
1 parent 94b4741
commit c1da150
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 11 deletions.
diff --git a/dumpling/cmd/dumpling/main.go b/dumpling/cmd/dumpling/main.go
@@ -67,6 +67,7 @@ var (
 	csvSeparator  string
 	csvDelimiter  string
 
+	completeInsert       bool
 	dumpEmptyDatabase    bool
 	escapeBackslash      bool
 	tidbMemQuotaQuery    uint64
@@ -122,6 +123,7 @@ func main() {
 	pflag.StringVar(&csvSeparator, "csv-separator", ",", "The separator for csv files, default ','")
 	pflag.StringVar(&csvDelimiter, "csv-delimiter", "\"", "The delimiter for values in csv files, default '\"'")
 	pflag.StringVar(&outputFilenameFormat, "output-filename-template", "", "The output filename template (without file extension), default '{{.DB}}.{{.Table}}.{{.Index}}'")
+	pflag.BoolVar(&completeInsert, "complete-insert", false, "Use complete INSERT statements that include column names")
 
 	printVersion := pflag.BoolP("version", "V", false, "Print Dumpling version")
 
@@ -208,6 +210,7 @@ func main() {
 	conf.CsvSeparator = csvSeparator
 	conf.CsvDelimiter = csvDelimiter
 	conf.OutputFileTemplate = tmpl
+	conf.CompleteInsert = completeInsert
 
 	err = export.Dump(context.Background(), conf)
 	if err != nil {

diff --git a/dumpling/v4/export/config.go b/dumpling/v4/export/config.go
@@ -54,6 +54,7 @@ type Config struct {
 	Rows               uint64
 	Where              string
 	FileType           string
+	CompleteInsert     bool
 	EscapeBackslash    bool
 	DumpEmptyDatabase  bool
 	OutputFileTemplate *template.Template

diff --git a/dumpling/v4/export/ir_impl.go b/dumpling/v4/export/ir_impl.go
@@ -315,7 +315,7 @@ func splitTableDataIntoChunks(
 	estimatedStep := (max-min)/estimatedChunks + 1
 	cutoff := min
 
-	selectedField, err := buildSelectField(db, dbName, tableName)
+	selectedField, err := buildSelectField(db, dbName, tableName, conf.CompleteInsert)
 	if err != nil {
 		errCh <- withStack(err)
 		return

diff --git a/dumpling/v4/export/sql.go b/dumpling/v4/export/sql.go
@@ -133,7 +133,7 @@ func SelectVersion(db *sql.DB) (string, error) {
 }
 
 func SelectAllFromTable(conf *Config, db *sql.Conn, database, table string) (TableDataIR, error) {
-	selectedField, err := buildSelectField(db, database, table)
+	selectedField, err := buildSelectField(db, database, table, conf.CompleteInsert)
 	if err != nil {
 		return nil, err
 	}
@@ -455,7 +455,7 @@ func createConnWithConsistency(ctx context.Context, db *sql.DB) (*sql.Conn, erro
 	return conn, nil
 }
 
-func buildSelectField(db *sql.Conn, dbName, tableName string) (string, error) {
+func buildSelectField(db *sql.Conn, dbName, tableName string, completeInsert bool) (string, error) {
 	query := `SELECT COLUMN_NAME,EXTRA FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA=? AND TABLE_NAME=?;`
 	rows, err := db.QueryContext(context.Background(), query, dbName, tableName)
 	if err != nil {
@@ -479,7 +479,7 @@ func buildSelectField(db *sql.Conn, dbName, tableName string) (string, error) {
 		}
 		availableFields = append(availableFields, wrapBackTicks(escapeString(fieldName)))
 	}
-	if hasGenerateColumn {
+	if completeInsert || hasGenerateColumn {
 		return strings.Join(availableFields, ","), nil
 	}
 	return "*", nil

diff --git a/dumpling/v4/export/sql_test.go b/dumpling/v4/export/sql_test.go
@@ -75,7 +75,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 		WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-	selectedField, err := buildSelectField(conn, "test", "t")
+	selectedField, err := buildSelectField(conn, "test", "t", false)
 	c.Assert(err, IsNil)
 	q := buildSelectQuery("test", "t", selectedField, "", orderByClause)
 	c.Assert(q, Equals, "SELECT * FROM `test`.`t` ORDER BY _tidb_rowid")
@@ -91,7 +91,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 		WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-	selectedField, err = buildSelectField(conn, "test", "t")
+	selectedField, err = buildSelectField(conn, "test", "t", false)
 	c.Assert(err, IsNil)
 	q = buildSelectQuery("test", "t", selectedField, "", orderByClause)
 	c.Assert(q, Equals, "SELECT * FROM `test`.`t`")
@@ -114,7 +114,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 			WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 			WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-		selectedField, err = buildSelectField(conn, "test", "t")
+		selectedField, err = buildSelectField(conn, "test", "t", false)
 		c.Assert(err, IsNil)
 		q = buildSelectQuery("test", "t", selectedField, "", orderByClause)
 		c.Assert(q, Equals, "SELECT * FROM `test`.`t` ORDER BY `id`", cmt)
@@ -138,7 +138,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 			WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 			WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-		selectedField, err = buildSelectField(conn, "test", "t")
+		selectedField, err = buildSelectField(conn, "test", "t", false)
 		c.Assert(err, IsNil)
 		q := buildSelectQuery("test", "t", selectedField, "", orderByClause)
 		c.Assert(q, Equals, "SELECT * FROM `test`.`t`", cmt)
@@ -157,7 +157,7 @@ func (s *testDumpSuite) TestBuildSelectAllQuery(c *C) {
 			WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 			WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-		selectedField, err := buildSelectField(conn, "test", "t")
+		selectedField, err := buildSelectField(conn, "test", "t", false)
 		c.Assert(err, IsNil)
 		q := buildSelectQuery("test", "t", selectedField, "", "")
 		c.Assert(q, Equals, "SELECT * FROM `test`.`t`", cmt)
@@ -177,18 +177,29 @@ func (s *testDumpSuite) TestBuildSelectField(c *C) {
 		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 		WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", ""))
 
-	selectedField, err := buildSelectField(conn, "test", "t")
+	selectedField, err := buildSelectField(conn, "test", "t", false)
 	c.Assert(selectedField, Equals, "*")
 	c.Assert(err, IsNil)
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 
+	// user assigns completeInsert
+	mock.ExpectQuery("SELECT COLUMN_NAME").
+		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
+		WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).AddRow("id", "").
+			AddRow("name", "").AddRow("quo`te", ""))
+
+	selectedField, err = buildSelectField(conn, "test", "t", true)
+	c.Assert(selectedField, Equals, "`id`,`name`,`quo``te`")
+	c.Assert(err, IsNil)
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+
 	// found generate columns, rest columns is `id`,`name`
 	mock.ExpectQuery("SELECT COLUMN_NAME").
 		WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).
 		WillReturnRows(sqlmock.NewRows([]string{"column_name", "extra"}).
 			AddRow("id", "").AddRow("name", "").AddRow("quo`te", "").AddRow("generated", "VIRTUAL GENERATED"))
 
-	selectedField, err = buildSelectField(conn, "test", "t")
+	selectedField, err = buildSelectField(conn, "test", "t", false)
 	c.Assert(selectedField, Equals, "`id`,`name`,`quo``te`")
 	c.Assert(err, IsNil)
 	c.Assert(mock.ExpectationsWereMet(), IsNil)