Add support to Set command.

The SET command changes run-time configuration parameters of PostgreSQL.
This commit adds support to the SET command. User may also run the SET
command with the LOCAL or SESSION annotation options to control the
lifespan of the coniguration.

One common use case of the Set command is to set statement timeout;
For that reason: a convenient method 'SetTimeout' was added to the
transaction struct.

Author: aquaomri

set.go

@@ -0,0 +1,109 @@
+package sqlz
+
+import (
+	"database/sql"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/jmoiron/sqlx"
+)
+
+// SetCmd represents a PostgreSQL SET command
+type SetCmd struct {
+	*Statement
+	level       string
+	configParam string
+	value       string
+	execer      Ext
+}
+
+// Set creates a new SetCmd object, with configuration parameter
+// and its value.
+func (db *DB) Set(configParam, value string) *SetCmd {
+	return &SetCmd{
+		configParam: configParam,
+		value:       value,
+		execer:      db.DB,
+		Statement:   &Statement{db.ErrHandlers},
+	}
+}
+
+// Set creates a new SetCmd object, with configuration parameter
+// and its value.
+func (tx *Tx) Set(configParam, value string) *SetCmd {
+	return &SetCmd{
+		configParam: configParam,
+		value:       value,
+		execer:      tx.Tx,
+		Statement:   &Statement{tx.ErrHandlers},
+	}
+}
+
+// SetTimeout sets a statement timeout. When set, any statement
+// (in the transaction) that takes more than the specified duration
+// will be aborted, starting from the time the command arrives
+// at the server from the client. A value of zero turns this off.
+func (tx *Tx) SetTimeout(d time.Duration) (res sql.Result, err error) {
+	stmt := &SetCmd{
+		configParam: "statement_timeout",
+		value:       fmt.Sprintf("\"%dms\"", d.Milliseconds()),
+		execer:      tx.Tx,
+		Statement:   &Statement{tx.ErrHandlers},
+	}
+
+	return stmt.Local().Exec()
+}
+
+// Local sets the configuration parameter locally in a transaction.
+//  The effect of SET LOCAL will last only till the end of the
+// current transaction, whether committed or not.
+func (cmd *SetCmd) Local() *SetCmd {
+	cmd.level = "LOCAL"
+	return cmd
+}
+
+// Session sets the configuration parameter to the entire session.
+//  The effect of SET SESSION will last only till the end of the
+// current session. if issued within a transaction that is later aborted,
+// the effects of the SET command disappear when the transaction is rolled
+// back. Once the surrounding transaction is committed, the effects will persist
+// until the end of the session, unless overridden by another SET.
+func (cmd *SetCmd) Session() *SetCmd {
+	cmd.level = "SESSION"
+	return cmd
+}
+
+// ToSQL generates the SET command SQL and returns a list of
+// bindings. It is used internally by Exec, but is exported if you
+// wish to use it directly.
+func (cmd *SetCmd) ToSQL(rebind bool) (string, []interface{}) {
+	clauses := []string{"SET"}
+	if cmd.level != "" {
+		clauses = append(clauses, cmd.level)
+	}
+
+	clauses = append(clauses, cmd.configParam, "TO", cmd.value)
+
+	asSQL := strings.Join(clauses, " ")
+
+	if rebind {
+		if db, ok := cmd.execer.(*sqlx.DB); ok {
+			asSQL = db.Rebind(asSQL)
+		} else if tx, ok := cmd.execer.(*sqlx.Tx); ok {
+			asSQL = tx.Rebind(asSQL)
+		}
+	}
+
+	return asSQL, []interface{}{}
+}
+
+// Exec executes the SET command, returning the standard
+// sql.Result struct and an error if the query failed.
+func (cmd *SetCmd) Exec() (res sql.Result, err error) {
+	asSQL, bindings := cmd.ToSQL(true)
+	res, err = cmd.execer.Exec(asSQL, bindings...)
+	cmd.Statement.HandleError(err)
+
+	return res, err
+}

set_test.go

@@ -0,0 +1,25 @@
+package sqlz
+
+import "testing"
+
+func TestSet(t *testing.T) {
+	runTests(t, func(dbz *DB) []test {
+		return []test{
+			{
+				name:        "simple set",
+				stmt:        dbz.Set("key", "value"),
+				expectedSQL: "SET key TO value",
+			},
+			{
+				name:        "local set",
+				stmt:        dbz.Set("key", "value").Local(),
+				expectedSQL: "SET LOCAL key TO value",
+			},
+			{
+				name:        "session set",
+				stmt:        dbz.Set("key", "value").Session(),
+				expectedSQL: "SET SESSION key TO value",
+			},
+		}
+	})
+}