diff --git a/lock.go b/lock.go new file mode 100644 index 0000000..f1bfc64 --- /dev/null +++ b/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") + } +} diff --git a/lock_test.go b/lock_test.go new file mode 100644 index 0000000..7fefbb6 --- /dev/null +++ b/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()) +} diff --git a/select.go b/select.go index 8780d3a..ced2ef6 100644 --- a/select.go +++ b/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 } diff --git a/select_test.go b/select_test.go index bc89fcb..9eb9b9f 100644 --- a/select_test.go +++ b/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 {