diff --git a/internal/jet/select_lock.go b/internal/jet/select_lock.go index 4694ee3..c9d431a 100644 --- a/internal/jet/select_lock.go +++ b/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") } diff --git a/mysql/select_statement_test.go b/mysql/select_statement_test.go index bd3a0e9..1136dd7 100644 --- a/mysql/select_statement_test.go +++ b/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; `) } diff --git a/postgres/select_statement_test.go b/postgres/select_statement_test.go index b487f90..5dcacf0 100644 --- a/postgres/select_statement_test.go +++ b/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; `) } diff --git a/tests/postgres/select_test.go b/tests/postgres/select_test.go index a6d9588..1f686b2 100644 --- a/tests/postgres/select_test.go +++ b/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 = `