From b03e504a800e00a5fef7e2c60b4f0186ce71f095 Mon Sep 17 00:00:00 2001 From: Mike Arpaia Date: Tue, 3 Jan 2017 09:54:24 -0700 Subject: [PATCH] MySQL's datastore.Drop should drop all tables (#738) * MySQL's datastore.Drop should drop all tables * Use the database name as a parameter * tx.Exec instead of tx.MustExec --- cli/prepare.go | 8 +--- cli/serve.go | 5 +-- server/datastore/mysql/datastore.go | 62 +++++++++++++++++++---------- server/datastore/mysql_test.go | 4 +- 4 files changed, 46 insertions(+), 33 deletions(-) diff --git a/cli/prepare.go b/cli/prepare.go index ae8cef5d54..0fe8256645 100644 --- a/cli/prepare.go +++ b/cli/prepare.go @@ -33,9 +33,7 @@ To setup kolide infrastructure, use one of the available commands. Long: ``, Run: func(cmd *cobra.Command, args []string) { config := configManager.LoadConfig() - connString := mysql.GetMysqlConnectionString(config.Mysql) - - ds, err := mysql.New(connString, clock.C) + ds, err := mysql.New(config.Mysql, clock.C) if err != nil { initFatal(err, "creating db connection") } @@ -54,9 +52,7 @@ To setup kolide infrastructure, use one of the available commands. Long: ``, Run: func(cmd *cobra.Command, arg []string) { config := configManager.LoadConfig() - connString := mysql.GetMysqlConnectionString(config.Mysql) - - ds, err := mysql.New(connString, clock.C) + ds, err := mysql.New(config.Mysql, clock.C) if err != nil { initFatal(err, "creating db connection") } diff --git a/cli/serve.go b/cli/serve.go index 032fab68c9..6aa4e733ad 100644 --- a/cli/serve.go +++ b/cli/serve.go @@ -82,10 +82,7 @@ the way that the kolide server works. mailService = mail.NewDevService() } else { const defaultMaxAttempts = 15 - - connString := mysql.GetMysqlConnectionString(config.Mysql) - ds, err = mysql.New(connString, clock.C, mysql.Logger(logger)) - + ds, err = mysql.New(config.Mysql, clock.C, mysql.Logger(logger)) if err != nil { initFatal(err, "initializing datastore") } diff --git a/server/datastore/mysql/datastore.go b/server/datastore/mysql/datastore.go index c4fbfa9b9c..b5a7080543 100644 --- a/server/datastore/mysql/datastore.go +++ b/server/datastore/mysql/datastore.go @@ -24,11 +24,11 @@ type Datastore struct { db *sqlx.DB logger log.Logger clock clock.Clock + config config.MysqlConfig } // New creates an MySQL datastore. -func New(dbConnectString string, c clock.Clock, opts ...DBOption) (*Datastore, error) { - +func New(config config.MysqlConfig, c clock.Clock, opts ...DBOption) (*Datastore, error) { options := &dbOptions{ maxAttempts: defaultMaxAttempts, logger: log.NewNopLogger(), @@ -38,7 +38,7 @@ func New(dbConnectString string, c clock.Clock, opts ...DBOption) (*Datastore, e setOpt(options) } - db, err := sqlx.Open("mysql", dbConnectString) + db, err := sqlx.Open("mysql", generateMysqlConnectionString(config)) if err != nil { return nil, err } @@ -60,7 +60,12 @@ func New(dbConnectString string, c clock.Clock, opts ...DBOption) (*Datastore, e return nil, dbError } - ds := &Datastore{db, options.logger, c} + ds := &Datastore{ + db: db, + logger: options.logger, + clock: c, + config: config, + } return ds, nil @@ -93,24 +98,41 @@ func (d *Datastore) Initialize() error { // Drop removes database func (d *Datastore) Drop() error { - goose.SetDialect("mysql") + tables := []struct { + Name string `db:"TABLE_NAME"` + }{} - for { - version, err := goose.EnsureDBVersion(d.db.DB) - if err != nil { - return err - } + sql := ` + SELECT TABLE_NAME + FROM INFORMATION_SCHEMA.TABLES + WHERE TABLE_SCHEMA = ?; + ` - if version == 0 { - d.db.Exec("DROP TABLE IF EXISTS `goose_db_version`;") - return nil - } - - if err = goose.Run("down", d.db.DB, "."); err != nil { - return err - } + if err := d.db.Select(&tables, sql, d.config.Database); err != nil { + return err } + tx, err := d.db.Begin() + if err != nil { + return err + } + + _, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0") + if err != nil { + return tx.Rollback() + } + + for _, table := range tables { + _, err = tx.Exec(fmt.Sprintf("DROP TABLE %s;", table.Name)) + if err != nil { + return tx.Rollback() + } + } + _, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 1") + if err != nil { + return tx.Rollback() + } + return tx.Commit() } // HealthCheck returns an error if the MySQL backend is not healthy. @@ -155,9 +177,9 @@ func appendListOptionsToSQL(sql string, opts kolide.ListOptions) string { return sql } -// GetMysqlConnectionString returns a MySQL connection string using the +// generateMysqlConnectionString returns a MySQL connection string using the // provided configuration. -func GetMysqlConnectionString(conf config.MysqlConfig) string { +func generateMysqlConnectionString(conf config.MysqlConfig) string { return fmt.Sprintf( "%s:%s@(%s)/%s?charset=utf8&parseTime=true&loc=UTC", conf.Username, diff --git a/server/datastore/mysql_test.go b/server/datastore/mysql_test.go index 47722634b1..222ec0bbfd 100644 --- a/server/datastore/mysql_test.go +++ b/server/datastore/mysql_test.go @@ -25,9 +25,7 @@ func setupMySQL(t *testing.T) (ds *mysql.Datastore, teardown func()) { config.Address = h + ":3306" } - connString := mysql.GetMysqlConnectionString(config) - - ds, err := mysql.New(connString, clock.NewMockClock(), mysql.Logger(log.NewNopLogger()), mysql.LimitAttempts(1)) + ds, err := mysql.New(config, clock.NewMockClock(), mysql.Logger(log.NewNopLogger()), mysql.LimitAttempts(1)) require.Nil(t, err) teardown = func() { ds.Close()