migrate/database/postgres/postgres.go

package postgres

import (
	"database/sql"
	"fmt"
	"hash/crc32"
	"io"
	"io/ioutil"
	nurl "net/url"

	"github.com/lib/pq"
	"github.com/mattes/migrate/database"
)

func init() {
	database.Register("postgres", &Postgres{})
}

var MigrationsTable = "schema_migrations"

var (
	ErrNilConfig      = fmt.Errorf("no config")
	ErrNoDatabaseName = fmt.Errorf("no database name")
	ErrNoSchema       = fmt.Errorf("no schema")
	ErrDatabaseDirty  = fmt.Errorf("database is dirty")
)

type Config struct {
	DatabaseName string
}

func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
	if config == nil {
		return nil, ErrNilConfig
	}

	query := `SELECT CURRENT_DATABASE()`
	var databaseName string
	if err := instance.QueryRow(query).Scan(&databaseName); err != nil {
		return nil, &database.Error{OrigErr: err, Query: []byte(query)}
	}

	if len(databaseName) == 0 {
		return nil, ErrNoDatabaseName
	}

	config.DatabaseName = databaseName

	px := &Postgres{
		db:     instance,
		config: config,
	}

	if err := px.ensureVersionTable(); err != nil {
		return nil, err
	}

	return px, nil
}

type Postgres struct {
	db       *sql.DB
	isLocked bool

	// Open and WithInstance need to garantuee that config is never nil
	config *Config
}

func (p *Postgres) Open(url string) (database.Driver, error) {
	purl, err := nurl.Parse(url)
	if err != nil {
		return nil, err
	}

	db, err := sql.Open("postgres", url)
	if err != nil {
		return nil, err
	}

	px, err := WithInstance(db, &Config{
		DatabaseName: purl.Path,
	})
	if err != nil {
		return nil, err
	}

	if err := db.Ping(); err != nil {
		return nil, err
	}

	return px, nil
}

func (p *Postgres) Close() error {
	return p.db.Close()
}

// https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
func (p *Postgres) Lock() error {
	if p.isLocked {
		return database.ErrLocked
	}

	aid, err := p.generateAdvisoryLockId()
	if err != nil {
		return err
	}

	// This will either obtain the lock immediately and return true,
	// or return false if the lock cannot be acquired immediately.
	query := `SELECT pg_try_advisory_lock($1)`
	var success bool
	if err := p.db.QueryRow(query, aid).Scan(&success); err != nil {
		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
	}

	if success {
		p.isLocked = true
		return nil
	}

	return database.ErrLocked
}

func (p *Postgres) Unlock() error {
	if !p.isLocked {
		return nil
	}

	aid, err := p.generateAdvisoryLockId()
	if err != nil {
		return err
	}

	query := `SELECT pg_advisory_unlock($1)`
	if _, err := p.db.Exec(query, aid); err != nil {
		return &database.Error{OrigErr: err, Query: []byte(query)}
	}
	p.isLocked = false
	return nil
}

func (p *Postgres) Run(version int, migration io.Reader) error {
	if dirty, err := p.isDirty(); err != nil {
		return err
	} else if dirty {
		return ErrDatabaseDirty
	}

	if migration == nil {
		// just apply version
		return p.saveVersion(version, false)
	}

	migr, err := ioutil.ReadAll(migration)
	if err != nil {
		return err
	}

	// set dirty flag and set version
	if err := p.saveVersion(version, true); err != nil {
		return err
	}

	// run migration
	query := string(migr[:])
	if _, err := p.db.Exec(query); err != nil {
		// TODO: cast to postgress error and get line number
		return database.Error{OrigErr: err, Err: "migration failed", Query: migr}
	}

	// remove dirty flag
	return p.saveVersion(version, false)
}

