Add support for OF in row lock clauses

This adds support for statements such as `SELECT ... FOR UPDATE OF table NOWAIT` where `OF table`
could not be specified previously. Fixes #285.

internal/jet/select_lock.go

@@ -4,12 +4,14 @@ package jet
 type RowLock interface {
 	Serializer
 
+	OF(...Table) RowLock
 	NOWAIT() RowLock
 	SKIP_LOCKED() RowLock
 }
 
 type selectLockImpl struct {
 	lockStrength       string
+	of                 []Table
 	noWait, skipLocked bool
 }
 
@@ -20,10 +22,15 @@ func NewRowLock(name string) func() RowLock {
 	}
 }
 
-func newSelectLock(lockStrength string) RowLock {
+func newSelectLock(lockStrength string) *selectLockImpl {
 	return &selectLockImpl{lockStrength: lockStrength}
 }
 
+func (s *selectLockImpl) OF(tables ...Table) RowLock {
+	s.of = tables
+	return s
+}
+
 func (s *selectLockImpl) NOWAIT() RowLock {
 	s.noWait = true
 	return s
@@ -37,6 +44,23 @@ func (s *selectLockImpl) SKIP_LOCKED() RowLock {
 func (s *selectLockImpl) serialize(statement StatementType, out *SQLBuilder, options ...SerializeOption) {
 	out.WriteString(s.lockStrength)
 
+	if len(s.of) > 0 {
+		out.WriteString("OF")
+
+		for i, of := range s.of {
+			if i > 0 {
+				out.WriteString(", ")
+			}
+
+			table := of.Alias()
+			if table == "" {
+				table = of.TableName()
+			}
+
+			out.WriteIdentifier(table)
+		}
+	}
+
 	if s.noWait {
 		out.WriteString("NOWAIT")
 	}

mysql/select_statement_test.go

@@ -122,6 +122,11 @@ FOR UPDATE;
 SELECT table1.col_bool AS "table1.col_bool"
 FROM db.table1
 FOR SHARE NOWAIT;
+`)
+	testutils.AssertStatementSql(t, SELECT(table1ColBool).FROM(table1).FOR(UPDATE().OF(table1).NOWAIT()), `
+SELECT table1.col_bool AS "table1.col_bool"
+FROM db.table1
+FOR UPDATE OF table1 NOWAIT;
 `)
 }
 

postgres/select_statement_test.go

@@ -132,5 +132,11 @@ FOR KEY SHARE NOWAIT;
 SELECT table1.col_bool AS "table1.col_bool"
 FROM db.table1
 FOR NO KEY UPDATE SKIP LOCKED;
+`)
+
+	assertStatementSql(t, SELECT(table1ColBool).FROM(table1).FOR(UPDATE().OF(table1, table2).NOWAIT()), `
+SELECT table1.col_bool AS "table1.col_bool"
+FROM db.table1
+FOR UPDATE OF table1, table2 NOWAIT;
 `)
 }

tests/postgres/select_test.go

@@ -2251,6 +2251,18 @@ FOR`
 	}
 }
 
+func TestRowLockWithJoins(t *testing.T) {
+	query := SELECT(STAR).
+		FROM(
+			Film.
+				INNER_JOIN(FilmCategory, FilmCategory.FilmID.EQ(Film.FilmID)).
+				LEFT_JOIN(FilmActor, FilmActor.FilmID.EQ(Film.FilmID))).
+		LIMIT(1).
+		FOR(UPDATE().OF(Film, FilmCategory).NOWAIT())
+
+	testutils.AssertExecAndRollback(t, query, db, 1)
+}
+
 func TestQuickStart(t *testing.T) {
 
 	var expectedSQL = `