diff --git a/internal/dumper/dumper_test.go b/internal/dumper/dumper_test.go index 41b03c77..bd2f9f6e 100644 --- a/internal/dumper/dumper_test.go +++ b/internal/dumper/dumper_test.go @@ -5,6 +5,7 @@ import ( "encoding/csv" "encoding/json" "os" + "path/filepath" "strings" "testing" @@ -1421,6 +1422,110 @@ func TestDumperInvertMatch(t *testing.T) { c.Assert(want_test2, qt.IsTrue) } +// TestDumperKeepsRowsAtChunkBoundary checks that rows that are not yet part +// of an INSERT statement when a chunk file is flushed are still written, even +// when the table ends right after that flush. +func TestDumperKeepsRowsAtChunkBoundary(t *testing.T) { + c := qt.New(t) + + log := xlog.NewStdLog(xlog.Level(xlog.INFO)) + fakedbs := driver.NewTestHandler(log) + server, err := driver.MockMysqlServer(log, fakedbs) + c.Assert(err, qt.IsNil) + c.Cleanup(func() { server.Close() }) + + payload := strings.Repeat("x", 1000) + rowLen := len(`(1,"` + payload + `")`) + stmtSize := 10 * rowLen + // The last row is the first one to reach the 1 MiB chunk size, and it + // does not complete an INSERT statement. + numRows := (1024*1024 + rowLen - 1) / rowLen + c.Assert(numRows%10, qt.Not(qt.Equals), 0) + + selectResult := &sqltypes.Result{ + Fields: []*querypb.Field{ + {Name: "id", Type: querypb.Type_INT32}, + {Name: "name", Type: querypb.Type_VARCHAR}, + }, + } + for i := 0; i < numRows; i++ { + selectResult.Rows = append(selectResult.Rows, []sqltypes.Value{ + sqltypes.MakeTrusted(querypb.Type_INT32, []byte("1")), + sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte(payload)), + }) + } + + schemaResult := &sqltypes.Result{ + Fields: []*querypb.Field{ + {Name: "Table", Type: querypb.Type_VARCHAR}, + {Name: "Create Table", Type: querypb.Type_VARCHAR}, + }, + Rows: [][]sqltypes.Value{ + { + sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte("t1")), + sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte("CREATE TABLE `t1` (`id` int, `name` text) ENGINE=InnoDB")), + }, + }, + } + + tablesResult := &sqltypes.Result{ + Fields: []*querypb.Field{{Name: "Tables_in_test", Type: querypb.Type_VARCHAR}}, + Rows: [][]sqltypes.Value{ + {sqltypes.MakeTrusted(querypb.Type_VARCHAR, []byte("t1"))}, + }, + } + + fieldsResult := &sqltypes.Result{ + Fields: []*querypb.Field{ + {Name: "Field", Type: querypb.Type_VARCHAR}, + {Name: "Type", Type: querypb.Type_VARCHAR}, + {Name: "Null", Type: querypb.Type_VARCHAR}, + {Name: "Key", Type: querypb.Type_VARCHAR}, + {Name: "Default", Type: querypb.Type_VARCHAR}, + {Name: "Extra", Type: querypb.Type_VARCHAR}, + }, + Rows: [][]sqltypes.Value{ + testRow("id", ""), + testRow("name", ""), + }, + } + + fakedbs.AddQueryPattern("use .*", &sqltypes.Result{}) + fakedbs.AddQueryPattern("show create table .*", schemaResult) + fakedbs.AddQueryPattern("show tables from .*", tablesResult) + fakedbs.AddQueryPattern("select table_name \n\t\t\t from information_schema.tables \n\t\t\t where table_schema like 'test' \n\t\t\t and table_type = 'view'\n\t\t\t", &sqltypes.Result{}) + fakedbs.AddQueryPattern("show fields from .*", fieldsResult) + fakedbs.AddQueryPattern("select .* from `test`\\..* .*", selectResult) + fakedbs.AddQueryPattern("set .*", &sqltypes.Result{}) + + cfg := &Config{ + Database: "test", + Outdir: c.TempDir(), + User: "mock", + Password: "mock", + Address: server.Addr(), + ChunksizeInMB: 1, + Threads: 1, + StmtSize: stmtSize, + IntervalMs: 500, + } + + d, err := NewDumper(cfg) + c.Assert(err, qt.IsNil) + c.Assert(d.Run(context.Background()), qt.IsNil) + + files, err := filepath.Glob(filepath.Join(cfg.Outdir, "test.t1.*.sql")) + c.Assert(err, qt.IsNil) + + dumpedRows := 0 + for _, file := range files { + dat, err := os.ReadFile(file) + c.Assert(err, qt.IsNil) + dumpedRows += strings.Count(string(dat), `(1,"`+payload+`")`) + } + c.Assert(dumpedRows, qt.Equals, numRows) +} + func TestWriteFile(t *testing.T) { c := qt.New(t) diff --git a/internal/dumper/sql_writer.go b/internal/dumper/sql_writer.go index 5c8cc1f3..f180c60a 100644 --- a/internal/dumper/sql_writer.go +++ b/internal/dumper/sql_writer.go @@ -58,10 +58,7 @@ func (w *sqlWriter) WriteRow(row []sqltypes.Value) (int, error) { w.chunkbytes += rowBytes if w.stmtsize >= w.cfg.StmtSize { - insertone := fmt.Sprintf("INSERT INTO %s(%s) VALUES\n%s", w.table, strings.Join(w.fields, ","), strings.Join(w.rows, ",\n")) - w.inserts = append(w.inserts, insertone) - w.rows = w.rows[:0] - w.stmtsize = 0 + w.appendInsert() } return rowBytes, nil @@ -71,7 +68,21 @@ func (w *sqlWriter) ShouldFlush() bool { return (w.chunkbytes / 1024 / 1024) >= w.cfg.ChunksizeInMB } +// appendInsert turns the buffered rows into an INSERT statement. +func (w *sqlWriter) appendInsert() { + insertone := fmt.Sprintf("INSERT INTO %s(%s) VALUES\n%s", w.table, strings.Join(w.fields, ","), strings.Join(w.rows, ",\n")) + w.inserts = append(w.inserts, insertone) + w.rows = w.rows[:0] + w.stmtsize = 0 +} + func (w *sqlWriter) Flush(outdir, database, table string, fileNo int) error { + // Write the rows that are not part of an INSERT statement yet, so they + // are not lost when the table ends right after this flush. + if len(w.rows) > 0 { + w.appendInsert() + } + query := strings.Join(w.inserts, ";\n") + ";\n" file, err := dumpOutputPath(outdir, database, table, fmt.Sprintf(".%05d.sql", fileNo)) if err != nil { @@ -89,10 +100,6 @@ func (w *sqlWriter) Flush(outdir, database, table string, fileNo int) error { func (w *sqlWriter) Close(outdir, database, table string, fileNo int) error { if w.chunkbytes > 0 { - if len(w.rows) > 0 { - insertone := fmt.Sprintf("INSERT INTO %s(%s) VALUES\n%s", w.table, strings.Join(w.fields, ","), strings.Join(w.rows, ",\n")) - w.inserts = append(w.inserts, insertone) - } return w.Flush(outdir, database, table, fileNo) } return nil