func (p *Postgres) saveVersion(version int, dirty bool) error {
	tx, err := p.db.Begin()
	if err != nil {
		return &database.Error{OrigErr: err, Err: "transaction start failed"}
	}

	query := `TRUNCATE "` + MigrationsTable + `"`
	if _, err := p.db.Exec(query); err != nil {
		tx.Rollback()
		return &database.Error{OrigErr: err, Query: []byte(query)}
	}

	if version >= 0 {
		query = `INSERT INTO "` + MigrationsTable + `" (version, dirty) VALUES ($1, $2)`
		if _, err := p.db.Exec(query, version, dirty); err != nil {
			tx.Rollback()
			return &database.Error{OrigErr: err, Query: []byte(query)}
		}
	}

	if err := tx.Commit(); err != nil {
		return &database.Error{OrigErr: err, Err: "transaction commit failed"}
	}

	return nil
}

func (p *Postgres) isDirty() (bool, error) {
	query := `SELECT dirty FROM "` + MigrationsTable + `" LIMIT 1`
	var dirty bool
	err := p.db.QueryRow(query).Scan(&dirty)
	switch {
	case err == sql.ErrNoRows:
		return false, nil

	case err != nil:
		if e, ok := err.(*pq.Error); ok {
			if e.Code.Name() == "undefined_table" {
				return false, nil
			}
		}
		return false, &database.Error{OrigErr: err, Query: []byte(query)}

	default:
		return dirty, nil
	}
}

func (p *Postgres) Version() (int, error) {
	query := `SELECT version FROM "` + MigrationsTable + `" LIMIT 1`
	var version uint64
	err := p.db.QueryRow(query).Scan(&version)
	switch {
	case err == sql.ErrNoRows:
		return database.NilVersion, nil

	case err != nil:
		if e, ok := err.(*pq.Error); ok {
			if e.Code.Name() == "undefined_table" {
				return database.NilVersion, nil
			}
		}
		return 0, &database.Error{OrigErr: err, Query: []byte(query)}

	default:
		return int(version), nil
	}
}

func (p *Postgres) Drop() error {
	// select all tables in current schema
	query := `SELECT table_name FROM information_schema.tables WHERE table_schema=(SELECT current_schema())`
	tables, err := p.db.Query(query)
	if err != nil {
		return &database.Error{OrigErr: err, Query: []byte(query)}
	}
	defer tables.Close()

	// delete one table after another
	tableNames := make([]string, 0)
	for tables.Next() {
		var tableName string
		if err := tables.Scan(&tableName); err != nil {
			return err
		}
		if len(tableName) > 0 {
			tableNames = append(tableNames, tableName)
		}
	}

	if len(tableNames) > 0 {
		// delete one by one ...
		for _, t := range tableNames {
			query = `DROP TABLE IF EXISTS ` + t + ` CASCADE`
			if _, err := p.db.Exec(query); err != nil {
				return &database.Error{OrigErr: err, Query: []byte(query)}
			}
		}
		if err := p.ensureVersionTable(); err != nil {
			return err
		}
	}

	return nil
}

func (p *Postgres) ensureVersionTable() error {
	// check if migration table exists
	var count int
	query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
	if err := p.db.QueryRow(query, MigrationsTable).Scan(&count); err != nil {
		return &database.Error{OrigErr: err, Query: []byte(query)}
	}
	if count == 1 {
		return nil
	}

	// if not, create the empty migration table
	query = `CREATE TABLE "` + MigrationsTable + `" (version bigint not null primary key, dirty boolean not null)`
	if _, err := p.db.Exec(query); err != nil {
		return &database.Error{OrigErr: err, Query: []byte(query)}
	}
	return nil
}

const AdvisoryLockIdSalt uint = 1486364155

// inspired by rails migrations, see https://goo.gl/8o9bCT
func (p *Postgres) generateAdvisoryLockId() (string, error) {
	sum := crc32.ChecksumIEEE([]byte(p.config.DatabaseName))
	sum = sum * uint32(AdvisoryLockIdSalt)
	return fmt.Sprintf("%v", sum), nil
}