feat(setup): Add migration tables and functions to set them up

.travis.yml

@@ -1,6 +1,8 @@
 dist: trusty
 sudo: false
 language: go
+services:
+  - postgresql
 go: 1.10.x
 install:
   - make setup

Makefile

@@ -1,8 +1,15 @@
 DIRS ?= $(shell find . -name '*.go' | grep --invert-match 'vendor' | xargs -n 1 dirname | sort --unique)
 
+TFLAGS ?=
+
 COVERAGE_PROFILE ?= coverage.out
 HTML_OUTPUT      ?= coverage.html
 
+PSQL := $(shell command -v psql 2> /dev/null)
+
+TEST_DATABASE_USER ?= go_pg_migrations_user
+TEST_DATABASE_NAME ?= go_pg_migrations
+
 default: install
 
 .PHONY: clean
@@ -36,8 +43,16 @@ setup:
 	@echo "--> Setting up"
 	go get -u -v github.com/alecthomas/gometalinter github.com/golang/dep/cmd/dep
 	gometalinter --install
+ifdef PSQL
+	dropuser --if-exists $(TEST_DATABASE_USER)
+	dropdb --if-exists $(TEST_DATABASE_NAME)
+	createuser --createdb $(TEST_DATABASE_USER)
+	createdb -U $(TEST_DATABASE_USER) $(TEST_DATABASE_NAME)
+else
+	$(error Postgres should be installed)
+endif
 
 .PHONY: test
 test:
 	@echo "---> Testing"
-	go test ./... -coverprofile $(COVERAGE_PROFILE)
+	TEST_DATABASE_USER=$(TEST_DATABASE_USER) TEST_DATABASE_NAME=$(TEST_DATABASE_NAME) go test ./... -coverprofile $(COVERAGE_PROFILE) $(TFLAGS)

migrations.go

@@ -1,17 +1,50 @@
+// Package migrations provides a robust mechanism for registering, creating, and
+// running migrations using go-pg-pg.
 package migrations
 
 import (
 	"errors"
 	"fmt"
+	"time"
+
+	"github.com/go-pg/pg"
+	"github.com/go-pg/pg/orm"
 )
 
 // Errors that can be returned from Run.
 var (
 	ErrCreateRequiresName = errors.New("migration name is required for create")
 )
 
+type migration struct {
+	tableName struct{} `sql:"migrations,alias:migrations"`
+
+	ID          int32
+	Name        string
+	Batch       int32
+	CompletedAt time.Time
+	Up          func(orm.DB) error `sql:"-"`
+	Down        func(orm.DB) error `sql:"-"`
+
+	DisableTransaction bool `sql:"-"`
+}
+
+type lock struct {
+	tableName struct{} `sql:"migration_lock,alias:migration_lock"`
+
+	ID       string
+	IsLocked bool `sql:",notnull"`
+}
+
+const lockID = "lock"
+
 // Run takes in a directory and an argument slice and runs the appropriate command.
