Skip to content

Commit

Permalink
Added support for pgx locking table
Browse files Browse the repository at this point in the history 
In order to support running migrations through PgBouncer which does not
support advisory locks.
  • Loading branch information
@fortnox-andreas
fortnox-andreas committed Sep 20, 2023
1 parent 856ea12 commit 129922a
Showing 1 changed file with 139 additions and 19 deletions.
158 changes: 139 additions & 19 deletions database/pgx/pgx.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,11 @@ import (
_ "github.com/jackc/pgx/v4/stdlib"
)

const (
LockStrategyAdvisory = "advisory"
LockStrategyTable = "table"
)

func init() {
db := Postgres{}
database.Register("pgx", &db)
Expand All @@ -36,6 +41,8 @@ var (

DefaultMigrationsTable = "schema_migrations"
DefaultMultiStatementMaxSize = 10 * 1 << 20 // 10 MB
DefaultLockTable = "schema_lock"
DefaultLockStrategy = "advisory"
)

var (
Expand All @@ -49,6 +56,8 @@ type Config struct {
MigrationsTable string
DatabaseName string
SchemaName string
LockTable string
LockStrategy string
migrationsSchemaName string
migrationsTableName string
StatementTimeout time.Duration
Expand Down Expand Up @@ -108,6 +117,14 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
config.MigrationsTable = DefaultMigrationsTable
}

if len(config.LockTable) == 0 {
config.LockTable = DefaultLockTable
}

if len(config.LockStrategy) == 0 {
config.LockStrategy = DefaultLockStrategy
}

config.migrationsSchemaName = config.SchemaName
config.migrationsTableName = config.MigrationsTable
if config.MigrationsTableQuoted {
Expand All @@ -133,6 +150,10 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
config: config,
}

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

if err := px.ensureVersionTable(); err != nil {
return nil, err
}
Expand Down Expand Up @@ -196,13 +217,16 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
}
}

lockStrategy := purl.Query().Get("x-lock-strategy")

px, err := WithInstance(db, &Config{
DatabaseName: purl.Path,
MigrationsTable: migrationsTable,
MigrationsTableQuoted: migrationsTableQuoted,
StatementTimeout: time.Duration(statementTimeout) * time.Millisecond,
MultiStatementEnabled: multiStatementEnabled,
MultiStatementMaxSize: multiStatementMaxSize,
LockStrategy: lockStrategy,
})

if err != nil {
Expand All @@ -221,36 +245,110 @@ func (p *Postgres) Close() error {
return nil
}

// https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
func (p *Postgres) Lock() error {
return database.CasRestoreOnErr(&p.isLocked, false, true, database.ErrLocked, func() error {
aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
if err != nil {
return err
}

// This will wait indefinitely until the lock can be acquired.
query := `SELECT pg_advisory_lock($1)`
if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
switch p.config.LockStrategy {
case LockStrategyAdvisory:
return p.applyAdvisoryLock()
case LockStrategyTable:
return p.applyTableLock()
default:
return fmt.Errorf("unknown lock strategy \"%s\"", p.config.LockStrategy)
}
return nil
})
}

func (p *Postgres) Unlock() error {
return database.CasRestoreOnErr(&p.isLocked, true, false, database.ErrNotLocked, func() error {
aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
if err != nil {
return err
switch p.config.LockStrategy {
case LockStrategyAdvisory:
return p.releaseAdvisoryLock()
case LockStrategyTable:
return p.releaseTableLock()
default:
return fmt.Errorf("unknown lock strategy \"%s\"", p.config.LockStrategy)
}
})
}

query := `SELECT pg_advisory_unlock($1)`
if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
// https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
func (p *Postgres) applyAdvisoryLock() error {
aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
if err != nil {
return err
}

// This will wait indefinitely until the lock can be acquired.
query := `SELECT pg_advisory_lock($1)`
if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
}
return nil
}

func (p *Postgres) applyTableLock() error {
tx, err := p.conn.BeginTx(context.Background(), &sql.TxOptions{})
if err != nil {
return &database.Error{OrigErr: err, Err: "transaction start failed"}
}

aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName)
if err != nil {
return err
}

query := "SELECT * FROM " + p.config.LockTable + " WHERE lock_id = $1"
rows, err := tx.Query(query, aid)
if err != nil {
return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
}

defer func() {
if errClose := rows.Close(); errClose != nil {
err = multierror.Append(err, errClose)
}
return nil
})
}()

// If row exists at all, lock is present
locked := rows.Next()
if locked {
return database.ErrLocked
}

query = "INSERT INTO " + p.config.LockTable + " (lock_id) VALUES ($1)"
if _, err := tx.Exec(query, aid); err != nil {
return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
}

return tx.Commit()
}

func (p *Postgres) releaseAdvisoryLock() error {
aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
if err != nil {
return err
}

query := `SELECT pg_advisory_unlock($1)`
if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
}

return nil
}

func (p *Postgres) releaseTableLock() error {
aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName)
if err != nil {
return err
}

query := "DELETE FROM " + p.config.LockTable + " WHERE lock_id = $1"
if _, err := p.db.Exec(query, aid); err != nil {
return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
}

return nil
}

func (p *Postgres) Run(migration io.Reader) error {
Expand Down Expand Up @@ -478,6 +576,28 @@ func (p *Postgres) ensureVersionTable() (err error) {
return nil
}

func (p *Postgres) ensureLockTable() error {
if p.config.LockStrategy != LockStrategyTable {
return nil
}

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.LockTable).Scan(&count); err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
}
if count == 1 {
return nil
}

query = `CREATE TABLE "` + p.config.LockTable + `" (lock_id BIGINT NOT NULL PRIMARY KEY)`
if _, err := p.db.Exec(query); err != nil {
return &database.Error{OrigErr: err, Query: []byte(query)}
}

return nil
}

// Copied from lib/pq implementation: https://github.com/lib/pq/blob/v1.9.0/conn.go#L1611
func quoteIdentifier(name string) string {
end := strings.IndexRune(name, 0)
Expand Down

0 comments on commit 129922a

Please sign in to comment.