Merge branch 'postgres-lock-fix'

Addresses: https://github.com/golang-migrate/migrate/pull/4

database/postgres/postgres.go

@@ -1,3 +1,5 @@
+// +build go1.9
+
 package postgres
 
 import (
@@ -7,9 +9,10 @@ import (
 	"io/ioutil"
 	nurl "net/url"
 
-	"github.com/lib/pq"
+	"context"
 	"github.com/golang-migrate/migrate"
 	"github.com/golang-migrate/migrate/database"
+	"github.com/lib/pq"
 )
 
 func init() {
@@ -33,7 +36,8 @@ type Config struct {
 }
 
 type Postgres struct {
-	db       *sql.DB
+	// Locking and unlocking need to use the same connection
+	conn     *sql.Conn
 	isLocked bool
 
 	// Open and WithInstance need to garantuee that config is never nil
@@ -65,8 +69,14 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
 		config.MigrationsTable = DefaultMigrationsTable
 	}
 
+	conn, err := instance.Conn(context.Background())
+
+	if err != nil {
+		return nil, err
+	}
+
 	px := &Postgres{
-		db:     instance,
+		conn:   conn,
 		config: config,
 	}
 
@@ -105,7 +115,7 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
 }
 
 func (p *Postgres) Close() error {
-	return p.db.Close()
+	return p.conn.Close()
 }
 
 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
@@ -123,7 +133,7 @@ func (p *Postgres) Lock() error {
 	// 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 {
+	if err := p.conn.QueryRowContext(context.Background(), query, aid).Scan(&success); err != nil {
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	}
 
@@ -146,7 +156,7 @@ func (p *Postgres) Unlock() error {
 	}
 
 	query := `SELECT pg_advisory_unlock($1)`
-	if _, err := p.db.Exec(query, aid); err != nil {
+	if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 	p.isLocked = false
@@ -161,7 +171,7 @@ func (p *Postgres) Run(migration io.Reader) error {
 
 	// run migration
 	query := string(migr[:])
-	if _, err := p.db.Exec(query); err != nil {
+	if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
 		// TODO: cast to postgress error and get line number
 		return database.Error{OrigErr: err, Err: "migration failed", Query: migr}
 	}
@@ -170,7 +180,7 @@ func (p *Postgres) Run(migration io.Reader) error {
 }
 
 func (p *Postgres) SetVersion(version int, dirty bool) error {
-	tx, err := p.db.Begin()
+	tx, err := p.conn.BeginTx(context.Background(), &sql.TxOptions{})
 	if err != nil {
 		return &database.Error{OrigErr: err, Err: "transaction start failed"}
 	}
@@ -198,7 +208,7 @@ func (p *Postgres) SetVersion(version int, dirty bool) error {
 
 func (p *Postgres) Version() (version int, dirty bool, err error) {
 	query := `SELECT version, dirty FROM "` + p.config.MigrationsTable + `" LIMIT 1`
-	err = p.db.QueryRow(query).Scan(&version, &dirty)
+	err = p.conn.QueryRowContext(context.Background(), query).Scan(&version, &dirty)
 	switch {
 	case err == sql.ErrNoRows:
 		return database.NilVersion, false, nil
@@ -219,7 +229,7 @@ func (p *Postgres) Version() (version int, dirty bool, err error) {
 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)
+	tables, err := p.conn.QueryContext(context.Background(), query)
 	if err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
@@ -241,7 +251,7 @@ func (p *Postgres) Drop() error {
 		// delete one by one ...
 		for _, t := range tableNames {
 			query = `DROP TABLE IF EXISTS ` + t + ` CASCADE`
-			if _, err := p.db.Exec(query); err != nil {
+			if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
 				return &database.Error{OrigErr: err, Query: []byte(query)}
 			}
 		}
@@ -257,7 +267,7 @@ 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, p.config.MigrationsTable).Scan(&count); err != nil {
+	if err := p.conn.QueryRowContext(context.Background(), query, p.config.MigrationsTable).Scan(&count); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 	if count == 1 {
@@ -266,7 +276,7 @@ func (p *Postgres) ensureVersionTable() error {
 
 	// if not, create the empty migration table
 	query = `CREATE TABLE "` + p.config.MigrationsTable + `" (version bigint not null primary key, dirty boolean not null)`
-	if _, err := p.db.Exec(query); err != nil {
+	if _, err := p.conn.ExecContext(context.Background(), query); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 	return nil

database/postgres/postgres_test.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"testing"
 
+	"context"
 	dt "github.com/golang-migrate/migrate/database/testing"
 	mt "github.com/golang-migrate/migrate/testing"
 	// "github.com/lib/pq"
@@ -72,7 +73,7 @@ func TestMultiStatement(t *testing.T) {
 
 			// make sure second table exists
 			var exists bool
-			if err := d.(*Postgres).db.QueryRow("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'bar' AND table_schema = (SELECT current_schema()))").Scan(&exists); err != nil {
+			if err := d.(*Postgres).conn.QueryRowContext(context.Background(), "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'bar' AND table_schema = (SELECT current_schema()))").Scan(&exists); err != nil {
 				t.Fatal(err)
 			}
 			if !exists {
@@ -154,3 +155,39 @@ func TestWithSchema(t *testing.T) {
 func TestWithInstance(t *testing.T) {
 
 }
+
+func TestPostgres_Lock(t *testing.T) {
+	mt.ParallelTest(t, versions, isReady,
+		func(t *testing.T, i mt.Instance) {
+			p := &Postgres{}
+			addr := fmt.Sprintf("postgres://postgres@%v:%v/postgres?sslmode=disable", i.Host(), i.Port())
+			d, err := p.Open(addr)
+			if err != nil {
+				t.Fatalf("%v", err)
+			}
+
+			dt.Test(t, d, []byte("SELECT 1"))
+
+			ps := d.(*Postgres)
+
+			err = ps.Lock()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			err = ps.Unlock()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			err = ps.Lock()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			err = ps.Unlock()
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+}