-func Run(directory string, args []string) error {
+func Run(db *pg.DB, directory string, args []string) error {
+	err := ensureMigrationTables(db)
+	if err != nil {
+		return err
+	}
+
 	cmd := ""
 
 	if len(args) > 1 {

migrations_test.go

@@ -4,24 +4,30 @@ import (
 	"os"
 	"testing"
 
+	"github.com/go-pg/pg"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestRun(t *testing.T) {
 	tmp := os.TempDir()
+	db := pg.Connect(&pg.Options{
+		Addr:     "localhost:5432",
+		User:     os.Getenv("TEST_DATABASE_USER"),
+		Database: os.Getenv("TEST_DATABASE_NAME"),
+	})
 
-	err := Run(tmp, []string{"cmd"})
+	err := Run(db, tmp, []string{"cmd"})
 	assert.Nil(t, err)
 
-	err = Run(tmp, []string{"cmd", "migrate"})
+	err = Run(db, tmp, []string{"cmd", "migrate"})
 	assert.Nil(t, err)
 
-	err = Run(tmp, []string{"cmd", "create"})
+	err = Run(db, tmp, []string{"cmd", "create"})
 	assert.Equal(t, ErrCreateRequiresName, err)
 
-	err = Run(tmp, []string{"cmd", "create", "test_migration"})
+	err = Run(db, tmp, []string{"cmd", "create", "test_migration"})
 	assert.Nil(t, err)
 
-	err = Run(tmp, []string{"cmd", "rollback"})
+	err = Run(db, tmp, []string{"cmd", "rollback"})
 	assert.Nil(t, err)
 }

setup.go

@@ -0,0 +1,59 @@
+package migrations
+
+import "github.com/go-pg/pg/orm"
+
+func ensureMigrationTables(db orm.DB) error {
+	exists, err := checkIfTableExists("migrations", db)
+	if err != nil {
+		return err
+	}
+	if !exists {
+		err = createTable(&migration{}, db)
+		if err != nil {
+			return err
+		}
+	}
+
+	exists, err = checkIfTableExists("migration_lock", db)
+	if err != nil {
+		return err
+	}
+	if !exists {
+		err = createTable(&lock{}, db)
+		if err != nil {
+			return err
+		}
+	}
+
+	count, err := db.Model(&lock{}).Count()
+	if err != nil {
+		return err
+	}
+	if count == 0 {
+		l := lock{ID: lockID, IsLocked: false}
+		err = db.Insert(&l)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func checkIfTableExists(name string, db orm.DB) (bool, error) {
+	count, err := orm.NewQuery(db).
+		Table("information_schema.tables").
+		Where("table_name = ?", name).
+		Where("table_schema = current_schema").
+		Count()
+	if err != nil {
+		return false, err
+	}
+	return count > 0, nil
+}
+
+func createTable(model interface{}, db orm.DB) error {
+	opts := orm.CreateTableOptions{IfNotExists: true}
+	_, err := orm.CreateTable(db, model, &opts)
+	return err
+}

setup_test.go

@@ -0,0 +1,73 @@
+package migrations
+
+import (
+	"os"
+	"testing"
+
+	"github.com/go-pg/pg"
+	"github.com/go-pg/pg/orm"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEnsureMigrationTables(t *testing.T) {
+	db := pg.Connect(&pg.Options{
+		Addr:     "localhost:5432",
+		User:     os.Getenv("TEST_DATABASE_USER"),
+		Database: os.Getenv("TEST_DATABASE_NAME"),
+	})
+
+	// drop tables to start from a clean database
+	dropMigrationTables(t, db)
+
+	err := ensureMigrationTables(db)
+	assert.Nil(t, err)
+
+	tables := []string{"migrations", "migration_lock"}
+
+	for _, table := range tables {
+		assertTable(t, db, table)
+	}
+
+	assertOneLock(t, db)
+
+	// with existing tables, ensureMigrationTables should do anything
+	err = ensureMigrationTables(db)
+	assert.Nil(t, err)
+
+	for _, table := range tables {
+		assertTable(t, db, table)
+	}
+
+	assertOneLock(t, db)
+}
+
+func dropMigrationTables(t *testing.T, db *pg.DB) {
+	t.Helper()
+
+	_, err := db.Exec("DROP TABLE migrations")
+	assert.Nil(t, err)
+	_, err = db.Exec("DROP TABLE migration_lock")
+	assert.Nil(t, err)
+}
+
+func assertTable(t *testing.T, db *pg.DB, table string) {
+	t.Helper()
+
+	count, err := orm.NewQuery(db).
+		Table("information_schema.tables").
+		Where("table_name = ?", table).
+		Where("table_schema = current_schema").
+		Count()
+	assert.Nil(t, err)
+	assert.Equalf(t, 1, count, "expected %q table to exist", table)
+}
+
+func assertOneLock(t *testing.T, db *pg.DB) {
+	t.Helper()
+
+	count, err := orm.NewQuery(db).
+		Table("migration_lock").
+		Count()
+	assert.Nil(t, err)
+	assert.Equal(t, 1, count, "expected migraions_lock to have a row")
+}