add lock clause support in select queries

lock.go

@@ -0,0 +1,41 @@
+package sqlbuilder
+
+type LockClause interface {
+	NoWait(noWait bool) LockClause
+	serialize(bldr *builder)
+}
+
+type lockClauseImpl struct {
+	strength string
+	tables   []Table
+	noWait   bool
+}
+
+func Lock(strength string, tables ...Table) LockClause {
+	return &lockClauseImpl{
+		strength: strength,
+		tables:   tables,
+	}
+}
+
+func (l *lockClauseImpl) NoWait(noWait bool) LockClause {
+	l.noWait = noWait
+
+	return l
+}
+
+func (l *lockClauseImpl) serialize(bldr *builder) {
+	t := make([]serializable, len(l.tables))
+	for i, v := range l.tables {
+		t[i] = v
+	}
+
+	bldr.Append(l.strength)
+	if len(l.tables) > 0 {
+		bldr.Append(" OF ")
+		bldr.AppendItems(t, ", ")
+	}
+	if l.noWait {
+		bldr.Append(" NOWAIT")
+	}
+}

lock_test.go

@@ -0,0 +1,35 @@
+package sqlbuilder
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestLockClauseImplements(t *testing.T) {
+	a := assert.New(t)
+	a.Implements(new(LockClause), &lockClauseImpl{})
+}
+
+func TestLockClause(t *testing.T) {
+	a := assert.New(t)
+	b := newBuilder()
+	table1 := NewTable(
+		"TABLE_A",
+		&TableOption{},
+		IntColumn("id", &ColumnOption{
+			PrimaryKey: true,
+		}),
+	)
+	table2 := NewTable(
+		"TABLE_B",
+		&TableOption{},
+		IntColumn("id", &ColumnOption{
+			PrimaryKey: true,
+		}),
+	)
+
+	Lock("UPDATE", table1, table2).serialize(b)
+	a.Equal(`UPDATE OF "TABLE_A", "TABLE_B"`, b.query.String())
+	a.Equal([]interface{}{}, b.Args())
+	a.NoError(b.Err())
+}

select.go

@@ -11,6 +11,7 @@ type SelectStatement struct {
 	limit    int
 	offset   int
 	having   Condition
+	locks    []LockClause
 
 	err error
 }
@@ -82,7 +83,7 @@ func (b *SelectStatement) GroupBy(columns ...Column) *SelectStatement {
 	return b
 }
 
-// GroupBy sets "HAVING" clause with the cond.
+// Having sets "HAVING" clause with the cond.
 func (b *SelectStatement) Having(cond Condition) *SelectStatement {
 	if b.err != nil {
 		return b
@@ -124,6 +125,12 @@ func (b *SelectStatement) Offset(offset int) *SelectStatement {
 	return b
 }
 
+// Lock sets LOCK clause(s).
+func (b *SelectStatement) Locks(locks ...LockClause) *SelectStatement {
+	b.locks = locks
+	return b
+}
+
 func (b *SelectStatement) serialize(bldr *builder) {
 	if b.err != nil {
 		bldr.SetError(b.err)
@@ -179,6 +186,17 @@ func (b *SelectStatement) serialize(bldr *builder) {
 		bldr.Append(" OFFSET ")
 		bldr.AppendValue(b.offset)
 	}
+
+	// FOR lock_strength [ OF table_name [, ...] ] [ NOWAIT ]
+	if len(b.locks) > 0 {
+		bldr.Append(" FOR ")
+		l := make([]serializable, len(b.locks))
+		for i, v := range b.locks {
+			l[i] = v
+		}
+		bldr.AppendItems(l, " ")
+	}
+
 	return
 }
 

select_test.go

@@ -94,6 +94,12 @@ func TestSelect(t *testing.T) {
 		query:  ``,
 		args:   []interface{}{},
 		errmsg: "sqlbuilder: GROUP BY clause is not found.",
+	}, {
+		stmt: Select(table1).
+			Locks(Lock("UPDATE", table1), Lock("SHARE", table1)),
+		query:  `SELECT * FROM "TABLE_A" FOR UPDATE OF "TABLE_A" SHARE OF "TABLE_A";`,
+		args:   []interface{}{},
+		errmsg: "",
 	}}
 
 	for num, c := range cases {