githaven-fork/vendor/gopkg.in/testfixtures.v2/generate.go

package testfixtures

import (
	"database/sql"
	"fmt"
	"os"
	"path"
	"unicode/utf8"

	"gopkg.in/yaml.v2"
)

// TableInfo is settings for generating a fixture for table.
type TableInfo struct {
	Name  string // Table name
	Where string // A condition for extracting records. If this value is empty, extracts all records.
}

func (ti *TableInfo) whereClause() string {
	if ti.Where == "" {
		return ""
	}
	return fmt.Sprintf(" WHERE %s", ti.Where)
}

// GenerateFixtures generates fixtures for the current contents of a database, and saves
// them to the specified directory
func GenerateFixtures(db *sql.DB, helper Helper, dir string) error {
	tables, err := helper.tableNames(db)
	if err != nil {
		return err
	}
	for _, table := range tables {
		filename := path.Join(dir, table+".yml")
		if err := generateFixturesForTable(db, helper, &TableInfo{Name: table}, filename); err != nil {
			return err
		}
	}
	return nil
}

// GenerateFixturesForTables generates fixtures for the current contents of specified tables in a database, and saves
// them to the specified directory
func GenerateFixturesForTables(db *sql.DB, tables []*TableInfo, helper Helper, dir string) error {
	for _, table := range tables {
		filename := path.Join(dir, table.Name+".yml")
		if err := generateFixturesForTable(db, helper, table, filename); err != nil {
			return err
		}
	}
	return nil
}

func generateFixturesForTable(db *sql.DB, h Helper, table *TableInfo, filename string) error {
	query := fmt.Sprintf("SELECT * FROM %s%s", h.quoteKeyword(table.Name), table.whereClause())
	rows, err := db.Query(query)
	if err != nil {
		return err
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	fixtures := make([]interface{}, 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(map[string]interface{}, len(entries))
		for i, column := range columns {
			entryMap[column] = convertValue(entries[i])
		}
		fixtures = append(fixtures, entryMap)
	}
	if err = rows.Err(); err != nil {
		return err
	}

	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	marshaled, err := yaml.Marshal(fixtures)
	if err != nil {
		return err
	}
	_, err = f.Write(marshaled)
	return err
}

func convertValue(value interface{}) interface{} {
	switch v := value.(type) {
	case []byte:
		if utf8.Valid(v) {
			return string(v)
		}
	}
	return value
}