feat: implement ConnAcquirer and DedicatedConn interfaces for MySQL and PostgreSQL drivers

Author: juicycleff

driver/driver.go

@@ -121,3 +121,25 @@ type Stmt interface {
 	Exec(ctx context.Context, args ...any) (Result, error)
 	Close() error
 }
+
+// ConnAcquirer is an optional interface that pool-based drivers implement
+// to provide a dedicated connection from the pool. This is needed for
+// operations that require session-level state (e.g., advisory locks)
+// to remain on a single connection across multiple queries.
+type ConnAcquirer interface {
+	// AcquireConn acquires a dedicated connection from the pool.
+	// The returned DedicatedConn provides Exec/Query/QueryRow methods
+	// that are guaranteed to run on the same underlying connection.
+	// The caller MUST call Release() when done.
+	AcquireConn(ctx context.Context) (DedicatedConn, error)
+}
+
+// DedicatedConn represents a single, dedicated database connection
+// acquired from a pool. All operations execute on the same underlying
+// connection, making it safe for session-level state like advisory locks.
+type DedicatedConn interface {
+	Exec(ctx context.Context, query string, args ...any) (Result, error)
+	Query(ctx context.Context, query string, args ...any) (Rows, error)
+	QueryRow(ctx context.Context, query string, args ...any) Row
+	Release()
+}

drivers/mysqldriver/mysql.go

@@ -29,6 +29,7 @@ type MysqlDB struct {
 
 var _ driver.Driver = (*MysqlDB)(nil)
 var _ driver.Preparer = (*MysqlDB)(nil)
+var _ driver.ConnAcquirer = (*MysqlDB)(nil)
 
 // New creates a new unconnected MysqlDB. Call Open to establish a connection pool.
 func New() *MysqlDB {
@@ -178,6 +179,54 @@ func (db *MysqlDB) Prepare(ctx context.Context, query string) (driver.Stmt, erro
 	return &mysqlStmt{stmt: stmt}, nil
 }
 
+// AcquireConn acquires a dedicated connection from the pool.
+// All operations on the returned DedicatedConn execute on the same
+// underlying MySQL session, making it safe for session-level state
+// such as GET_LOCK advisory locks.
+func (db *MysqlDB) AcquireConn(ctx context.Context) (driver.DedicatedConn, error) {
+	conn, err := db.db.Conn(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("mysqldriver: acquire dedicated conn: %w", err)
+	}
+	return &mysqlDedicatedConn{conn: conn}, nil
+}
+
+// mysqlDedicatedConn wraps a single *sql.Conn to implement driver.DedicatedConn.
+// All queries execute on the same underlying connection.
+type mysqlDedicatedConn struct {
+	conn *sql.Conn
+}
+
+var _ driver.DedicatedConn = (*mysqlDedicatedConn)(nil)
+
+func (c *mysqlDedicatedConn) Exec(ctx context.Context, query string, args ...any) (driver.Result, error) {
+	res, err := c.conn.ExecContext(ctx, query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("mysqldriver: dedicated exec: %w", err)
+	}
+	return &mysqlResult{res: res}, nil
+}
+
+func (c *mysqlDedicatedConn) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) {
+	rows, err := c.conn.QueryContext(ctx, query, args...) //nolint:rowserrcheck // caller checks Err() via driver.Rows
+	if err != nil {
+		return nil, fmt.Errorf("mysqldriver: dedicated query: %w", err)
+	}
+	return &mysqlRows{rows: rows}, nil
+}
+
+func (c *mysqlDedicatedConn) QueryRow(ctx context.Context, query string, args ...any) driver.Row {
+	row := c.conn.QueryRowContext(ctx, query, args...)
+	return &mysqlRow{row: row}
+}
+
+func (c *mysqlDedicatedConn) Release() {
+	if c.conn != nil {
+		_ = c.conn.Close()
+		c.conn = nil
+	}
+}
+
 // mapIsolationLevel converts a driver.IsolationLevel to the corresponding
 // sql.IsolationLevel constant.
 func mapIsolationLevel(level driver.IsolationLevel) sql.IsolationLevel {

drivers/mysqldriver/mysqlmigrate/executor.go

@@ -24,8 +24,15 @@ const (
 )
 
 // Executor implements migrate.Executor for MySQL.
+//
+// When the underlying driver supports [driver.ConnAcquirer], the executor
+// acquires a dedicated connection in [AcquireLock] and routes ALL subsequent
+// operations through it until [ReleaseLock]. This guarantees that the
+// session-level GET_LOCK is acquired and released on the same connection,
+// preventing advisory-lock leaks in pooled environments.
 type Executor struct {
-	drv driver.Driver
+	drv       driver.Driver
+	dedicated driver.DedicatedConn // non-nil while migration lock is held
 }
 
 var _ migrate.Executor = (*Executor)(nil)
@@ -35,14 +42,44 @@ func New(drv driver.Driver) *Executor {
 	return &Executor{drv: drv}
 }
 
+// --- routing helpers: prefer dedicated conn when available ---
+
+func (e *Executor) exec(ctx context.Context, query string, args ...any) (driver.Result, error) {
+	if e.dedicated != nil {
+		return e.dedicated.Exec(ctx, query, args...)
+	}
+	return e.drv.Exec(ctx, query, args...)
+}
+
+func (e *Executor) query(ctx context.Context, query string, args ...any) (driver.Rows, error) {
+	if e.dedicated != nil {
+		return e.dedicated.Query(ctx, query, args...)
+	}
+	return e.drv.Query(ctx, query, args...)
+}
+
+func (e *Executor) queryRow(ctx context.Context, query string, args ...any) driver.Row {
+	if e.dedicated != nil {
+		return e.dedicated.QueryRow(ctx, query, args...)
+	}
+	return e.drv.QueryRow(ctx, query, args...)
+}
+
+func (e *Executor) releaseDedicated() {
+	if e.dedicated != nil {
+		e.dedicated.Release()
+		e.dedicated = nil
+	}
+}
+
 // Exec executes a SQL statement that does not return rows.
 func (e *Executor) Exec(ctx context.Context, query string, args ...any) (driver.Result, error) {
-	return e.drv.Exec(ctx, query, args...)
+	return e.exec(ctx, query, args...)
 }
 
 // Query executes a SQL statement that returns rows.
 func (e *Executor) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) {
-	return e.drv.Query(ctx, query, args...)
+	return e.query(ctx, query, args...)
 }
 
 // EnsureMigrationTable creates the grove_migrations table if it doesn't exist.
@@ -55,7 +92,7 @@ func (e *Executor) EnsureMigrationTable(ctx context.Context) error {
 		"`migrated_at` DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), "+
 		"UNIQUE KEY `uq_version_group` (`version`, `group`)"+
 		") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4", migrationTableName)
-	_, err := e.drv.Exec(ctx, query)
+	_, err := e.exec(ctx, query)
 	return err
 }
 
