diff --git a/postgresql/helpers.go b/postgresql/helpers.go index bff91f2b..1cb95355 100644 --- a/postgresql/helpers.go +++ b/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( diff --git a/postgresql/resource_postgresql_role.go b/postgresql/resource_postgresql_role.go index 94795a90..ff6c4951 100644 --- a/postgresql/resource_postgresql_role.go +++ b/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 +} diff --git a/postgresql/resource_postgresql_role_test.go b/postgresql/resource_postgresql_role_test.go index 26156b48..9b361c06 100644 --- a/postgresql/resource_postgresql_role_test.go +++ b/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" {