Merge pull request #9 from bogunenko/support-select-for-update

FEATURE: support SELECT ... FOR UPDATE

select.go

@@ -17,6 +17,7 @@ type SelectStmt struct {
 
 	LimitCount  int64
 	OffsetCount int64
+	IsForUpdate bool
 }
 
 // Build builds `SELECT ...` in dialect
@@ -113,6 +114,10 @@ func (b *SelectStmt) Build(d Dialect, buf Buffer) error {
 		buf.WriteString(" ")
 		buf.WriteString(d.Limit(b.OffsetCount, b.LimitCount))
 	}
+
+	if b.IsForUpdate {
+		buf.WriteString(" FOR UPDATE")
+	}
 	return nil
 }
 
@@ -203,6 +208,12 @@ func (b *SelectStmt) Offset(n uint64) *SelectStmt {
 	return b
 }
 
+// ForUpdate adds `FOR UPDATE`
+func (b *SelectStmt) ForUpdate() *SelectStmt {
+	b.IsForUpdate = true
+	return b
+}
+
 // Join joins table on condition
 func (b *SelectStmt) Join(table, on interface{}) *SelectStmt {
 	b.JoinTable = append(b.JoinTable, join(inner, table, on))

select_builder.go

@@ -184,3 +184,9 @@ func (b *SelectBuilder) Where(query interface{}, value ...interface{}) *SelectBu
 	b.SelectStmt.Where(query, value...)
 	return b
 }
+
+// ForUpdate adds lock via FOR UPDATE
+func (b *SelectBuilder) ForUpdate() *SelectBuilder {
+	b.SelectStmt.ForUpdate()
+	return b
+}

select_test.go

@@ -18,10 +18,11 @@ func TestSelectStmt(t *testing.T) {
 		Having(Eq("e", 2)).
 		OrderAsc("f").
 		Limit(3).
-		Offset(4)
+		Offset(4).
+		ForUpdate()
 	err := builder.Build(dialect.MySQL, buf)
 	assert.NoError(t, err)
-	assert.Equal(t, "SELECT DISTINCT a, b FROM ? LEFT JOIN `table2` ON table.a1 = table.a2 WHERE (`c` = ?) GROUP BY d HAVING (`e` = ?) ORDER BY f ASC LIMIT 4,3", buf.String())
+	assert.Equal(t, "SELECT DISTINCT a, b FROM ? LEFT JOIN `table2` ON table.a1 = table.a2 WHERE (`c` = ?) GROUP BY d HAVING (`e` = ?) ORDER BY f ASC LIMIT 4,3 FOR UPDATE", buf.String())
 	// two functions cannot be compared
 	assert.Equal(t, 3, len(buf.Value()))
 }