@@ -67,21 +104,37 @@ func (e *Executor) EnsureLockTable(ctx context.Context) error {
 		"`locked_by` VARCHAR(255), "+
 		"CONSTRAINT `single_lock` CHECK (`id` = 1)"+
 		") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4", lockTableName)
-	_, err := e.drv.Exec(ctx, query)
+	_, err := e.exec(ctx, query)
 	return err
 }
 
 // AcquireLock attempts to acquire the distributed migration lock using
 // MySQL advisory locks (GET_LOCK) for immediate feedback.
+//
+// If the driver implements [driver.ConnAcquirer], a dedicated connection is
+// acquired first so that the advisory lock and all subsequent migration
+// operations share the same session. This prevents lock leaks when using
+// connection pools.
 func (e *Executor) AcquireLock(ctx context.Context, lockedBy string) error {
+	// Acquire a dedicated connection if the driver supports it.
+	if acq, ok := e.drv.(driver.ConnAcquirer); ok {
+		conn, err := acq.AcquireConn(ctx)
+		if err != nil {
+			return fmt.Errorf("mysqlmigrate: acquire dedicated conn: %w", err)
+		}
+		e.dedicated = conn
+	}
+
 	// Use MySQL advisory lock. GET_LOCK returns 1 if acquired, 0 if timeout.
 	// Timeout of 0 means try immediately without waiting.
-	row := e.drv.QueryRow(ctx, "SELECT GET_LOCK(?, 0)", advisoryLockName)
+	row := e.queryRow(ctx, "SELECT GET_LOCK(?, 0)", advisoryLockName)
 	var acquired int
 	if err := row.Scan(&acquired); err != nil {
+		e.releaseDedicated()
 		return fmt.Errorf("mysqlmigrate: advisory lock: %w", err)
 	}
 	if acquired != 1 {
+		e.releaseDedicated()
 		return fmt.Errorf("mysqlmigrate: migration lock is held by another process")
 	}
 
@@ -90,19 +143,32 @@ func (e *Executor) AcquireLock(ctx context.Context, lockedBy string) error {
 		"VALUES (1, NOW(6), ?) "+
 		"ON DUPLICATE KEY UPDATE `locked_at` = NOW(6), `locked_by` = ?",
 		lockTableName)
-	_, err := e.drv.Exec(ctx, query, lockedBy, lockedBy)
-	return err
+	if _, err := e.exec(ctx, query, lockedBy, lockedBy); err != nil {
+		// Best-effort unlock before releasing the connection.
+		_, _ = e.exec(ctx, "SELECT RELEASE_LOCK(?)", advisoryLockName) //nolint:errcheck // best-effort cleanup on error path
+		e.releaseDedicated()
+		return err
+	}
+
+	return nil
 }
 
 // ReleaseLock releases the distributed migration lock.
