Repository: ariga/atlas
Pull request #646 (2022-03-10 15:44:42)

sql: introduce schema.Locker


Patch:
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index 028e96cc..cabda15c 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -6,8 +6,11 @@ package mysql
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
+	"io"
 	"strings"
+	"time"
 
 	"ariga.io/atlas/sql/internal/sqlx"
 	"ariga.io/atlas/sql/migrate"
@@ -76,6 +79,65 @@ 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) (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()))
+	if err != nil {
+		return nil, err
+	}
+	var acquired sql.NullBool
+	switch err := sqlx.ScanOne(rows, &acquired); {
+	case err != nil:
+		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)
+		if err != nil {
+			return err
+		}
+		var released sql.NullBool
+		switch err := sqlx.ScanOne(rows, &released); {
+		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
+		}
+	}, 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/schema/migrate.go b/sql/schema/migrate.go
index e15dd853..afda08de 100644
--- a/sql/schema/migrate.go
+++ b/sql/schema/migrate.go
@@ -4,6 +4,12 @@
 
 package schema
 
+import (
+	"context"
+	"errors"
+	"time"
+)
+
 type (
 	// A Change represents a schema change. The types below implement this
 	// interface and can be used for describing schema changes.
@@ -240,6 +246,19 @@ 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)
+}
+
 // changes.
 func (*AddAttr) change()          {}
 func (*DropAttr) change()         {}


Test patch:
diff --git a/sql/mysql/driver_test.go b/sql/mysql/driver_test.go
new file mode 100644
index 00000000..bf25895e
--- /dev/null
+++ b/sql/mysql/driver_test.go
@@ -0,0 +1,138 @@
+// 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 mysql
+
+import (
+	"context"
+	"database/sql"
+	"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)
+	lock := func(l schema.Locker) {
+		name, sec := "name", 1
+		m.ExpectQuery(sqltest.Escape("SELECT GET_LOCK(?, ?)")).
+			WithArgs(name, sec).
+			WillReturnRows(sqlmock.NewRows([]string{"acquired"}).AddRow(1)).
+			RowsWillBeClosed()
+		m.ExpectQuery(sqltest.Escape("SELECT RELEASE_LOCK(?)")).
+			WithArgs(name).
+			WillReturnRows(sqlmock.NewRows([]string{"released"}).AddRow(1)).
+			RowsWillBeClosed()
+		unlock, err := l.Lock(context.Background(), name, time.Second)
+		require.NoError(t, err)
+		require.NoError(t, unlock())
+	}
+
+	t.Run("OnPool", func(t *testing.T) {
+		d := &Driver{}
+		d.ExecQuerier = &mockOpener{DB: db}
+		lock(d)
+		require.EqualValues(t, 1, d.ExecQuerier.(*mockOpener).opened)
+	})
+
+	t.Run("OnConn", func(t *testing.T) {
+		conn, err := db.Conn(context.Background())
+		require.NoError(t, err)
+		d := &Driver{}
+		d.ExecQuerier = conn
+		lock(d)
+	})
+
+	t.Run("OnTx", func(t *testing.T) {
+		m.ExpectBegin()
+		tx, err := db.Begin()
+		require.NoError(t, err)
+		d := &Driver{}
+		d.ExecQuerier = tx
+		lock(d)
+	})
+	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
+
+	t.Run("Timeout", func(t *testing.T) {
+		name, sec := "name", 1
+		m.ExpectQuery(sqltest.Escape("SELECT GET_LOCK(?, ?)")).
+			WithArgs(name, sec).
+			WillReturnRows(sqlmock.NewRows([]string{"acquired"}).AddRow(0)).
+			RowsWillBeClosed()
+		unlock, err := d.Lock(context.Background(), name, time.Second)
+		require.Equal(t, schema.ErrLockTimeout, err)
+		require.Nil(t, unlock)
+	})
+
+	t.Run("Internal", func(t *testing.T) {
+		name, sec := "migrate", -1
+		m.ExpectQuery(sqltest.Escape("SELECT GET_LOCK(?, ?)")).
+			WithArgs(name, sec).
+			WillReturnRows(sqlmock.NewRows([]string{"acquired"}).AddRow(nil)).
+			RowsWillBeClosed()
+		unlock, err := d.Lock(context.Background(), name, -time.Second)
+		require.EqualError(t, err, `sql/mysql: unexpected internal error on Lock("migrate", -1s)`)
+		require.Nil(t, unlock)
+	})
+}
+
+func TestDriver_UnlockError(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	d := &Driver{}
+	d.ExecQuerier = db
+	acquired := func(name string, sec int) {
+		m.ExpectQuery(sqltest.Escape("SELECT GET_LOCK(?, ?)")).
+			WithArgs(name, sec).
+			WillReturnRows(sqlmock.NewRows([]string{"acquired"}).AddRow(1)).
+			RowsWillBeClosed()
+	}
+
+	t.Run("NotHeld", func(t *testing.T) {
+		name, sec := "unknown_lock", 0
+		acquired(name, sec)
+		unlock, err := d.Lock(context.Background(), name, time.Millisecond)
+		require.NoError(t, err)
+		m.ExpectQuery(sqltest.Escape("SELECT RELEASE_LOCK(?)")).
+			WithArgs(name).
+			WillReturnRows(sqlmock.NewRows([]string{"released"}).AddRow(0)).
+			RowsWillBeClosed()
+		require.Error(t, unlock())
+	})
+
+	t.Run("Internal", func(t *testing.T) {
+		name, sec := "unknown_error", 1
+		acquired(name, sec)
+		unlock, err := d.Lock(context.Background(), name, time.Second+time.Millisecond)
+		require.NoError(t, err)
+		m.ExpectQuery(sqltest.Escape("SELECT RELEASE_LOCK(?)")).
+			WithArgs(name).
+			WillReturnRows(sqlmock.NewRows([]string{"released"}).AddRow(nil)).
+			RowsWillBeClosed()
+		require.Error(t, unlock())
+	})
+}
+
+type mockOpener struct {
+	*sql.DB
+	opened uint
+}
+
+func (m *mockOpener) Conn(ctx context.Context) (*sql.Conn, error) {
+	m.opened++
+	return m.DB.Conn(ctx)
+}