From df7824d944c5ffbe28b913384d13c68583a26c6a Mon Sep 17 00:00:00 2001 From: Reid Beels Date: Wed, 17 Aug 2022 17:36:41 -0700 Subject: [PATCH] Add --data-only to copy data without structure This was prompted as a workaround for the issue presented in #139. By allowing data-only transfers to be run, we can use multiple config files in order to extract multiple subsets from the same table. An initial run, without `--data-only`, copies the DB structure and anonmized users while a second run using `--data-only` and a different config file copies unmodified versions of admin users to the same table. It's likely useful in many other caeses as well. --- cmd/steal.go | 4 +++- features/mysql_test.go | 2 +- features/postgres_test.go | 2 +- pkg/dumper/dumper.go | 2 +- pkg/dumper/engine/engine.go | 8 +++++--- pkg/dumper/query/dumper.go | 16 +++++++++------- 6 files changed, 20 insertions(+), 14 deletions(-) diff --git a/cmd/steal.go b/cmd/steal.go index 48a1610..1418fde 100644 --- a/cmd/steal.go +++ b/cmd/steal.go @@ -33,6 +33,7 @@ type ( concurrency int readOpts connOpts writeOpts connOpts + dataOnly bool } connOpts struct { timeout time.Duration @@ -76,6 +77,7 @@ func NewStealCmd() *cobra.Command { persistentFlags.DurationVar(&opts.writeOpts.maxConnLifetime, "write-conn-lifetime", 0, "Sets the maximum amount of time a connection may be reused on the write database") persistentFlags.IntVar(&opts.writeOpts.maxConns, "write-max-conns", 5, "Sets the maximum number of open connections to the write database") persistentFlags.IntVar(&opts.writeOpts.maxIdleConns, "write-max-idle-conns", 0, "Sets the maximum number of connections in the idle connection pool for the write database") + persistentFlags.BoolVar(&opts.dataOnly, "data-only", false, "Only steal data; requires that the target database structure already exists") return cmd } @@ -122,7 +124,7 @@ func RunSteal(opts *StealOptions) (err error) { defer close(done) start := time.Now() - if err := target.Dump(done, opts.cfgTables, opts.concurrency); err != nil { + if err := target.Dump(done, opts.cfgTables, opts.concurrency, opts.dataOnly); err != nil { return fmt.Errorf("error while dumping: %w", err) } diff --git a/features/mysql_test.go b/features/mysql_test.go index ce76bc0..269d4ee 100644 --- a/features/mysql_test.go +++ b/features/mysql_test.go @@ -53,7 +53,7 @@ func (s *MysqlTestSuite) TestExample() { done := make(chan struct{}) defer close(done) - s.Require().NoError(dmp.Dump(done, config.Tables{}, 4), "Failed to dump") + s.Require().NoError(dmp.Dump(done, config.Tables{}, 4, false), "Failed to dump") <-done diff --git a/features/postgres_test.go b/features/postgres_test.go index 347f526..84fd427 100644 --- a/features/postgres_test.go +++ b/features/postgres_test.go @@ -54,7 +54,7 @@ func (s *PostgresTestSuite) TestExample() { done := make(chan struct{}) defer close(done) - s.Require().NoError(dmp.Dump(done, config.Tables{}, 4), "Failed to dump") + s.Require().NoError(dmp.Dump(done, config.Tables{}, 4, false), "Failed to dump") <-done diff --git a/pkg/dumper/dumper.go b/pkg/dumper/dumper.go index 8b70b01..97d57ed 100644 --- a/pkg/dumper/dumper.go +++ b/pkg/dumper/dumper.go @@ -22,7 +22,7 @@ type ( // A Dumper writes a database's structure to the provided stream. Dumper interface { // Dump executes the dump process. - Dump(chan<- struct{}, config.Tables, int) error + Dump(chan<- struct{}, config.Tables, int, bool) error // Close closes the dumper resources and releases them. Close() error } diff --git a/pkg/dumper/engine/engine.go b/pkg/dumper/engine/engine.go index b6931bd..5bf6bb4 100644 --- a/pkg/dumper/engine/engine.go +++ b/pkg/dumper/engine/engine.go @@ -47,9 +47,11 @@ func New(rdr reader.Reader, dumper Dumper) dumper.Dumper { } // Dump executes the dump process. -func (e *Engine) Dump(done chan<- struct{}, cfgTables config.Tables, concurrency int) error { - if err := e.readAndDumpStructure(); err != nil { - return err +func (e *Engine) Dump(done chan<- struct{}, cfgTables config.Tables, concurrency int, dataOnly bool) error { + if !dataOnly { + if err := e.readAndDumpStructure(); err != nil { + return err + } } return e.readAndDumpTables(done, cfgTables, concurrency) diff --git a/pkg/dumper/query/dumper.go b/pkg/dumper/query/dumper.go index 8bfe673..d921e61 100644 --- a/pkg/dumper/query/dumper.go +++ b/pkg/dumper/query/dumper.go @@ -33,18 +33,20 @@ func NewDumper(output io.Writer, rdr reader.Reader) dumper.Dumper { } // Dump executes the dump stream process. -func (d *textDumper) Dump(done chan<- struct{}, cfgTables config.Tables, concurrency int) error { +func (d *textDumper) Dump(done chan<- struct{}, cfgTables config.Tables, concurrency int, dataOnly bool) error { tables, err := d.reader.GetTables() if err != nil { return fmt.Errorf("failed to get tables: %w", err) } - structure, err := d.reader.GetStructure() - if err != nil { - return fmt.Errorf("could not get database structure: %w", err) - } - if _, err := io.WriteString(d.output, structure); err != nil { - return fmt.Errorf("could not write structure to output: %w", err) + if !dataOnly { + structure, err := d.reader.GetStructure() + if err != nil { + return fmt.Errorf("could not get database structure: %w", err) + } + if _, err := io.WriteString(d.output, structure); err != nil { + return fmt.Errorf("could not write structure to output: %w", err) + } } var wg sync.WaitGroup