@@ -20,6 +20,7 @@ import (
2020 "flag"
2121 "fmt"
2222 "os"
23+ "os/exec"
2324 "path/filepath"
2425 "regexp"
2526 "sort"
@@ -29,6 +30,7 @@ import (
2930 "time"
3031
3132 "github.com/defined2014/mysql"
33+ "github.com/google/uuid"
3234 "github.com/pingcap/errors"
3335 log "github.com/sirupsen/logrus"
3436)
4749 retryConnCount int
4850 collationDisable bool
4951 checkErr bool
52+ pathBR string
53+ pathDumpling string
5054)
5155
5256func init () {
@@ -63,6 +67,8 @@ func init() {
6367 flag .IntVar (& retryConnCount , "retry-connection-count" , 120 , "The max number to retry to connect to the database." )
6468 flag .BoolVar (& checkErr , "check-error" , false , "if --error ERR does not match, return error instead of just warn" )
6569 flag .BoolVar (& collationDisable , "collation-disable" , false , "run collation related-test with new-collation disabled" )
70+ flag .StringVar (& pathBR , "path-br" , "" , "Path of BR" )
71+ flag .StringVar (& pathDumpling , "path-dumpling" , "" , "Path of Dumpling" )
6672}
6773
6874const (
@@ -98,6 +104,11 @@ type ReplaceRegex struct {
98104 replace string
99105}
100106
107+ type SourceAndTarget struct {
108+ sourceTable string
109+ targetTable string
110+ }
111+
101112type tester struct {
102113 mdb * sql.DB
103114 name string
@@ -148,6 +159,12 @@ type tester struct {
148159
149160 // replace output result through --replace_regex /\.dll/.so/
150161 replaceRegex []* ReplaceRegex
162+
163+ // backup and restore context through --backup_and_restore $BACKUP_TABLE as $RESTORE_TABLE'
164+ backupAndRestore * SourceAndTarget
165+
166+ // dump and import context through --dump_and_import $SOURCE_TABLE as $TARGET_TABLE'
167+ dumpAndImport * SourceAndTarget
151168}
152169
153170func newTester (name string ) * tester {
@@ -352,6 +369,58 @@ func (t *tester) addSuccess(testSuite *XUnitTestSuite, startTime *time.Time, cnt
352369 })
353370}
354371
372+ func generateBRStatements (source , target string ) (string , string ) {
373+ // Generate a random UUID
374+ uuid := uuid .NewString ()
375+
376+ // Create the TMP_DIR path
377+ tmpDir := fmt .Sprintf ("/tmp/%s_%s" , source , uuid )
378+
379+ // Generate the SQL statements
380+ backupSQL := fmt .Sprintf ("BACKUP TABLE `%s` TO '%s'" , source , tmpDir )
381+ restoreSQL := fmt .Sprintf ("RESTORE TABLE `%s` FROM '%s'" , source , tmpDir )
382+
383+ return backupSQL , restoreSQL
384+ }
385+
386+ func (t * tester ) dumpTable (source string ) (string , error ) {
387+ log .Warnf ("Start dumping table: %s" , source )
388+ path := "/tmp/" + source + "_" + uuid .NewString ()
389+ cmdArgs := []string {
390+ fmt .Sprintf ("-h%s" , host ),
391+ fmt .Sprintf ("-P%s" , port ),
392+ fmt .Sprintf ("-u%s" , user ),
393+ fmt .Sprintf ("-T%s.%s" , t .name , source ),
394+ fmt .Sprintf ("-o%s" , path ),
395+ "--no-header" ,
396+ "--filetype" ,
397+ "csv" ,
398+ }
399+
400+ if passwd != "" {
401+ cmdArgs = append (cmdArgs , fmt .Sprintf ("-p%s" , passwd ))
402+ }
403+
404+ cmd := exec .Command (pathDumpling , cmdArgs ... )
405+
406+ output , err := cmd .CombinedOutput ()
407+ if err != nil {
408+ log .Warnf ("Failed executing commands: %s, output: %s)" ,
409+ cmd .String (), string (output ))
410+ return "" , err
411+ }
412+ log .Warnf ("Done executing commands: %s, output: %s)" ,
413+ cmd .String (), string (output ))
414+ return path , nil
415+ }
416+
417+ func (t * tester ) importTableStmt (path , target string ) string {
418+ return fmt .Sprintf (`
419+ IMPORT INTO %s
420+ FROM '%s/example.t.000000000.csv'
421+ ` , target , path )
422+ }
423+
355424func (t * tester ) Run () error {
356425 t .preProcess ()
357426 defer t .postProcess ()
@@ -523,6 +592,61 @@ func (t *tester) Run() error {
523592 return errors .Annotate (err , fmt .Sprintf ("Could not parse regex in --replace_regex: line: %d sql:%v" , q .Line , q .Query ))
524593 }
525594 t .replaceRegex = regex
595+ case Q_BACKUP_AND_RESTORE :
596+ t .backupAndRestore , err = parseSourceAndTarget (q .Query )
597+ if err != nil {
598+ return errors .Annotate (err , fmt .Sprintf ("Could not parse backup table and restore table name in --backup_and_restore, line: %d sql:%v" , q .Line , q .Query ))
599+ }
600+ backupStmt , restoreStmt := generateBRStatements (t .backupAndRestore .sourceTable , t .backupAndRestore .targetTable )
601+ log .WithFields (log.Fields {"stmt" : backupStmt , "line" : q .Line }).Warn ("Backup started" )
602+ if err := t .executeStmt (backupStmt ); err != nil {
603+ return err
604+ }
605+ log .WithFields (log.Fields {"stmt" : backupStmt , "line" : q .Line }).Warn ("Backup end" )
606+ tempTable := t .backupAndRestore .sourceTable + uuid .NewString ()
607+ renameStmt := fmt .Sprintf ("RENAME TABLE `%s` TO `%s`" , t .backupAndRestore .sourceTable , tempTable )
608+ if err := t .executeStmt (renameStmt ); err != nil {
609+ return err
610+ }
611+ dupTableStmt := fmt .Sprintf ("CREATE TABLE `%s` LIKE `%s`" , t .backupAndRestore .sourceTable , tempTable )
612+ if err := t .executeStmt (dupTableStmt ); err != nil {
613+ return err
614+ }
615+ log .WithFields (log.Fields {"stmt" : restoreStmt , "line" : q .Line }).Warn ("Restore start" )
616+ if err := t .executeStmt (restoreStmt ); err != nil {
617+ return err
618+ }
619+ log .WithFields (log.Fields {"stmt" : restoreStmt , "line" : q .Line }).Warn ("Restore end" )
620+ renameStmt = fmt .Sprintf ("RENAME TABLE `%s` TO `%s`" , t .backupAndRestore .sourceTable , t .backupAndRestore .targetTable )
621+ if err := t .executeStmt (renameStmt ); err != nil {
622+ return err
623+ }
624+ renameStmt = fmt .Sprintf ("RENAME TABLE `%s` TO `%s`" , tempTable , t .backupAndRestore .sourceTable )
625+ if err := t .executeStmt (renameStmt ); err != nil {
626+ return err
627+ }
628+ case Q_DUMP_AND_IMPORT :
629+ t .dumpAndImport , err = parseSourceAndTarget (q .Query )
630+ if err != nil {
631+ return err
632+ }
633+ path , err := t .dumpTable (t .dumpAndImport .sourceTable )
634+ if err != nil {
635+ return err
636+ }
637+
638+ dupTableStmt := fmt .Sprintf ("CREATE TABLE `%s` LIKE `%s`" , t .dumpAndImport .targetTable , t .backupAndRestore .sourceTable )
639+ if err := t .executeStmt (dupTableStmt ); err != nil {
640+ return err
641+ }
642+
643+ importStmt := t .importTableStmt (path , t .dumpAndImport .targetTable )
644+ log .WithFields (log.Fields {"stmt" : importStmt , "line" : q .Line }).Warn ("Import start" )
645+ if err = t .executeStmt (importStmt ); err != nil {
646+ return err
647+ }
648+ log .WithFields (log.Fields {"stmt" : importStmt , "line" : q .Line }).Warn ("Restore end" )
649+
526650 default :
527651 log .WithFields (log.Fields {"command" : q .firstWord , "arguments" : q .Query , "line" : q .Line }).Warn ("command not implemented" )
528652 }
0 commit comments