Repository: ariga/atlas
Pull request #648 (2022-03-11 15:20:58)

sql/postgres: implement schema.Locker
Please, don't review yet. Working on integration tests. 


Patch:
diff --git a/sql/internal/sqlx/dev.go b/sql/internal/sqlx/dev.go
index 0247e931..631f88ae 100644
--- a/sql/internal/sqlx/dev.go
+++ b/sql/internal/sqlx/dev.go
@@ -112,5 +112,6 @@ func (d *DevDriver) formatName(name string) string {
 		return dev
 	}
 	h := fnv.New128()
-	return fmt.Sprintf("%s_%x", dev[:d.MaxNameLen-1-h.Size()*2], h.Sum([]byte(dev)))
+	h.Write([]byte(dev))
+	return fmt.Sprintf("%s_%x", dev[:d.MaxNameLen-1-h.Size()*2], h.Sum(nil))
 }
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index d49f1ab6..7b601397 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -6,8 +6,11 @@ package sqlx
 
 import (
 	"bytes"
+	"context"
 	"database/sql"
+	"database/sql/driver"
 	"fmt"
+	"io"
 	"reflect"
 	"strconv"
 	"strings"
@@ -15,6 +18,41 @@ import (
 	"ariga.io/atlas/sql/schema"
 )
 
+type (
+	// ExecQueryCloser is the interface that groups
+	// Close with the schema.ExecQuerier methods.
+	ExecQueryCloser interface {
+		schema.ExecQuerier
+		io.Closer
+	}
+	nopCloser struct {
+		schema.ExecQuerier
+	}
+)
+
+// Close implements the io.Closer interface.
+func (nopCloser) Close() error { return nil }
+
+// SingleConn returns a closable single connection from the given ExecQuerier.
+// If the ExecQuerier is already bound to a single connection (e.g. Tx, Conn),
+// the connection will return as-is with a NopCloser.
+func SingleConn(ctx context.Context, conn schema.ExecQuerier) (ExecQueryCloser, error) {
+	// A standard sql.DB or a wrapper of it.
+	if opener, ok := conn.(interface {
+		Conn(context.Context) (*sql.Conn, error)
+	}); ok {
+		return opener.Conn(ctx)
+	}
+	// Tx and Conn are bounded to a single connection.
+	// We use sql/driver.Tx to cover also custom Tx structs.
+	_, ok1 := conn.(driver.Tx)
+	_, ok2 := conn.(*sql.Conn)
+	if ok1 || ok2 {
+		return nopCloser{ExecQuerier: conn}, nil
+	}
+	return nil, fmt.Errorf("cannot obtain a single connection from %T", conn)
+}
+
 // ValidString reports if the given string is not null and valid.
 func ValidString(s sql.NullString) bool {
 	return s.Valid && s.String != "" && strings.ToLower(s.String) != "null"
@@ -32,6 +70,12 @@ func ScanOne(rows *sql.Rows, dest ...interface{}) error {
 	return rows.Close()
 }
 
+// ScanNullBool scans one sql.NullBool record and closes the rows at the end.
+func ScanNullBool(rows *sql.Rows) (sql.NullBool, error) {
+	var b sql.NullBool
+	return b, ScanOne(rows, &b)
+}
+
 // SchemaFKs scans the rows and adds the foreign-key to the schema table.
 // Reference elements are added as stubs and should be linked manually by the
 // caller.
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index cabda15c..22a42630 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -6,9 +6,7 @@ package mysql
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
-	"io"
 	"strings"
 	"time"
 
@@ -80,64 +78,48 @@ func (d *Driver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*schema
 }
 
 // Lock implements the schema.Locker interface.
-func (d *Driver) Lock(ctx context.Context, name string, timeout time.Duration) (unlock func() error, err error) {
-	var (
-		closer io.Closer
-		conn   = d.ExecQuerier
-	)
-	// If the connection is sql.DB (a connection-pool),
-	// acquire a connection and use it to obtain the lock.
-	if opener, ok := conn.(interface {
-		Conn(context.Context) (*sql.Conn, error)
-	}); ok {
-		c, err := opener.Conn(ctx)
-		if err != nil {
-			return nil, err
-		}
-		defer func() {
-			if err != nil {
-				// Return the connection to the pool,
-				// if we failed to obtain the lock.
-				closer.Close()
-			}
-		}()
-		conn, closer = c, c
-	}
-	rows, err := conn.QueryContext(ctx, "SELECT GET_LOCK(?, ?)", name, int(timeout.Seconds()))
+func (d *Driver) Lock(ctx context.Context, name string, timeout time.Duration) (schema.UnlockFunc, error) {
+	conn, err := sqlx.SingleConn(ctx, d.ExecQuerier)
 	if err != nil {
 		return nil, err
 	}
-	var acquired sql.NullBool
-	switch err := sqlx.ScanOne(rows, &acquired); {
-	case err != nil:
+	if err := acquire(ctx, conn, name, timeout); err != nil {
+		conn.Close()
 		return nil, err
-	case !acquired.Valid:
-		// NULL is returned in case of an unexpected internal error.
-		return nil, fmt.Errorf("sql/mysql: unexpected internal error on Lock(%q, %s)", name, timeout)
-	case !acquired.Bool:
-		return nil, schema.ErrLockTimeout
 	}
-	// Acquired.
 	return func() error {
-		if closer != nil {
-			defer closer.Close()
-		}
-		rows, err = conn.QueryContext(ctx, "SELECT RELEASE_LOCK(?)", name)
+		defer conn.Close()
+		rows, err := conn.QueryContext(ctx, "SELECT RELEASE_LOCK(?)", name)
 		if err != nil {
 			return err
 		}
-		var released sql.NullBool
-		switch err := sqlx.ScanOne(rows, &released); {
+		switch released, err := sqlx.ScanNullBool(rows); {
 		case err != nil:
 			return err
 		case !released.Valid || !released.Bool:
 			return fmt.Errorf("sql/mysql: failed releasing a named lock %q", name)
-		default: // Released.
-			return nil
 		}
+		return nil
 	}, nil
 }
 
+func acquire(ctx context.Context, conn schema.ExecQuerier, name string, timeout time.Duration) error {
+	rows, err := conn.QueryContext(ctx, "SELECT GET_LOCK(?, ?)", name, int(timeout.Seconds()))
+	if err != nil {
+		return err
+	}
+	switch acquired, err := sqlx.ScanNullBool(rows); {
+	case err != nil:
+		return err
+	case !acquired.Valid:
+		// NULL is returned in case of an unexpected internal error.
+		return fmt.Errorf("sql/mysql: unexpected internal error on Lock(%q, %s)", name, timeout)
+	case !acquired.Bool:
+		return schema.ErrLocked
+	}
+	return nil
+}
+
 // supportsCheck reports if the connected database supports
 // the CHECK clause, and return the querying for getting them.
 func (d *conn) supportsCheck() (string, bool) {
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index b8b5ef00..0235209d 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -6,7 +6,10 @@ package postgres
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"hash/fnv"
+	"time"
 
 	"ariga.io/atlas/sql/internal/sqlx"
 	"ariga.io/atlas/sql/migrate"
@@ -79,6 +82,70 @@ func (d *Driver) NormalizeSchema(ctx context.Context, s *schema.Schema) (*schema
 	return d.dev().NormalizeSchema(ctx, s)
 }
 
+// Lock implements the schema.Locker interface.
+func (d *Driver) Lock(ctx context.Context, name string, timeout time.Duration) (schema.UnlockFunc, error) {
+	conn, err := sqlx.SingleConn(ctx, d.ExecQuerier)
+	if err != nil {
+		return nil, err
+	}
+	h := fnv.New32()
+	h.Write([]byte(name))
+	id := h.Sum32()
+	if err := acquire(ctx, conn, id, timeout); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	return func() error {
+		defer conn.Close()
+		rows, err := conn.QueryContext(ctx, "SELECT pg_advisory_unlock($1)", id)
+		if err != nil {
+			return err
+		}
+		switch released, err := sqlx.ScanNullBool(rows); {
+		case err != nil:
+			return err
+		case !released.Valid || !released.Bool:
+			return fmt.Errorf("sql/postgres: failed releasing lock %d", id)
+		}
+		return nil
+	}, nil
+}
+
+func acquire(ctx context.Context, conn schema.ExecQuerier, id uint32, timeout time.Duration) error {
+	switch {
+	// With timeout (context-based).
+	case timeout > 0:
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+		fallthrough
+	// Infinite timeout.
+	case timeout < 0:
+		rows, err := conn.QueryContext(ctx, "SELECT pg_advisory_lock($1)", id)
+		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+			err = schema.ErrLocked
+		}
+		if err != nil {
+			return err
+		}
+		return rows.Close()
+	// No timeout.
+	default:
+		rows, err := conn.QueryContext(ctx, "SELECT pg_try_advisory_lock($1)", id)
+		if err != nil {
+			return err
+		}
+		acquired, err := sqlx.ScanNullBool(rows)
+		if err != nil {
+			return err
+		}
+		if !acquired.Bool {
+			return schema.ErrLocked
+		}
+		return nil
+	}
+}
+
 // Standard column types (and their aliases) as defined in
 // PostgreSQL codebase/website.
 const (
diff --git a/sql/schema/migrate.go b/sql/schema/migrate.go
index afda08de..7ff84672 100644
--- a/sql/schema/migrate.go
+++ b/sql/schema/migrate.go
@@ -246,18 +246,25 @@ type Differ interface {
 	TableDiff(from, to *Table) ([]Change, error)
 }
 
-// ErrLockTimeout is returned on Lock calls which have timed out.
-var ErrLockTimeout = errors.New("sql/schema: lock timeout")
-
-// Locker is an interface that is optionally implemented by the different drivers
-// for obtaining an "advisory lock" with the given name.
-type Locker interface {
-	// Lock acquires a named "advisory lock", using the timeout. Negative value means no timeout.
-	// The returned unlock function is used to release the advisory lock acquired by the session.
-	//
-	// An ErrLockTimeout is returned if the operation timed out before obtaining the lock.
-	Lock(ctx context.Context, name string, timeout time.Duration) (unlock func() error, err error)
-}
+// ErrLocked is returned on Lock calls which have failed to obtain the lock.
+var ErrLocked = errors.New("sql/schema: lock is held by other session")
+
+type (
+	// UnlockFunc is returned by the Locker to explicitly
+	// release the named "advisory lock".
+	UnlockFunc func() error
+
+	// Locker is an interface that is optionally implemented by the different drivers
+	// for obtaining an "advisory lock" with the given name.
+	Locker interface {
+		// Lock acquires a named "advisory lock", using the given timeout. Negative value means no timeout,
+		// and the zero value means a "try lock" mode. i.e. return immediately if the lock is already taken.
+		// The returned unlock function is used to release the advisory lock acquired by the session.
+		//
+		// An ErrLocked is returned if the operation failed to obtain the lock in all different timeout modes.
+		Lock(ctx context.Context, name string, timeout time.Duration) (UnlockFunc, error)
+	}
+)
 
 // changes.
 func (*AddAttr) change()          {}


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 61bc558a..f33c547c 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -431,3 +431,24 @@ func sanity(c *ent.Client) {
 		SetInfo(c.GroupInfo.Create().SetDesc("desc").SaveX(ctx)).
 		SaveX(ctx)
 }
+
+func testAdvisoryLock(t *testing.T, l schema.Locker) {
+	t.Run("One", func(t *testing.T) {
+		unlock, err := l.Lock(context.Background(), "migrate", 0)
+		require.NoError(t, err)
+		_, err = l.Lock(context.Background(), "migrate", 0)
+		require.Equal(t, schema.ErrLocked, err)
+		require.NoError(t, unlock())
+	})
+	t.Run("Multi", func(t *testing.T) {
+		var unlocks []schema.UnlockFunc
+		for _, name := range []string{"a", "b", "c"} {
+			unlock, err := l.Lock(context.Background(), name, 0)
+			require.NoError(t, err)
+			unlocks = append(unlocks, unlock)
+		}
+		for _, unlock := range unlocks {
+			require.NoError(t, unlock())
+		}
+	})
+}
diff --git a/internal/integration/mysql_test.go b/internal/integration/mysql_test.go
index f1f43e1c..2a071b69 100644
--- a/internal/integration/mysql_test.go
+++ b/internal/integration/mysql_test.go
@@ -460,6 +460,12 @@ func TestMySQL_Ent(t *testing.T) {
 	})
 }
 
+func TestMySQL_AdvisoryLock(t *testing.T) {
+	myRun(t, func(t *myTest) {
+		testAdvisoryLock(t.T, t.drv)
+	})
+}
+
 func TestMySQL_HCL(t *testing.T) {
 	full := `
 schema "test" {
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index 7aa16d52..8b6dc47b 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -368,6 +368,12 @@ func TestPostgres_Ent(t *testing.T) {
 	})
 }
 
+func TestPostgres_AdvisoryLock(t *testing.T) {
+	pgRun(t, func(t *pgTest) {
+		testAdvisoryLock(t.T, t.drv)
+	})
+}
+
 func TestPostgres_HCL(t *testing.T) {
 	full := `
 schema "public" {
diff --git a/sql/mysql/driver_test.go b/sql/mysql/driver_test.go
index bf25895e..cd3661a2 100644
--- a/sql/mysql/driver_test.go
+++ b/sql/mysql/driver_test.go
@@ -74,7 +74,7 @@ func TestDriver_LockError(t *testing.T) {
 			WillReturnRows(sqlmock.NewRows([]string{"acquired"}).AddRow(0)).
 			RowsWillBeClosed()
 		unlock, err := d.Lock(context.Background(), name, time.Second)
-		require.Equal(t, schema.ErrLockTimeout, err)
+		require.Equal(t, schema.ErrLocked, err)
 		require.Nil(t, unlock)
 	})
 
diff --git a/sql/postgres/driver_test.go b/sql/postgres/driver_test.go
new file mode 100644
index 00000000..3b60b80e
--- /dev/null
+++ b/sql/postgres/driver_test.go
@@ -0,0 +1,103 @@
+// Copyright 2021-present The Atlas Authors. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package postgres
+
+import (
+	"context"
+	"io"
+	"testing"
+	"time"
+
+	"ariga.io/atlas/sql/internal/sqltest"
+	"ariga.io/atlas/sql/schema"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDriver_LockAcquired(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	name, hash := "name", 797654004
+	m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
+		WithArgs(hash).
+		WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_lock"}).AddRow(1)).
+		RowsWillBeClosed()
+	m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
+		WithArgs(hash).
+		WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(1)).
+		RowsWillBeClosed()
+
+	d := &Driver{}
+	d.ExecQuerier = db
+	unlock, err := d.Lock(context.Background(), name, 0)
+	require.NoError(t, err)
+	require.NoError(t, unlock())
+	require.NoError(t, m.ExpectationsWereMet())
+}
+
+func TestDriver_LockError(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	d := &Driver{}
+	d.ExecQuerier = db
+	name, hash := "migrate", 979249972
+
+	t.Run("Timeout", func(t *testing.T) {
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnError(context.DeadlineExceeded).
+			RowsWillBeClosed()
+		unlock, err := d.Lock(context.Background(), name, time.Millisecond)
+		require.Equal(t, schema.ErrLocked, err)
+		require.Nil(t, unlock)
+	})
+
+	t.Run("Internal", func(t *testing.T) {
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnError(io.EOF).
+			RowsWillBeClosed()
+		unlock, err := d.Lock(context.Background(), name, time.Minute)
+		require.Equal(t, io.EOF, err)
+		require.Nil(t, unlock)
+	})
+}
+
+func TestDriver_UnlockError(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	d := &Driver{}
+	d.ExecQuerier = db
+	name, hash := "up", 1551306158
+	acquired := func() {
+		m.ExpectQuery(sqltest.Escape("SELECT pg_try_advisory_lock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(1)).
+			RowsWillBeClosed()
+	}
+
+	t.Run("NotHeld", func(t *testing.T) {
+		acquired()
+		unlock, err := d.Lock(context.Background(), name, 0)
+		require.NoError(t, err)
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(0)).
+			RowsWillBeClosed()
+		require.Error(t, unlock())
+	})
+
+	t.Run("Internal", func(t *testing.T) {
+		acquired()
+		unlock, err := d.Lock(context.Background(), name, 0)
+		require.NoError(t, err)
+		m.ExpectQuery(sqltest.Escape("SELECT pg_advisory_unlock($1)")).
+			WithArgs(hash).
+			WillReturnRows(sqlmock.NewRows([]string{"pg_advisory_unlock"}).AddRow(nil)).
+			RowsWillBeClosed()
+		require.Error(t, unlock())
+	})
+}