+// If a dedicated connection was acquired in [AcquireLock], the advisory lock
+// is released on that same connection before the connection is returned to
+// the pool.
 func (e *Executor) ReleaseLock(ctx context.Context) error {
 	// Clear the lock record.
 	query := fmt.Sprintf("UPDATE `%s` SET `locked_at` = NULL, `locked_by` = NULL WHERE `id` = 1",
 		lockTableName)
-	_, _ = e.drv.Exec(ctx, query)
+	_, _ = e.exec(ctx, query) //nolint:errcheck // best-effort lock record clearing
+
+	// Release the advisory lock (on the SAME connection that acquired it).
+	_, err := e.exec(ctx, "SELECT RELEASE_LOCK(?)", advisoryLockName)
+
+	// Release the dedicated connection back to the pool.
+	e.releaseDedicated()
 
-	// Release the advisory lock.
-	_, err := e.drv.Exec(ctx, "SELECT RELEASE_LOCK(?)", advisoryLockName)
 	return err
 }
 
@@ -113,7 +179,7 @@ func (e *Executor) ListApplied(ctx context.Context) ([]*migrate.AppliedMigration
 		"SELECT `id`, `version`, `name`, `group`, CAST(`migrated_at` AS CHAR) FROM `%s` ORDER BY `id` ASC",
 		migrationTableName)
 
-	rows, err := e.drv.Query(ctx, query)
+	rows, err := e.query(ctx, query)
 	if err != nil {
 		return nil, err
 	}
@@ -135,7 +201,7 @@ func (e *Executor) RecordApplied(ctx context.Context, m *migrate.Migration) erro
 	query := fmt.Sprintf(
 		"INSERT INTO `%s` (`version`, `name`, `group`) VALUES (?, ?, ?)",
 		migrationTableName)
-	_, err := e.drv.Exec(ctx, query, m.Version, m.Name, m.Group)
+	_, err := e.exec(ctx, query, m.Version, m.Name, m.Group)
 	return err
 }
 
@@ -144,6 +210,6 @@ func (e *Executor) RemoveApplied(ctx context.Context, m *migrate.Migration) erro
 	query := fmt.Sprintf(
 		"DELETE FROM `%s` WHERE `version` = ? AND `group` = ?",
 		migrationTableName)
-	_, err := e.drv.Exec(ctx, query, m.Version, m.Group)
+	_, err := e.exec(ctx, query, m.Version, m.Group)
 	return err
 }

drivers/pgdriver/pg.go

@@ -30,6 +30,7 @@ type PgDB struct {
 var _ driver.Driver = (*PgDB)(nil)
 var _ driver.StreamCapable = (*PgDB)(nil)
 var _ driver.Preparer = (*PgDB)(nil)
+var _ driver.ConnAcquirer = (*PgDB)(nil)
 
 // New creates a new unconnected PgDB. Call Open to establish a connection pool.
 func New() *PgDB {
@@ -201,6 +202,54 @@ func (db *PgDB) Prepare(ctx context.Context, query string) (driver.Stmt, error)
 	return &pgPoolStmt{conn: conn, sd: sd}, nil
 }
 
+// AcquireConn acquires a dedicated connection from the pool.
+// All operations on the returned DedicatedConn execute on the same
+// underlying PostgreSQL session, making it safe for session-level state
+// such as advisory locks.
+func (db *PgDB) AcquireConn(ctx context.Context) (driver.DedicatedConn, error) {
+	conn, err := db.pool.Acquire(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("pgdriver: acquire dedicated conn: %w", err)
+	}
+	return &pgDedicatedConn{conn: conn}, nil
+}
+
+// pgDedicatedConn wraps a single pgxpool.Conn to implement driver.DedicatedConn.
+// All queries execute on the same underlying connection.
+type pgDedicatedConn struct {
+	conn *pgxpool.Conn
+}
+
+var _ driver.DedicatedConn = (*pgDedicatedConn)(nil)
+
+func (c *pgDedicatedConn) Exec(ctx context.Context, query string, args ...any) (driver.Result, error) {
+	ct, err := c.conn.Exec(ctx, query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("pgdriver: dedicated exec: %w", err)
+	}
+	return &pgResult{ct: ct}, nil
+}
+
+func (c *pgDedicatedConn) Query(ctx context.Context, query string, args ...any) (driver.Rows, error) {
+	rows, err := c.conn.Query(ctx, query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("pgdriver: dedicated query: %w", err)
+	}
+	return &pgRows{rows: rows}, nil
+}
+
+func (c *pgDedicatedConn) QueryRow(ctx context.Context, query string, args ...any) driver.Row {
+	row := c.conn.QueryRow(ctx, query, args...)
+	return &pgRow{row: row}
+}
+
+func (c *pgDedicatedConn) Release() {
+	if c.conn != nil {
+		c.conn.Release()
+		c.conn = nil
+	}
+}
+
 // mapIsolationLevel converts a driver.IsolationLevel to the corresponding
 // pgx.TxIsoLevel string constant.
 func mapIsolationLevel(level driver.IsolationLevel) pgx.TxIsoLevel {