diff --git a/postgresql/resource_postgresql_role.go b/postgresql/resource_postgresql_role.go index b47a5886..3d8f19a9 100644 --- a/postgresql/resource_postgresql_role.go +++ b/postgresql/resource_postgresql_role.go @@ -33,6 +33,7 @@ const ( roleRolesAttr = "roles" roleSearchPathAttr = "search_path" roleStatementTimeoutAttr = "statement_timeout" + roleWorkMemAttr = "work_mem" // Deprecated options roleDepEncryptedAttr = "encrypted" @@ -160,6 +161,12 @@ func resourcePostgreSQLRole() *schema.Resource { Description: "Abort any statement that takes more than the specified number of milliseconds", ValidateFunc: validation.IntAtLeast(0), }, + roleWorkMemAttr: { + Type: schema.TypeInt, + Optional: true, + Description: "Specify the amount of memory in bytes to be used for each query for the role", + ValidateFunc: validation.IntAtLeast(0), + }, }, } } @@ -294,6 +301,10 @@ func resourcePostgreSQLRoleCreate(d *schema.ResourceData, meta interface{}) erro return err } + if err = setWorkMem(txn, d); err != nil { + return err + } + if err = txn.Commit(); err != nil { return fmt.Errorf("could not commit transaction: %w", err) } @@ -461,6 +472,13 @@ func resourcePostgreSQLRoleReadImpl(c *Client, d *schema.ResourceData) error { d.Set(roleStatementTimeoutAttr, statementTimeout) + workMem, err := readWorkMem(roleConfig) + if err != nil { + return err + } + + d.Set(roleWorkMemAttr, workMem) + d.SetId(roleName) password, err := readRolePassword(c, d, roleCanLogin) @@ -502,6 +520,23 @@ func readStatementTimeout(roleConfig pq.ByteaArray) (int, error) { return 0, nil } +// readWorkMem searches for a work_mem entry in the rolconfig array. +// In case no such value is present, it returns nil. +func readWorkMem(roleConfig pq.ByteaArray) (int, error) { + for _, v := range roleConfig { + config := string(v) + if strings.HasPrefix(config, roleWorkMemAttr) { + var result = strings.Split(strings.TrimPrefix(config, roleWorkMemAttr+"="), ", ") + res, err := strconv.Atoi(result[0]) + if err != nil { + return -1, fmt.Errorf("Error reading work_mem: %w", 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) { @@ -633,6 +668,10 @@ func resourcePostgreSQLRoleUpdate(d *schema.ResourceData, meta interface{}) erro return err } + if err = setWorkMem(txn, d); err != nil { + return err + } + if err = txn.Commit(); err != nil { return fmt.Errorf("could not commit transaction: %w", err) } @@ -956,3 +995,28 @@ func setStatementTimeout(txn *sql.Tx, d *schema.ResourceData) error { } return nil } + +func setWorkMem(txn *sql.Tx, d *schema.ResourceData) error { + if !d.HasChange(roleWorkMemAttr) { + return nil + } + + roleName := d.Get(roleNameAttr).(string) + workMem := d.Get(roleWorkMemAttr).(int) + if workMem != 0 { + sql := fmt.Sprintf( + "ALTER ROLE %s SET work_mem TO %d", pq.QuoteIdentifier(roleName), workMem, + ) + if _, err := txn.Exec(sql); err != nil { + return fmt.Errorf("could not set work_mem %d for %s: %w", workMem, roleName, err) + } + } else { + sql := fmt.Sprintf( + "ALTER ROLE %s RESET work_mem", pq.QuoteIdentifier(roleName), + ) + if _, err := txn.Exec(sql); err != nil { + return fmt.Errorf("could not reset work_mem for %s: %w", roleName, err) + } + } + return nil +} diff --git a/postgresql/resource_postgresql_role_test.go b/postgresql/resource_postgresql_role_test.go index 35879623..b3a9b07a 100644 --- a/postgresql/resource_postgresql_role_test.go +++ b/postgresql/resource_postgresql_role_test.go @@ -44,6 +44,7 @@ func TestAccPostgresqlRole_Basic(t *testing.T) { 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_defaults", "work_mem", "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"), @@ -119,6 +120,7 @@ resource "postgresql_role" "update_role" { roles = ["${postgresql_role.group_role.name}"] search_path = ["mysearchpath"] statement_timeout = 30000 + work_mem = 1000 } ` resource.Test(t, resource.TestCase{ @@ -141,6 +143,7 @@ resource "postgresql_role" "update_role" { 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"), + resource.TestCheckResourceAttr("postgresql_role.update_role", "work_mem", "0"), testAccCheckRoleCanLogin(t, "update_role", "toto"), ), }, @@ -163,6 +166,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"), + resource.TestCheckResourceAttr("postgresql_role.update_role", "work_mem", "1000"), testAccCheckRoleCanLogin(t, "update_role2", "titi"), ), }, @@ -179,6 +183,7 @@ resource "postgresql_role" "update_role" { 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"), + resource.TestCheckResourceAttr("postgresql_role.update_role", "work_mem", "0"), testAccCheckRoleCanLogin(t, "update_role", "toto"), ), }, @@ -357,6 +362,7 @@ resource "postgresql_role" "role_with_defaults" { skip_drop_role = false valid_until = "infinity" statement_timeout = 0 + work_mem = 0 } resource "postgresql_role" "role_with_create_database" { diff --git a/website/docs/r/postgresql_role.html.markdown b/website/docs/r/postgresql_role.html.markdown index 97df8cb8..1e7a1e85 100644 --- a/website/docs/r/postgresql_role.html.markdown +++ b/website/docs/r/postgresql_role.html.markdown @@ -116,6 +116,8 @@ resource "postgresql_role" "my_replication_role" { * `statement_timeout` - (Optional) Defines [`statement_timeout`](https://www.postgresql.org/docs/current/runtime-config-client.html#RUNTIME-CONFIG-CLIENT-STATEMENT) setting for this role which allows to abort any statement that takes more than the specified amount of time. +* `work_mem` - (Optional) Defines [`work_mem`](https://www.postgresql.org/docs/current/runtime-config-client.html#RUNTIME-CONFIG-CLIENT-STATEMENT) setting for this role which allows to specify the amount of memory in bytes to be used for each query. + ## Import Example `postgresql_role` supports importing resources. Supposing the following