Merge pull request #105 from mced/master

[add] set statement_timeout to role

postgresql/helpers.go

@@ -37,6 +37,14 @@ func validateConnLimit(v interface{}, key string) (warnings []string, errors []e
 	return
 }
 
+func validateStatementTimeout(v interface{}, key string) (warnings []string, errors []error) {
+	value := v.(int)
+	if value < 0 {
+		errors = append(errors, fmt.Errorf("%s can not be less than 0", key))
+	}
+	return
+}
+
 func isRoleMember(db QueryAble, role, member string) (bool, error) {
 	var _rez int
 	err := db.QueryRow(

postgresql/resource_postgresql_role.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"strconv"
 	"strings"
 
 	"github.com/hashicorp/errwrap"
@@ -31,6 +32,7 @@ const (
 	roleValidUntilAttr        = "valid_until"
 	roleRolesAttr             = "roles"
 	roleSearchPathAttr        = "search_path"
+	roleStatementTimeoutAttr  = "statement_timeout"
 
 	// Deprecated options
 	roleDepEncryptedAttr = "encrypted"
@@ -152,6 +154,12 @@ func resourcePostgreSQLRole() *schema.Resource {
 				Default:     false,
 				Description: "Skip actually running the REASSIGN OWNED command when removing a role from PostgreSQL",
 			},
+			roleStatementTimeoutAttr: {
+				Type:         schema.TypeInt,
+				Optional:     true,
+				Description:  "Abort any statement that takes more than the specified number of milliseconds",
+				ValidateFunc: validateStatementTimeout,
+			},
 		},
 	}
 }
@@ -282,6 +290,10 @@ func resourcePostgreSQLRoleCreate(d *schema.ResourceData, meta interface{}) erro
 		return err
 	}
 
+	if err = setStatementTimeout(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return errwrap.Wrapf("could not commit transaction: {{err}}", err)
 	}
@@ -437,6 +449,13 @@ func resourcePostgreSQLRoleReadImpl(c *Client, d *schema.ResourceData) error {
 	d.Set(roleRolesAttr, pgArrayToSet(roleRoles))
 	d.Set(roleSearchPathAttr, readSearchPath(roleConfig))
 
+	statementTimeout, err := readStatementTimeout(roleConfig)
+	if err != nil {
+		return err
+	}
+
+	d.Set(roleStatementTimeoutAttr, statementTimeout)
+
 	d.SetId(roleName)
 
 	password, err := readRolePassword(c, d, roleCanLogin)
@@ -454,12 +473,30 @@ func readSearchPath(roleConfig pq.ByteaArray) []string {
 	for _, v := range roleConfig {
 		config := string(v)
 		if strings.HasPrefix(config, roleSearchPathAttr) {
-			return strings.Split(strings.TrimPrefix(config, roleSearchPathAttr+"="), ", ")
+			var result = strings.Split(strings.TrimPrefix(config, roleSearchPathAttr+"="), ", ")
+			return result
 		}
 	}
 	return nil
 }
 
+// readStatementTimeout searches for a statement_timeout entry in the rolconfig array.
+// In case no such value is present, it returns nil.
+func readStatementTimeout(roleConfig pq.ByteaArray) (int, error) {
+	for _, v := range roleConfig {
+		config := string(v)
+		if strings.HasPrefix(config, roleStatementTimeoutAttr) {
+			var result = strings.Split(strings.TrimPrefix(config, roleStatementTimeoutAttr+"="), ", ")
+			res, err := strconv.Atoi(result[0])
+			if err != nil {
+				return -1, errwrap.Wrapf("Error reading statement_timeout: {{err}}", err)
+			}
+			return res, nil
+		}
+	}
+	return 0, nil
+}
+
 // readRolePassword reads password either from Postgres if admin user is a superuser
 // or only from Terraform state.
 func readRolePassword(c *Client, d *schema.ResourceData, roleCanLogin bool) (string, error) {
@@ -582,6 +619,10 @@ func resourcePostgreSQLRoleUpdate(d *schema.ResourceData, meta interface{}) erro
 		return err
 	}
 
+	if err = setStatementTimeout(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return errwrap.Wrapf("could not commit transaction: {{err}}", err)
 	}
@@ -880,3 +921,28 @@ func alterSearchPath(txn *sql.Tx, d *schema.ResourceData) error {
 	}
 	return nil
 }
+
+func setStatementTimeout(txn *sql.Tx, d *schema.ResourceData) error {
+	if !d.HasChange(roleStatementTimeoutAttr) {
+		return nil
+	}
+
+	roleName := d.Get(roleNameAttr).(string)
+	statementTimeout := d.Get(roleStatementTimeoutAttr).(int)
+	if statementTimeout != 0 {
+		sql := fmt.Sprintf(
+			"ALTER ROLE %s SET statement_timeout TO %d", pq.QuoteIdentifier(roleName), statementTimeout,
+		)
+		if _, err := txn.Exec(sql); err != nil {
+			return errwrap.Wrapf(fmt.Sprintf("could not set statement_timeout %d for %s: {{err}}", statementTimeout, roleName), err)
+		}
+	} else {
+		sql := fmt.Sprintf(
+			"ALTER ROLE %s RESET statement_timeout", pq.QuoteIdentifier(roleName),
+		)
+		if _, err := txn.Exec(sql); err != nil {
+			return errwrap.Wrapf(fmt.Sprintf("could not reset statement_timeout for %s: {{err}}", roleName), err)
+		}
+	}
+	return nil
+}

postgresql/resource_postgresql_role_test.go

@@ -43,6 +43,7 @@ func TestAccPostgresqlRole_Basic(t *testing.T) {
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "valid_until", "infinity"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "skip_drop_role", "false"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "skip_reassign_owned", "false"),
+					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "statement_timeout", "0"),
 
 					resource.TestCheckResourceAttr("postgresql_role.role_with_create_database", "name", "role_with_create_database"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_create_database", "create_database", "true"),
@@ -88,6 +89,7 @@ resource "postgresql_role" "update_role" {
   password = "titi"
   roles = ["${postgresql_role.group_role.name}"]
   search_path = ["mysearchpath"]
+  statement_timeout = 30000
 }
 `
 	resource.Test(t, resource.TestCase{
@@ -109,6 +111,7 @@ resource "postgresql_role" "update_role" {
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "valid_until", "2099-05-04 12:00:00+00"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "roles.#", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.#", "0"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "0"),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -130,6 +133,7 @@ resource "postgresql_role" "update_role" {
 					),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.#", "1"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.0", "mysearchpath"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "30000"),
 					testAccCheckRoleCanLogin(t, "update_role2", "titi"),
 				),
 			},
@@ -145,6 +149,7 @@ resource "postgresql_role" "update_role" {
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "password", "toto"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "roles.#", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.#", "0"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "0"),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -323,6 +328,7 @@ resource "postgresql_role" "role_with_defaults" {
   skip_drop_role = false
   skip_reassign_owned = false
   valid_until = "infinity"
+  statement_timeout = 0
 }
 
 resource "postgresql_role" "role_with_create_database" {