Additional tests for row lock UPDATE OF use case.

tests/mysql/select_test.go

@@ -2,6 +2,7 @@ package mysql
 
 import (
 	"context"
+	"database/sql"
 	"strings"
 	"testing"
 	"time"
@@ -632,6 +633,86 @@ FOR`
 	}
 }
 
+func TestRowLockWithUpdateOf(t *testing.T) {
+	skipForMariaDB(t) // MariaDB does not support UPDATE OF
+
+	stmt := SELECT(
+		Film.FilmID,
+		Film.Title,
+		Actor.ActorID,
+		Actor.FirstName,
+	).FROM(
+		Film.
+			INNER_JOIN(FilmCategory, FilmCategory.FilmID.EQ(Film.FilmID)).
+			INNER_JOIN(FilmActor, FilmActor.FilmID.EQ(Film.FilmID)).
+			INNER_JOIN(Actor, Actor.ActorID.EQ(FilmActor.ActorID)),
+	).LIMIT(
+		1,
+	).FOR(
+		UPDATE().OF(Film, Actor).NOWAIT(),
+	)
+
+	testutils.AssertDebugStatementSql(t, stmt, `
+SELECT film.film_id AS "film.film_id",
+     film.title AS "film.title",
+     actor.actor_id AS "actor.actor_id",
+     actor.first_name AS "actor.first_name"
+FROM dvds.film
+     INNER JOIN dvds.film_category ON (film_category.film_id = film.film_id)
+     INNER JOIN dvds.film_actor ON (film_actor.film_id = film.film_id)
+     INNER JOIN dvds.actor ON (actor.actor_id = film_actor.actor_id)
+LIMIT 1
+FOR UPDATE OF film, actor NOWAIT;
+`)
+
+	testutils.ExecuteInTxAndRollback(t, db, func(tx *sql.Tx) {
+		var dest []struct {
+			model.Film
+			CategoryID int
+			Actor      []model.Actor
+		}
+
+		err := stmt.Query(tx, &dest)
+		require.NoError(t, err)
+		require.Len(t, dest, 1)
+	})
+}
+
+func TestRowLockWithUpdateOfAliasedTable(t *testing.T) {
+	skipForMariaDB(t) // MariaDB does not support UPDATE OF
+
+	myFilm := Film.AS("myFilm")
+
+	stmt := SELECT(
+		myFilm.FilmID,
+		myFilm.Title,
+	).FROM(
+		myFilm,
+	).LIMIT(
+		1,
+	).FOR(
+		UPDATE().OF(myFilm),
+	)
+
+	testutils.AssertDebugStatementSql(t, stmt, strings.ReplaceAll(`
+SELECT ''myFilm''.film_id AS "myFilm.film_id",
+     ''myFilm''.title AS "myFilm.title"
+FROM dvds.film AS ''myFilm''
+LIMIT 1
+FOR UPDATE OF ''myFilm'';
+`, "''", "`"))
+
+	testutils.ExecuteInTxAndRollback(t, db, func(tx *sql.Tx) {
+		var dest []struct {
+			model.Film `alias:"myFilm.*"`
+		}
+
+		err := stmt.Query(db, &dest)
+		require.NoError(t, err)
+		require.Len(t, dest, 1)
+	})
+}
+
 func TestExpressionWrappers(t *testing.T) {
 	query := SELECT(
 		BoolExp(Raw("true")),

tests/postgres/select_test.go

@@ -2251,16 +2251,81 @@ FOR`
 	}
 }
 
-func TestRowLockWithJoins(t *testing.T) {
+func TestRowLockWithUpdateOf(t *testing.T) {
-	query := SELECT(STAR).
+	stmt := SELECT(
-		FROM(
+		Film.FilmID,
-			Film.
+		Film.Title,
-				INNER_JOIN(FilmCategory, FilmCategory.FilmID.EQ(Film.FilmID)).
+		Actor.ActorID,
-				LEFT_JOIN(FilmActor, FilmActor.FilmID.EQ(Film.FilmID))).
+		Actor.FirstName,
-		LIMIT(1).
+	).FROM(
-		FOR(UPDATE().OF(Film, FilmCategory).NOWAIT())
+		Film.
+			INNER_JOIN(FilmCategory, FilmCategory.FilmID.EQ(Film.FilmID)).
+			INNER_JOIN(FilmActor, FilmActor.FilmID.EQ(Film.FilmID)).
+			INNER_JOIN(Actor, Actor.ActorID.EQ(FilmActor.ActorID)),
+	).LIMIT(
+		1,
+	).FOR(
+		UPDATE().OF(Film, Actor).NOWAIT(),
+	)
 
-	testutils.AssertExecAndRollback(t, query, db, 1)
+	testutils.AssertDebugStatementSql(t, stmt, `
+SELECT film.film_id AS "film.film_id",
+     film.title AS "film.title",
+     actor.actor_id AS "actor.actor_id",
+     actor.first_name AS "actor.first_name"
+FROM dvds.film
+     INNER JOIN dvds.film_category ON (film_category.film_id = film.film_id)
+     INNER JOIN dvds.film_actor ON (film_actor.film_id = film.film_id)
+     INNER JOIN dvds.actor ON (actor.actor_id = film_actor.actor_id)
+LIMIT 1
+FOR UPDATE OF film, actor NOWAIT;
+`)
+
+	testutils.ExecuteInTxAndRollback(t, db, func(tx *sql.Tx) {
+		var dest []struct {
+			model.Film
+			CategoryID int
+			Actor      []model.Actor
+		}
+
+		err := stmt.Query(tx, &dest)
+		require.NoError(t, err)
+		require.Len(t, dest, 1)
+	})
+}
+
+func TestRowLockWithUpdateOfAliasedTable(t *testing.T) {
+
+	myFilm := Film.AS("myFilm")
+
+	stmt := SELECT(
+		myFilm.FilmID,
+		myFilm.Title,
+	).FROM(
+		myFilm,
+	).LIMIT(
+		1,
+	).FOR(
+		UPDATE().OF(myFilm),
+	)
+
+	testutils.AssertDebugStatementSql(t, stmt, `
+SELECT "myFilm".film_id AS "myFilm.film_id",
+     "myFilm".title AS "myFilm.title"
+FROM dvds.film AS "myFilm"
+LIMIT 1
+FOR UPDATE OF "myFilm";
+`)
+
+	testutils.ExecuteInTxAndRollback(t, db, func(tx *sql.Tx) {
+		var dest []struct {
+			model.Film `alias:"myFilm.*"`
+		}
+
+		err := stmt.Query(db, &dest)
+		require.NoError(t, err)
+		require.Len(t, dest, 1)
+	})
 }
 
 func TestQuickStart(t *testing.T) {