123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165 |
- package testfixtures
-
- import (
- "database/sql"
- "fmt"
- "os"
- "path/filepath"
- "unicode/utf8"
-
- "gopkg.in/yaml.v2"
- )
-
- // Dumper is resposible for dumping fixtures from the database into a
- // directory.
- type Dumper struct {
- db *sql.DB
- helper helper
- dir string
-
- tables []string
- }
-
- // NewDumper creates a new dumper with the given options.
- //
- // The "DumpDatabase", "DumpDialect" and "DumpDirectory" options are required.
- func NewDumper(options ...func(*Dumper) error) (*Dumper, error) {
- d := &Dumper{}
-
- for _, option := range options {
- if err := option(d); err != nil {
- return nil, err
- }
- }
-
- return d, nil
- }
-
- // DumpDatabase sets the database to be dumped.
- func DumpDatabase(db *sql.DB) func(*Dumper) error {
- return func(d *Dumper) error {
- d.db = db
- return nil
- }
- }
-
- // DumpDialect informs Loader about which database dialect you're using.
- //
- // Possible options are "postgresql", "timescaledb", "mysql", "mariadb",
- // "sqlite" and "sqlserver".
- func DumpDialect(dialect string) func(*Dumper) error {
- return func(d *Dumper) error {
- h, err := helperForDialect(dialect)
- if err != nil {
- return err
- }
- d.helper = h
- return nil
- }
- }
-
- // DumpDirectory sets the directory where the fixtures files will be created.
- func DumpDirectory(dir string) func(*Dumper) error {
- return func(d *Dumper) error {
- d.dir = dir
- return nil
- }
- }
-
- // DumpTables allows you to choose which tables you want to dump.
- //
- // If not informed, Dumper will dump all tables by default.
- func DumpTables(tables ...string) func(*Dumper) error {
- return func(d *Dumper) error {
- d.tables = tables
- return nil
- }
- }
-
- // Dump dumps the databases as YAML fixtures.
- func (d *Dumper) Dump() error {
- tables := d.tables
- if len(tables) == 0 {
- var err error
- tables, err = d.helper.tableNames(d.db)
- if err != nil {
- return err
- }
- }
-
- for _, table := range tables {
- if err := d.dumpTable(table); err != nil {
- return err
- }
- }
- return nil
- }
-
- func (d *Dumper) dumpTable(table string) error {
- query := fmt.Sprintf("SELECT * FROM %s", d.helper.quoteKeyword(table))
-
- stmt, err := d.db.Prepare(query)
- if err != nil {
- return err
- }
- defer stmt.Close()
-
- rows, err := stmt.Query()
- if err != nil {
- return err
- }
- defer rows.Close()
-
- columns, err := rows.Columns()
- if err != nil {
- return err
- }
-
- fixtures := make([]yaml.MapSlice, 0, 10)
- for rows.Next() {
- entries := make([]interface{}, len(columns))
- entryPtrs := make([]interface{}, len(entries))
- for i := range entries {
- entryPtrs[i] = &entries[i]
- }
- if err := rows.Scan(entryPtrs...); err != nil {
- return err
- }
-
- entryMap := make([]yaml.MapItem, len(entries))
- for i, column := range columns {
- entryMap[i] = yaml.MapItem{
- Key: column,
- Value: convertValue(entries[i]),
- }
- }
- fixtures = append(fixtures, entryMap)
- }
- if err = rows.Err(); err != nil {
- return err
- }
-
- filePath := filepath.Join(d.dir, table+".yml")
- f, err := os.Create(filePath)
- if err != nil {
- return err
- }
- defer f.Close()
-
- data, err := yaml.Marshal(fixtures)
- if err != nil {
- return err
- }
- _, err = f.Write(data)
- return err
- }
-
- func convertValue(value interface{}) interface{} {
- switch v := value.(type) {
- case []byte:
- if utf8.Valid(v) {
- return string(v)
- }
- }
- return value
- }
|