From 12fc7d231aff14108c6cdcd582b170adb405389b Mon Sep 17 00:00:00 2001
From: Qishuai Liu <lqs@lqs.me>
Date: Tue, 22 Aug 2023 22:29:11 +0900
Subject: [PATCH] add ForUpdateNoWait & ForUpdateSkipLocked

---
 select.go      | 12 ++++++++++++
 select_test.go | 12 +++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/select.go b/select.go
index a216ba7..4d32201 100644
--- a/select.go
+++ b/select.go
@@ -105,6 +105,8 @@ type selectWithOffset interface {
 type toSelectWithLock interface {
 	LockInShareMode() selectWithLock
 	ForUpdate() selectWithLock
+	ForUpdateNoWait() selectWithLock
+	ForUpdateSkipLocked() selectWithLock
 }
 
 type selectWithLock interface {
@@ -392,6 +394,16 @@ func (s selectStatus) ForUpdate() selectWithLock {
 	return s
 }
 
+func (s selectStatus) ForUpdateNoWait() selectWithLock {
+	s.lock = " FOR UPDATE NOWAIT"
+	return s
+}
+
+func (s selectStatus) ForUpdateSkipLocked() selectWithLock {
+	s.lock = " FOR UPDATE SKIP LOCKED"
+	return s
+}
+
 func (s selectStatus) asDerivedTable(name string) Table {
 	return derivedTable{
 		name:         name,
diff --git a/select_test.go b/select_test.go
index 5619735..d4103f4 100644
--- a/select_test.go
+++ b/select_test.go
@@ -223,10 +223,16 @@ func TestFetchAll(t *testing.T) {
 }
 
 func TestLock(t *testing.T) {
-	db := database{}
+	db := newMockDatabase()
 	table1 := NewTable("table1")
-	db.Select(1).From(table1).LockInShareMode()
-	db.Select(1).From(table1).ForUpdate()
+	_, _ = db.Select(1).From(table1).LockInShareMode().FetchAll()
+	assertLastSql(t, "SELECT 1 FROM `table1` LOCK IN SHARE MODE")
+	_, _ = db.Select(1).From(table1).ForUpdate().FetchAll()
+	assertLastSql(t, "SELECT 1 FROM `table1` FOR UPDATE")
+	_, _ = db.Select(1).From(table1).ForUpdateNoWait().FetchAll()
+	assertLastSql(t, "SELECT 1 FROM `table1` FOR UPDATE NOWAIT")
+	_, _ = db.Select(1).From(table1).ForUpdateSkipLocked().FetchAll()
+	assertLastSql(t, "SELECT 1 FROM `table1` FOR UPDATE SKIP LOCKED")
 }
 
 func TestUnion(t *testing.T) {