Add set work_mem to role

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
+}

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" {

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