-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* add v2 * add docs * add script to create database * restructure * add postgres * update readme * update readme * update package doc * remove SetGlobalEngine * update api * update readme
- Loading branch information
Showing
23 changed files
with
338 additions
and
720 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,122 +1,77 @@ | ||
| // Package dbcleaner helps cleaning up database's tables upon unit test. | ||
| // With the help of https://github.com/stretchr/testify/tree/master/suite, we can easily | ||
| // acquire the tables using in the test in SetupTest or SetupSuite, and cleanup all data | ||
| // using TearDownTest or TearDownSuite | ||
| package dbcleaner | ||
|
|
||
| import ( | ||
| "database/sql" | ||
| "errors" | ||
| "log" | ||
| "sync" | ||
|
|
||
| "github.com/khaiql/dbcleaner/helper" | ||
| "github.com/khaiql/dbcleaner/utils" | ||
| "github.com/khaiql/dbcleaner/engine" | ||
| ) | ||
|
|
||
| // DBCleaner instance of cleaner that can perform cleaning tables data | ||
| type DBCleaner struct { | ||
| db *sql.DB | ||
| driver string | ||
| } | ||
| // DbCleaner interface | ||
| type DbCleaner interface { | ||
| // SetEngine sets dbEngine, can be mysql, postgres... | ||
| SetEngine(dbEngine engine.Engine) | ||
|
|
||
| var ( | ||
| mutex sync.Mutex | ||
| registeredHelpers = make(map[string]helper.Helper) | ||
| // Acquire will lock tables passed in params so data in the table would not be deleted by other test cases | ||
| Acquire(tables ...string) | ||
|
|
||
| // ErrHelperNotFound return when calling an unregistered Helper | ||
| ErrHelperNotFound = errors.New("Helper has not been registered") | ||
| ) | ||
| // Clean calls Truncate the tables | ||
| Clean(tables ...string) error | ||
|
|
||
| // RegisterHelper register an Helper instance for a particular driver | ||
| func RegisterHelper(driverName string, helper helper.Helper) { | ||
| mutex.Lock() | ||
| defer mutex.Unlock() | ||
| registeredHelpers[driverName] = helper | ||
| // Close calls corresponding method on dbEngine to release connection to db | ||
| Close() error | ||
| } | ||
|
|
||
| // New returns a Cleaner instance for a particular driver using provided | ||
| // connectionString | ||
| func New(driver, connectionString string) (*DBCleaner, error) { | ||
| db, err := sql.Open(driver, connectionString) | ||
| // ErrTableNeverLockBefore is paniced if calling Release on table that havent' been acquired before | ||
| var ErrTableNeverLockBefore = errors.New("Table has never been locked before") | ||
|
|
||
| if err != nil { | ||
| return nil, err | ||
| // New returns a default Cleaner with Noop Engine. Call SetEngine to set an actual working engine | ||
| func New() DbCleaner { | ||
| return &cleanerImpl{ | ||
| locks: map[string]*sync.RWMutex{}, | ||
| dbEngine: &engine.NoOp{}, | ||
| } | ||
|
|
||
| return &DBCleaner{db, driver}, err | ||
| } | ||
|
|
||
| // FindHelper return a registered Helper using driver name | ||
| func FindHelper(driver string) (helper.Helper, error) { | ||
| if helper, ok := registeredHelpers[driver]; ok { | ||
| return helper, nil | ||
| } | ||
|
|
||
| return nil, ErrHelperNotFound | ||
| type cleanerImpl struct { | ||
| locks map[string]*sync.RWMutex | ||
| dbEngine engine.Engine | ||
| } | ||
|
|
||
| // Close closes connection to database | ||
| func (c *DBCleaner) Close() error { | ||
| return c.db.Close() | ||
| func (c *cleanerImpl) SetEngine(dbEngine engine.Engine) { | ||
| c.dbEngine = dbEngine | ||
| } | ||
|
|
||
| // TruncateTables truncates data of all tables | ||
| func (c *DBCleaner) TruncateTables() error { | ||
| return c.TruncateTablesExclude() | ||
| } | ||
| func (c *cleanerImpl) Acquire(tables ...string) { | ||
| for _, table := range tables { | ||
| if c.locks[table] == nil { | ||
| c.locks[table] = new(sync.RWMutex) | ||
| } | ||
|
|
||
| // TruncateTablesExclude truncates data of all tables but exclude some specify | ||
| // in the list | ||
| func (c *DBCleaner) TruncateTablesExclude(excludedTables ...string) error { | ||
| tables, err := c.getTables() | ||
| if err != nil { | ||
| return err | ||
| c.locks[table].RLock() | ||
| } | ||
|
|
||
| tables = utils.SubtractStringArray(tables, excludedTables) | ||
| return c.TruncateSelectedTables(tables...) | ||
| } | ||
|
|
||
| // TruncateSelectedTables truncates data of included tables | ||
| func (c *DBCleaner) TruncateSelectedTables(tables ...string) error { | ||
| helper, err := FindHelper(c.driver) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| var wg sync.WaitGroup | ||
| wg.Add(len(tables)) | ||
|
|
||
| func (c *cleanerImpl) Clean(tables ...string) error { | ||
| for _, table := range tables { | ||
| go func(tbl string) { | ||
| cmd := helper.TruncateTableCommand(tbl) | ||
| if _, err := c.db.Exec(cmd); err != nil { | ||
| log.Fatalf("Failed to truncate table %s. Error: %s", tbl, err.Error()) | ||
| } | ||
| wg.Done() | ||
| }(table) | ||
| } | ||
| if c.locks[table] != nil { | ||
| c.locks[table].RUnlock() | ||
| c.locks[table].Lock() | ||
| defer c.locks[table].Unlock() | ||
| } | ||
|
|
||
| wg.Wait() | ||
| if err := c.dbEngine.Truncate(table); err != nil { | ||
| return err | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func (c *DBCleaner) getTables() ([]string, error) { | ||
| tables := make([]string, 0) | ||
| helper, err := FindHelper(c.driver) | ||
| if err != nil { | ||
| return tables, err | ||
| } | ||
|
|
||
| rows, err := c.db.Query(helper.GetTablesQuery()) | ||
| if err != nil { | ||
| return tables, err | ||
| } | ||
| defer rows.Close() | ||
| for rows.Next() { | ||
| var value string | ||
| if err = rows.Scan(&value); err == nil { | ||
| tables = append(tables, value) | ||
| } | ||
| } | ||
|
|
||
| return tables, nil | ||
| func (c *cleanerImpl) Close() error { | ||
| return c.dbEngine.Close() | ||
| } |
Oops, something went wrong.